generator.js

const STORAGE_KEY = 'pxAutomationState';
const $ = (id) => document.getElementById(id);

let csvRows = null;
let csvHeaders = [];
let generatedPayload = [];
let generatedChunks = [];
let fullArrayJson = '';

function setStatus(message, kind = 'info') {
  const el = $('status');
  el.textContent = message;
  el.className = `status ${kind}`;
}

function normalizeSequence(raw) {
  return String(raw || '').replace(/\s+/g, '').toUpperCase();
}

function validateProteinSequence(sequence, label) {
  if (!sequence) throw new Error(`${label} has an empty sequence.`);
  if (!/^[ACDEFGHIKLMNPQRSTVWYX]+$/.test(sequence)) {
    throw new Error(`${label} contains unsupported characters. Use the 20 standard amino-acid codes or X.`);
  }
}

function addChainRow(name = '', sequence = '') {
  const row = document.createElement('div');
  row.className = 'chain-row';

  const nameInput = document.createElement('input');
  nameInput.type = 'text';
  nameInput.placeholder = 'Partner name';
  nameInput.value = name;
  nameInput.className = 'chain-name';

  const sequenceInput = document.createElement('textarea');
  sequenceInput.placeholder = 'Paste amino-acid sequence…';
  sequenceInput.value = sequence;
  sequenceInput.className = 'chain-seq';

  const remove = document.createElement('button');
  remove.className = 'remove';
  remove.textContent = '✕';
  remove.title = 'Remove chain';
  remove.onclick = () => row.remove();
  row.append(nameInput, sequenceInput, remove);
  $('chains').appendChild(row);
}

function getChains() {
  const chains = [];
  document.querySelectorAll('#chains .chain-row').forEach((row, index) => {
    const sequence = normalizeSequence(row.querySelector('.chain-seq').value);
    if (!sequence) return;
    const name = row.querySelector('.chain-name').value.trim() || `Partner${index + 1}`;
    validateProteinSequence(sequence, name);
    chains.push({ name, sequence });
  });
  return chains;
}

function parseCSV(text) {
  const rows = [];
  let row = [];
  let field = '';
  let inQuotes = false;
  for (let index = 0; index < text.length; index += 1) {
    const char = text[index];
    if (inQuotes) {
      if (char === '"') {
        if (text[index + 1] === '"') { field += '"'; index += 1; }
        else inQuotes = false;
      } else field += char;
    } else if (char === '"') inQuotes = true;
    else if (char === ',') { row.push(field); field = ''; }
    else if (char === '\n') { row.push(field); rows.push(row); row = []; field = ''; }
    else if (char !== '\r') field += char;
  }
  if (field.length || row.length) { row.push(field); rows.push(row); }
  return rows;
}

function loadCsvText(text) {
  const clean = text.charCodeAt(0) === 0xfeff ? text.slice(1) : text;
  const matrix = parseCSV(clean).filter((row) => row.some((cell) => cell.trim()));
  if (!matrix.length) throw new Error('The CSV appears to be empty.');
  csvHeaders = matrix[0].map((header) => header.trim());
  csvRows = matrix.slice(1).map((row) => Object.fromEntries(csvHeaders.map((header, index) => [header, row[index] || ''])));

  const dataList = $('csvCols');
  dataList.innerHTML = '';
  csvHeaders.forEach((header) => {
    const option = document.createElement('option');
    option.value = header;
    dataList.appendChild(option);
  });
  const guess = (pattern) => csvHeaders.find((header) => pattern.test(header));
  $('nameCol').value = guess(/name|id|target/i) || csvHeaders[0] || '';
  $('seqCol').value = guess(/seq/i) || csvHeaders[1] || '';
  $('csvInfo').textContent = `Loaded ${csvRows.length} row(s). Columns: ${csvHeaders.join(', ')}`;
}

$('csvFile').addEventListener('change', (event) => {
  const file = event.target.files?.[0];
  if (!file) return;
  const reader = new FileReader();
  reader.onload = () => {
    try { loadCsvText(String(reader.result)); setStatus('CSV loaded. Add the fixed partner sequence and generate.', 'ok'); }
    catch (error) { setStatus(error.message, 'err'); }
  };
  reader.onerror = () => setStatus('The CSV file could not be read.', 'err');
  reader.readAsText(file);
});

function loadTargets(nameColumn, sequenceColumn) {
  const targets = [];
  csvRows.forEach((row, index) => {
    const sequence = normalizeSequence(row[sequenceColumn]);
    if (!sequence) return;
    const name = String(row[nameColumn] || '').trim() || `Target${index + 1}`;
    validateProteinSequence(sequence, name);
    targets.push({ name, sequence });
  });
  return targets;
}

function safeName(value) {
  return String(value || '').replace(/[^a-z0-9_.-]+/gi, '_').replace(/^_+|_+$/g, '') || 'prediction';
}

function proteinChain(sequence) {
  return { proteinChain: { sequence, count: 1 } };
}

