frontend/src/App.jsx

import React from 'react';
import Upload from './screens/Upload';
import Pipeline from './screens/Pipeline';
import Reader from './screens/Reader';
import AnnotatorPanel, { apLoadAnnotator, apSaveAnnotator, AP_PROVIDERS } from './components/AnnotatorPanel';
import DisplayMenu, { PALETTES } from './components/DisplayMenu';
import Toast from './components/Toast';
import { startParse, jobEvents, annotateStream } from './lib/api';
import { hashPdf, exportSession, importSession } from './lib/session';
import FigurePopover from './components/FigurePopover';
import NotebookDrawer from './components/NotebookDrawer';

// Default annotation mode for fresh parses and legend-less notes.
const DEFAULT_MODE = 'history';

// ── Display / tweak defaults ──────────────────────────────────────────────────
// Palette [0] = Parchment — same hard-coded values that were previously in
// rootStyle, so the default visual output is unchanged.
const DISPLAY_DEFAULTS = {
  paperPalette: ['#f6f1e6', '#2b2418', '#9a3b20'],
  bodySize: 19,
  dropCaps: true,
  justify: true,
  showScores: true,
  notesSide: 'right',
  litCount: 3,
};

const DISPLAY_KEY = 'astroparse_display_v1';
const SESSION_KEY = 'astroparse_session_v2';

function loadDisplay() {
  try {
    const raw = localStorage.getItem(DISPLAY_KEY);
    if (raw) {
      const v = JSON.parse(raw);
      if (v && Array.isArray(v.paperPalette) && v.paperPalette.length === 3) return { ...DISPLAY_DEFAULTS, ...v };
    }
  } catch (e) { /* ignore corrupt */ }
  return { ...DISPLAY_DEFAULTS };
}