function buildPayload(targets, chains, targetFirst) {
  return targets.map((target) => {
    const targetEntity = proteinChain(target.sequence);
    const partnerEntities = chains.map((chain) => proteinChain(chain.sequence));
    return {
      name: safeName(`${target.name}_${chains.map((chain) => chain.name).join('_')}`),
      sequences: targetFirst ? [targetEntity, ...partnerEntities] : [...partnerEntities, targetEntity],
      covalent_bonds: []
    };
  });
}

function chunkPayload(payload, size) {
  const chunks = [];
  for (let start = 0; start < payload.length; start += size) {
    const slice = payload.slice(start, start + size);
    chunks.push({ slice, start: start + 1, end: start + slice.length });
  }
  return chunks;
}

$('generate').addEventListener('click', () => {
  try {
    const chains = getChains();
    if (!chains.length) throw new Error('Add at least one fixed partner chain with a sequence.');
    if (!csvRows) throw new Error('Choose a CSV file containing target sequences.');
    const nameColumn = $('nameCol').value.trim();
    const sequenceColumn = $('seqCol').value.trim();
    if (!csvHeaders.includes(nameColumn) || !csvHeaders.includes(sequenceColumn)) {
      throw new Error(`Choose valid CSV columns. Available: ${csvHeaders.join(', ')}`);
    }
    const targets = loadTargets(nameColumn, sequenceColumn);
    if (!targets.length) throw new Error('No usable target sequences were found.');

    generatedPayload = buildPayload(targets, chains, $('targetFirst').checked);
    fullArrayJson = JSON.stringify(generatedPayload, null, 2);
    const chunkSize = Math.max(1, Number($('chunkSize').value) || 100);
    const baseName = safeName($('baseName').value.replace(/\.json$/i, '') || 'protenix_predictions');
    generatedChunks = chunkPayload(generatedPayload, chunkSize).map((chunk) => ({
      name: `${baseName}_${chunk.start}-${chunk.end}.json`,
      json: JSON.stringify(chunk.slice, null, 2)
    }));
    $('preview').value = generatedChunks[0].json;
    ['queueJobs', 'downloadSplit', 'downloadSingle', 'copyAll'].forEach((id) => { $(id).disabled = false; });
    setStatus(`Generated ${generatedPayload.length} prediction job${generatedPayload.length === 1 ? '' : 's'} from ${targets.length} target${targets.length === 1 ? '' : 's'}.`, 'ok');
  } catch (error) {
    setStatus(error.message, 'err');
  }
});

function newId() {
  return crypto.randomUUID ? crypto.randomUUID() : `${Date.now()}-${Math.random().toString(16).slice(2)}`;
}

$('queueJobs').addEventListener('click', async () => {
  if (!generatedPayload.length) return;
  const stored = await chrome.storage.local.get(STORAGE_KEY);
  const state = stored[STORAGE_KEY] || {
    version: 1, queue: [], running: false, paused: false, currentId: null, submitUrl: '', downloadedKeys: [], note: '',
    settings: { modelName: 'protenix_base_default_v1.0.0', msa: 'unchanged', template: 'unchanged', delayMs: 3000, outcomeTimeoutMs: 180000 }
  };
  generatedPayload.forEach((entry) => {
    state.queue.push({
      id: newId(),
      name: entry.name,
      fileName: `${safeName(entry.name)}.json`,
      json: JSON.stringify([entry], null, 2),
      status: 'pending',
      attempts: 0,
      error: '',
      addedAt: Date.now()
    });
  });
  state.note = `Added ${generatedPayload.length} generated job${generatedPayload.length === 1 ? '' : 's'} to the queue.`;
  state.updatedAt = Date.now();
  await chrome.storage.local.set({ [STORAGE_KEY]: state });
  setStatus(`${state.note} Open the extension popup on the Protenix submission page to start.`, 'ok');
});

function downloadBlob(filename, text) {
  const blob = new Blob([text], { type: 'application/json' });
  const url = URL.createObjectURL(blob);
  const link = document.createElement('a');
  link.href = url;
  link.download = filename;
  document.body.appendChild(link);
  link.click();
  link.remove();
  setTimeout(() => URL.revokeObjectURL(url), 4000);
}

$('downloadSplit').addEventListener('click', async () => {
  for (const chunk of generatedChunks) {
    downloadBlob(chunk.name, chunk.json);
    await new Promise((resolve) => setTimeout(resolve, 350));
  }
  setStatus(`Downloaded ${generatedChunks.length} file(s). Allow multiple downloads if Chrome asks.`, 'ok');
});

$('downloadSingle').addEventListener('click', () => {
  if (!fullArrayJson) return;
  const baseName = safeName($('baseName').value.replace(/\.json$/i, '') || 'protenix_predictions');
  downloadBlob(`${baseName}.json`, fullArrayJson);
});

$('copyAll').addEventListener('click', async () => {
  try {
    await navigator.clipboard.writeText(fullArrayJson);
    setStatus('Copied the full Protenix JSON array to the clipboard.', 'ok');
  } catch (_) {
    setStatus('Clipboard access was blocked; use a download button instead.', 'err');
  }
});

addChainRow();
$('addChain').addEventListener('click', () => addChainRow());