export default function App() {
  const [screen, setScreen] = React.useState('upload');
  const [fileName, setFileName] = React.useState('');
  const [stageEvents, setStageEvents] = React.useState([]);
  const [result, setResult] = React.useState(null);
  const [pdfHash, setPdfHash] = React.useState('');

  // noteModes: { [paraId]: modeId }
  const [noteModes, setNoteModes] = React.useState({});
  // commentary: { ['paraId:modeId']: { text, source, error, message, streaming? } }
  const [commentary, setCommentary] = React.useState({});
  const commentaryRef = React.useRef(commentary);
  commentaryRef.current = commentary;
  // generating: { ['paraId:modeId']: true }
  const [generating, setGenerating] = React.useState({});

  // Figure popover state: { fig, n, pos } or null
  const [figPop, setFigPop] = React.useState(null);

  // User notes state.
  const [userNotes, setUserNotes] = React.useState({});
  const [paperNote, setPaperNote] = React.useState('');
  const [notebookOpen, setNotebookOpen] = React.useState(false);

  // Toast state.
  const [toast, setToast] = React.useState(null);
  const toastTimerRef = React.useRef(null);

  // Annotator config — loaded from localStorage on mount.
  const [annotator, setAnnotator] = React.useState(apLoadAnnotator);
  const [settingsOpen, setSettingsOpen] = React.useState(false);

  // Display/tweak state — loaded from localStorage on mount.
  const [tweaks, setTweaks] = React.useState(loadDisplay);
  const [displayOpen, setDisplayOpen] = React.useState(false);

  // Inflight set — dedupes concurrent requests for the same (paraId, modeId).
  const inflight = React.useRef(new Set());

  // Toolbar file-import ref for session load.
  const importRef = React.useRef(null);

  // lastFileName for Upload screen (from localStorage session on mount).
  const [lastFileName, setLastFileName] = React.useState(null);
  // Stash for the raw session text so onOpenLast can restore it.
  const storedSessionRef = React.useRef(null);

  // showToast: sets toast message and auto-clears after 2600ms.
  const showToast = React.useCallback((msg) => {
    setToast(msg);
    if (toastTimerRef.current) clearTimeout(toastTimerRef.current);
    toastTimerRef.current = setTimeout(() => setToast(null), 2600);
  }, []);

  const handleSaveNote = React.useCallback((paraId, text) => {
    setUserNotes((prev) => {
      if (!text.trim()) {
        const next = { ...prev }; delete next[paraId]; return next;
      }
      return { ...prev, [paraId]: { text, updatedAt: new Date().toISOString() } };
    });
  }, []);

  const handleSavePaperNote = React.useCallback((text) => {
    setPaperNote(text);
  }, []);

  // ── restoreSession ─────────────────────────────────────────────────────────
  // Shared restore path used by toolbar "load", upload-screen onSession, and
  // onOpenLast.  Reads text, parses, populates all state, switches to reader.
  const restoreSession = React.useCallback((text) => {
    try {
      const s = importSession(text);
      fromParseRef.current = false;
      setFileName(s.fileName || '');
      setResult({
        paper: s.paper,
        litPapers: s.litPapers,
        litByPara: s.litByPara,
        figures: s.figures || [],
        litByFig: s.litByFig || {},
        equations: s.equations || [],
        references: s.references || [],
        citeIndex: s.citeIndex || {},
      });
      setPdfHash(s.pdfHash || '');
      if (s.noteModes) {
        setNoteModes(s.noteModes);
      } else {
        const modes = {};
        (s.paper?.paragraphs || []).forEach((p) => { modes[p.id] = DEFAULT_MODE; });
        setNoteModes(modes);
      }
      setCommentary(s.commentary || {});
      setGenerating({});
      setUserNotes(s.userNotes || {});
      setPaperNote(s.paperNote || '');
      setScreen('reader');
      showToast('session restored');
    } catch {
      showToast('could not read that session file');
    }
  }, [showToast]);

  // ── On mount: read localStorage session for Upload "open last" button ──────
  React.useEffect(() => {
    try {
      const raw = localStorage.getItem(SESSION_KEY);
      if (raw) {
        const parsed = JSON.parse(raw);
        if (parsed && parsed.fileName) {
          storedSessionRef.current = raw;
          setLastFileName(parsed.fileName);
        }
      }
    } catch (e) { /* ignore corrupt */ }
  }, []); // eslint-disable-line react-hooks/exhaustive-deps

  // ── Autosave ───────────────────────────────────────────────────────────────
  // Write to localStorage whenever commentary or noteModes change, if in reader.
  React.useEffect(() => {
    if (screen !== 'reader' || !result) return;
    try {
      const json = exportSession({
        fileName,
        pdfHash,
        paper: result.paper,
        litPapers: result.litPapers,
        litByPara: result.litByPara,
        figures: result.figures || [],
        litByFig: result.litByFig || {},
        equations: result.equations || [],
        references: result.references || [],
        citeIndex: result.citeIndex || {},
        userNotes,
        paperNote,
        noteModes,
        commentary,
      });
      localStorage.setItem(SESSION_KEY, json);
    } catch (e) { /* ignore quota errors */ }
  }, [commentary, noteModes, userNotes, paperNote]); // eslint-disable-line react-hooks/exhaustive-deps

  // ── Persist display preferences ───────────────────────────────────────────
  React.useEffect(() => {
    try { localStorage.setItem(DISPLAY_KEY, JSON.stringify(tweaks)); } catch (e) { /* ignore */ }
  }, [tweaks]);

  // ── CSS custom properties from palette ────────────────────────────────────
  const pal = tweaks.paperPalette || DISPLAY_DEFAULTS.paperPalette;
  const rootStyle = {
    '--paper': pal[0],
    '--ink': pal[1],
    '--rubric': pal[2],
    '--body-size': tweaks.bodySize + 'px',
  };

  // generate: stream an annotation for (targetId, modeId) via the backend.
  // targetId may be a paragraph id OR a figure id — resolved below.
  const generate = React.useCallback(async (targetId, modeId) => {
    const k = `${targetId}:${modeId}`;
    if (inflight.current.has(k)) return;
    if (!result) return;
    inflight.current.add(k);
    setGenerating((g) => ({ ...g, [k]: true }));
    // Resolve target: paragraph first, then figure
    const para = result.paper.paragraphs.find((p) => p.id === targetId);
    let paragraphText, section, litIds;
    if (para) {
      paragraphText = para.text;
      section = para.section;
      litIds = result.litByPara[targetId] || [];
    } else {
      const fig = (result.figures || []).find((f) => f.id === targetId);
      if (!fig) {
        inflight.current.delete(k);
        setGenerating((g) => { const n = { ...g }; delete n[k]; return n; });
        return;
      }
      paragraphText = fig.caption;
      section = fig.label;
      litIds = (result.litByFig || {})[targetId] || [];
    }
    const slicedLitIds = litIds.slice(0, tweaks.litCount);
    const lit = slicedLitIds.map((id) => result.litPapers[id]).filter(Boolean);
    const paraId = targetId;

    let paperContext = null;
    if (modeId === 'referee' && para) {
      const paras = result.paper.paragraphs;
      const idx = paras.findIndex((p) => p.id === targetId);
      paperContext = {
        title: result.paper.title,
        sectionOutline: [...new Set(paras.map((p) => p.section))],
        opening: paras[0]?.text?.slice(0, 400) || '',
        prevParagraph: idx > 0 ? paras[idx - 1].text.slice(0, 300) : '',
        nextParagraph: idx < paras.length - 1 ? paras[idx + 1].text.slice(0, 300) : '',
      };
    }

    try {
      setCommentary((c) => ({ ...c, [k]: { text: '', cites: slicedLitIds, source: 'llm', streaming: true } }));
      for await (const ev of annotateStream({
        paragraphId: paraId,
        paragraph: paragraphText,
        section: section,
        mode: modeId,
        lit: lit.map((l) => ({
          id: l.id,
          short: l.short,
          title: l.title,
          journal: l.journal || 'arXiv',
          abstract: l.abstract,
        })),
        provider: annotator.provider,
        model: annotator.model,
        key: annotator.key,
        paperContext,
      })) {
        if (ev.event === 'token') {
          setCommentary((c) => ({ ...c, [k]: { ...c[k], text: (c[k]?.text || '') + ev.data.text } }));
        } else if (ev.event === 'done') {
          setCommentary((c) => ({ ...c, [k]: { text: ev.data.text, cites: slicedLitIds, source: 'llm' } }));
        } else if (ev.event === 'error') {
          setCommentary((c) => ({ ...c, [k]: { text: '', cites: [], source: 'llm', error: true, message: ev.data.message } }));
        }
      }
    } catch (e) {
      const msg = String(e.message || e).slice(0, 200);
      setCommentary((c) => ({ ...c, [k]: { text: '', cites: [], source: 'llm', error: true, message: msg } }));
      if (e instanceof TypeError || (e.message && e.message.startsWith('Failed to fetch'))) {
        showToast('could not reach the annotator backend');
      }
    } finally {
      inflight.current.delete(k);
      setGenerating((g) => { const n = { ...g }; delete n[k]; return n; });
    }
  }, [result, annotator, tweaks, showToast]);

  // Auto-generate on fresh parse (not session import).
  const fromParseRef = React.useRef(false);
  const prevResultRef = React.useRef(null);

  React.useEffect(() => {
    if (!result || result === prevResultRef.current) return;
    prevResultRef.current = result;
    if (!fromParseRef.current) return;
    result.paper.paragraphs.forEach((p) => {
      generate(p.id, DEFAULT_MODE);
    });
  }, [result, generate]);

  async function onPdf(file) {
    setFileName(file.name);
    setScreen('pipeline');
    setStageEvents([]);
    setResult(null);
    setNoteModes({});
    setCommentary({});
    setGenerating({});
    fromParseRef.current = true;
    try {
      setPdfHash(await hashPdf(await file.arrayBuffer()));
      const jobId = await startParse(file);
      for await (const ev of jobEvents(jobId)) {
        if (ev.event === 'stage') {
          setStageEvents((s) => [...s, ev.data]);
        } else if (ev.event === 'result') {
          const data = ev.data;
          // Prefer backend pdfHash (authoritative); fall back to frontend-computed hash.
          if (data.pdfHash) setPdfHash(data.pdfHash);
          setResult(data);
          const modes = {};
          (data.paper?.paragraphs || []).forEach((p) => { modes[p.id] = DEFAULT_MODE; });
          setNoteModes(modes);
          setCommentary({});
          setGenerating({});
          setScreen('reader');
        } else if (ev.event === 'error') {
          setStageEvents((s) => {
            const failedIndex = s.length ? Math.max(...s.map((x) => x.index)) : 0;
            return [...s, { index: failedIndex, name: '', status: 'failed', log: ev.data.message }];
          });
        }
      }
    } catch (e) {
      setStageEvents((s) => [
        ...s,
        { index: 0, name: '', status: 'failed', log: String(e.message || e) },
      ]);
    }
  }

  function onSetMode(paraId, modeId) {
    setNoteModes((m) => (m[paraId] === modeId ? m : { ...m, [paraId]: modeId }));
    if (!commentaryRef.current[`${paraId}:${modeId}`]) generate(paraId, modeId);
  }

  function onRegenerate(paraId, modeId) {
    generate(paraId, modeId);
  }

  function onSetAllModes(modeId) {
    if (!result?.paper?.paragraphs) return;
    const next = {};
    result.paper.paragraphs.forEach((p) => { next[p.id] = modeId; });
    setNoteModes(next);
    result.paper.paragraphs.forEach((p) => {
      if (!commentaryRef.current[`${p.id}:${modeId}`]) generate(p.id, modeId);
    });
  }

  function onReset() {
    setScreen('upload');
    setResult(null);
    setNoteModes({});
    setCommentary({});
    setGenerating({});
    setFileName('');
    fromParseRef.current = false;
  }

  // ── Save session ───────────────────────────────────────────────────────────
  function onSaveSession() {
    if (!result) return;
    const json = exportSession({
      fileName,
      pdfHash,
      paper: result.paper,
      litPapers: result.litPapers,
      litByPara: result.litByPara,
      figures: result.figures || [],
      litByFig: result.litByFig || {},
      equations: result.equations || [],
      references: result.references || [],
      citeIndex: result.citeIndex || {},
      userNotes,
      paperNote,
      noteModes,
      commentary,
    });
    const blob = new Blob([json], { type: 'application/json' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = fileName.replace(/\.pdf$/i, '') + '.beacon.json';
    a.click();
    URL.revokeObjectURL(url);
    showToast('session written to ‹' + a.download + '›');
  }

  // onOpenLast: restores the full session stored in localStorage on mount.
  function onOpenLast() {
    if (storedSessionRef.current) {
      restoreSession(storedSessionRef.current);
    }
  }

  // Toolbar rendered inside Reader via toolbar prop.
  const toolbar = (
    <div className="ap-toolbar">
      <span className="ap-toolbar-mark">beacon</span>
      <span className="ap-toolbar-file">{fileName}</span>
      <span className="ap-toolbar-spacer"></span>
      <button className="ap-textbtn" onClick={() => setSettingsOpen(true)}>
        annotator: {AP_PROVIDERS[annotator.provider]?.label.toLowerCase() || annotator.provider}
      </button>
      <button className="ap-textbtn" onClick={() => setDisplayOpen(true)}>display</button>
      <button className="ap-toolbar-btn" onClick={() => setNotebookOpen(true)} title="open notebook">&#x1F4D3;</button>
      <button className="ap-textbtn" onClick={onSaveSession}>
        save session &darr;
      </button>
      <button
        className="ap-textbtn"
        onClick={() => importRef.current && importRef.current.click()}
      >
        load
      </button>
      <button className="ap-textbtn" onClick={onReset}>new manuscript</button>
      <input
        ref={importRef}
        type="file"
        accept=".json,application/json"
        style={{ display: 'none' }}
        onChange={(e) => {
          const f = e.target.files[0]; if (!f) return;
          const r = new FileReader();
          r.onload = () => restoreSession(r.result);
          r.readAsText(f);
          e.target.value = '';
        }}
      />
    </div>
  );

  return (
    <div className="ap-root" style={rootStyle}>
      {screen === 'upload' && (
        <Upload
          onPdf={onPdf}
          lastFileName={lastFileName}
          onOpenLast={onOpenLast}
          onSession={restoreSession}
        />
      )}
      {screen === 'pipeline' && (
        <Pipeline
          fileName={fileName}
          events={stageEvents}
          onSkip={result ? () => setScreen('reader') : null}
          onReset={() => setScreen('upload')}
        />
      )}
      {screen === 'reader' && result && (
        <Reader
          paper={result.paper}
          litPapers={result.litPapers || {}}
          litByPara={result.litByPara || {}}
          figures={result.figures || []}
          pdfHash={pdfHash}
          equations={result.equations || []}
          commentary={commentary}
          noteModes={noteModes}
          generating={generating}
          onSetMode={onSetMode}
          onRegenerate={onRegenerate}
          onSetAllModes={onSetAllModes}
          onOpenFigure={(n, e) => {
            const figures = result.figures || [];
            const fig = figures.find((f) => {
              const labelNum = f.label.replace(/^(Figure|Fig\.?)\s*/i, '').trim();
              return labelNum === n;
            });
            if (!fig) return;
            const r = e.currentTarget.getBoundingClientRect();
            setFigPop({ fig, n, pos: { x: r.left, y: r.bottom } });
          }}
          tweaks={tweaks}
          toolbar={toolbar}
          citeIndex={result?.citeIndex || {}}
          references={result?.references || []}
          userNotes={userNotes}
          onSaveNote={handleSaveNote}
        />
      )}
      <Toast msg={toast} />
      {settingsOpen && (
        <AnnotatorPanel
          cfg={annotator}
          onChange={(c) => { setAnnotator(c); apSaveAnnotator(c); }}
          onClose={() => setSettingsOpen(false)}
        />
      )}
      {displayOpen && (
        <DisplayMenu
          tweaks={tweaks}
          onChange={(t) => setTweaks(t)}
          onClose={() => setDisplayOpen(false)}
        />
      )}
      {figPop && (
        <FigurePopover
          fig={figPop.fig}
          n={figPop.n}
          pdfHash={pdfHash}
          pos={figPop.pos}
          commentary={commentary}
          generating={generating}
          onGenerate={generate}
          onClose={() => setFigPop(null)}
        />
      )}
      <NotebookDrawer
        open={notebookOpen}
        onClose={() => setNotebookOpen(false)}
        paper={result?.paper}
        userNotes={userNotes}
        paperNote={paperNote}
        onSavePaperNote={handleSavePaperNote}
        onScrollToPara={(paraId) => {
          document.querySelector(`[data-comment-anchor="para-${paraId}"]`)
            ?.scrollIntoView({ behavior: 'smooth', block: 'center' });
          setNotebookOpen(false);
        }}
      />
    </div>
  );
}