/* ============================================
   WORKSPACE — Chat (RAG) + Document Viewer with bbox overlay
   ============================================ */
function WorkspaceScreen({ onNavigate }) {
  const [input, setInput] = useState("");
  const [sending, setSending] = useState(false);
  const [messages, setMessages] = useState([]);
  const [citations, setCitations] = useState([]);
  const [activeIndex, setActiveIndex] = useState(null);
  const [error, setError] = useState(null);
  const scrollRef = useRef(null);

  useEffect(() => {
    if (scrollRef.current) scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
  }, [messages]);

  // Auto-select citation [1] whenever a new set arrives.
  useEffect(() => {
    if (citations.length > 0) setActiveIndex(1);
    else setActiveIndex(null);
  }, [citations]);

  const send = async () => {
    const prompt = input.trim();
    if (!prompt || sending) return;
    setError(null);
    setInput("");
    setSending(true);
    setMessages((m) => [...m, { role: "user", text: prompt }, { role: "assistant", text: "" }]);
    setCitations([]);

    try {
      const res = await window.API.chat(prompt);
      const reader = res.body.getReader();
      const decoder = new TextDecoder();
      let buf = "";
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buf += decoder.decode(value, { stream: true });
        const events = buf.split("\n\n");
        buf = events.pop();
        for (const evt of events) {
          if (evt.startsWith("event: done")) continue;
          if (!evt.startsWith("data: ")) continue;
          const dataStr = evt.slice(6);
          if (dataStr === "[DONE]") continue;
          let payload;
          try { payload = JSON.parse(dataStr); } catch (_) { continue; }
          if (payload.citations) {
            setCitations(payload.citations);
          } else if (payload.token) {
            setMessages((m) => {
              const next = [...m];
              const last = next[next.length - 1];
              next[next.length - 1] = { ...last, text: last.text + payload.token };
              return next;
            });
          } else if (payload.error) {
            setError(payload.error);
          }
        }
      }
    } catch (err) {
      setError(err.message || "Chat failed.");
    } finally {
      setSending(false);
    }
  };

  const reset = () => {
    setMessages([]);
    setCitations([]);
    setActiveIndex(null);
    setError(null);
  };

  const activeCitation = citations.find((c) => c.index === activeIndex) || null;

  return (
    <div style={{
      display: "grid",
      gridTemplateColumns: "1fr 1fr",
      gap: 16,
      height: "calc(100vh - var(--topbar-h) - 48px)",
    }}>
      {/* ===== LEFT: Chat panel ===== */}
      <div className="card" style={{ display: "flex", flexDirection: "column", overflow: "hidden" }}>
        <div className="card__header">
          <div>
            <h3 className="card__title">
              <Icons.Chat size={16} />
              Ask the indexed manuals
            </h3>
            <p className="card__subtitle mono">
              {messages.length === 0 ? "No messages yet" : `${messages.length} message${messages.length === 1 ? "" : "s"} · ${citations.length} sources in scope`}
            </p>
          </div>
          <div className="row">
            <button className="btn btn--ghost btn--sm" onClick={reset} disabled={messages.length === 0 && !sending}>
              <Icons.Plus size={14} /> New chat
            </button>
          </div>
        </div>

        <div ref={scrollRef} style={{ flex: 1, overflowY: "auto", padding: "20px 22px", display: "flex", flexDirection: "column", gap: 18 }}>
          {messages.length === 0 && (
            <div style={{ textAlign: "center", color: "var(--ink-500)", padding: "60px 20px" }}>
              <div style={{ fontSize: 14, marginBottom: 8 }}>Ask a question about your ingested manuals.</div>
              <div className="mono text-xs muted">Citations appear as clickable pills — click one to see the source page.</div>
            </div>
          )}

          {messages.map((msg, i) => (
            <ChatBubble
              key={i}
              role={msg.role}
              name={msg.role === "user" ? "You" : "ORS Assistant"}
              model={msg.role === "assistant" ? "qwen-plus" : null}
            >
              {msg.role === "assistant" ? (
                <>
                  <MarkdownWithCitations text={msg.text} citations={citations} onCite={setActiveIndex} />
                  {!msg.text && sending && i === messages.length - 1 && <span className="muted">…</span>}
                </>
              ) : msg.text}
            </ChatBubble>
          ))}

          {error && (
            <div style={{
              padding: "10px 14px",
              background: "rgba(220, 38, 38, 0.08)",
              border: "1px solid rgba(220, 38, 38, 0.35)",
              borderRadius: 6, color: "#991b1b",
              fontSize: 13,
            }}>{error}</div>
          )}
        </div>

        {/* Composer */}
        <div style={{ padding: 14, borderTop: "1px solid var(--ink-200)", background: "var(--ink-50)" }}>
          <div style={{ background: "white", border: "1px solid var(--ink-300)", borderRadius: 8, padding: "10px 12px" }}>
            <textarea
              value={input}
              onChange={(e) => setInput(e.target.value)}
              placeholder="Ask a question about an indexed manual…"
              style={{ width: "100%", border: "none", outline: "none", resize: "none", fontFamily: "inherit", fontSize: 13.5, minHeight: 40, lineHeight: 1.5 }}
              onKeyDown={(e) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); send(); } }}
              disabled={sending}
            />
            <div style={{ display: "flex", gap: 6, alignItems: "center", marginTop: 6 }}>
              <span className="spacer" />
              <span className="text-xs muted mono">{input.length}/4096</span>
              <button className="btn btn--primary btn--sm" onClick={send} disabled={!input.trim() || sending}>
                <Icons.Send size={14} /> {sending ? "Streaming…" : "Send"}
              </button>
            </div>
          </div>
          <div className="text-xs muted" style={{ marginTop: 6, paddingLeft: 4 }}>
            Responses grounded in the indexed manuals. Verify safety-critical values against the printed copy.
          </div>
        </div>
      </div>

      {/* ===== RIGHT: Document Viewer ===== */}
      <div className="card" style={{ display: "flex", flexDirection: "column", overflow: "hidden" }}>
        <div className="card__header">
          <div style={{ minWidth: 0, flex: 1 }}>
            <h3 className="card__title">
              <Icons.File size={16} />
              Document Viewer
            </h3>
            <p className="card__subtitle" style={{ overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap" }}>
              {activeCitation
                ? <span><span className="mono" style={{ color: "var(--orange-700)" }}>[{activeCitation.index}]</span> <span className="muted">{activeCitation.source} · p.{activeCitation.page}</span></span>
                : <span className="muted">Send a question to see the cited page.</span>}
            </p>
          </div>
        </div>

        <div style={{ flex: 1, overflowY: "auto", padding: 18 }}>
          {citations.length === 0 && (
            <div style={{ textAlign: "center", color: "var(--ink-500)", padding: "60px 20px" }}>
              <div className="mono text-xs muted">Cited pages will render here after each query.</div>
            </div>
          )}

          {activeCitation && (
            <DocumentPage citation={activeCitation} />
          )}

          {citations.length > 0 && (
            <div style={{ marginTop: 16 }}>
              <div className="text-xs mono muted" style={{ textTransform: "uppercase", letterSpacing: "0.06em", fontWeight: 600, marginBottom: 8 }}>
                All retrieved sources ({citations.length})
              </div>
              {citations.map((c) => (
                <button
                  key={c.index}
                  onClick={() => setActiveIndex(c.index)}
                  style={{
                    display: "block",
                    width: "100%",
                    textAlign: "left",
                    padding: "8px 12px",
                    marginBottom: 6,
                    border: c.index === activeIndex ? "1px solid var(--orange-500)" : "1px solid var(--ink-200)",
                    borderRadius: 6,
                    background: c.index === activeIndex ? "var(--orange-50)" : "white",
                    cursor: "pointer",
                    fontSize: 12.5,
                  }}
                >
                  <div className="row" style={{ justifyContent: "space-between", marginBottom: 2 }}>
                    <span className="mono" style={{ color: "var(--orange-700)", fontWeight: 600 }}>[{c.index}] p.{c.page}</span>
                    <span className="mono text-xs muted">score {c.score}</span>
                  </div>
                  <div style={{ overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap", color: "var(--ink-700)" }}>{c.source}</div>
                </button>
              ))}
            </div>
          )}
        </div>
      </div>
    </div>
  );
}

function DocumentPage({ citation }) {
  const { task_id, page, bbox, text } = citation;
  const imgUrl = task_id ? `/api/document/${task_id}/page/${page}.png` : null;
  const token = window.API.getToken();
  // FileResponse via FastAPI honors Bearer in standard fetch, but <img src=...>
  // can't attach Authorization headers. Workaround: use a fetch + blob URL.
  const [blobUrl, setBlobUrl] = useState(null);
  const [imgError, setImgError] = useState(null);

  useEffect(() => {
    if (!imgUrl) return;
    let cancel = false;
    let createdUrl = null;
    (async () => {
      try {
        const res = await fetch(imgUrl, { headers: token ? { Authorization: `Bearer ${token}` } : {} });
        if (!res.ok) throw new Error(`Page image not found (${res.status})`);
        const blob = await res.blob();
        if (cancel) return;
        createdUrl = URL.createObjectURL(blob);
        setBlobUrl(createdUrl);
      } catch (err) {
        if (!cancel) setImgError(err.message);
      }
    })();
    return () => {
      cancel = true;
      if (createdUrl) URL.revokeObjectURL(createdUrl);
    };
  }, [imgUrl, token]);

  return (
    <div>
      <div style={{ position: "relative", border: "1px solid var(--ink-200)", borderRadius: 6, overflow: "hidden", background: "var(--ink-50)" }}>
        {imgError && (
          <div style={{ padding: 24, textAlign: "center", color: "var(--ink-500)", fontSize: 13 }}>{imgError}</div>
        )}
        {!imgError && !blobUrl && (
          <div style={{ padding: 60, textAlign: "center", color: "var(--ink-500)", fontSize: 13 }}>Loading page…</div>
        )}
        {blobUrl && (
          <div style={{ position: "relative", display: "block" }}>
            <img src={blobUrl} alt={`Page ${page}`} style={{ display: "block", width: "100%", height: "auto" }} />
            {bbox && (
              <div style={{
                position: "absolute",
                left: `${bbox.x0 * 100}%`,
                top: `${bbox.y0 * 100}%`,
                width: `${(bbox.x1 - bbox.x0) * 100}%`,
                height: `${(bbox.y1 - bbox.y0) * 100}%`,
                border: "2px solid #DC2626",
                background: "rgba(220, 38, 38, 0.05)",
                pointerEvents: "none",
              }} />
            )}
          </div>
        )}
      </div>

      {text && (
        <div style={{ marginTop: 12, padding: 12, background: "rgba(220, 38, 38, 0.04)", border: "1px solid rgba(220, 38, 38, 0.25)", borderRadius: 6 }}>
          <div className="text-xs mono" style={{ color: "#991b1b", fontWeight: 600, marginBottom: 6 }}>MATCHED TEXT</div>
          <div style={{ fontSize: 12.5, lineHeight: 1.55, color: "var(--ink-800)", whiteSpace: "pre-wrap", wordBreak: "break-word" }}>{text}</div>
        </div>
      )}
    </div>
  );
}

// Renders assistant Markdown text with [file.pdf p.N] citation markers turned
// into clickable orange pills that switch the Document Viewer.
const CITE_PATTERN = /\[([^\[\]]+?\.pdf)\s+p\.([\d.-]+)\]/g;

function MarkdownWithCitations({ text, citations, onCite }) {
  const containerRef = React.useRef(null);
  const citeMapRef = React.useRef([]);

  const html = React.useMemo(() => {
    if (!text) return "";
    citeMapRef.current = [];

    // Replace citation markers with styled <button> placeholders.
    // DOMPurify allows <button> + data-* by default; we add click via delegation.
    const withButtons = text.replace(CITE_PATTERN, (_, source, pageStr) => {
      const page = pageStr;
      const cite = citations.find(
        (c) => c.source === source.trim() && String(c.page) === String(page)
      );
      const idx = citeMapRef.current.length;
      citeMapRef.current.push({ cite });
      const cls = cite ? "cite-pill" : "cite-pill cite-pill--missing";
      return `<button class="${cls}" data-cite-idx="${idx}">p.${page}</button>`;
    });

    const markedLib = window.marked;
    const purifyLib = window.DOMPurify;

    let rawHtml;
    if (markedLib) {
      rawHtml = markedLib.parse(withButtons);
    } else {
      // Fallback: basic newline → <br> if marked hasn't loaded yet
      rawHtml = withButtons.replace(/\n/g, "<br>");
    }

    return purifyLib
      ? purifyLib.sanitize(rawHtml, { ADD_ATTR: ["data-cite-idx"] })
      : rawHtml;
  }, [text, citations]);

  // Delegated click handler — DOMPurify strips onclick, so we attach in JS.
  React.useEffect(() => {
    const el = containerRef.current;
    if (!el) return;
    const handler = (e) => {
      const btn = e.target.closest("[data-cite-idx]");
      if (!btn) return;
      const idx = parseInt(btn.getAttribute("data-cite-idx"), 10);
      const entry = citeMapRef.current[idx];
      if (entry && entry.cite) onCite(entry.cite.index);
    };
    el.addEventListener("click", handler);
    return () => el.removeEventListener("click", handler);
  }, [onCite]);

  if (!html) return null;
  return <div ref={containerRef} className="chat-md" dangerouslySetInnerHTML={{ __html: html }} />;
}

function ChatBubble({ role, name, model, children }) {
  const isUser = role === "user";
  return (
    <div style={{ display: "flex", gap: 12, alignItems: "flex-start" }}>
      <div style={{
        width: 32, height: 32, borderRadius: 6,
        background: isUser ? "var(--navy-600)" : "linear-gradient(135deg, var(--orange-500), var(--orange-700))",
        color: "white", display: "grid", placeItems: "center",
        fontWeight: 600, fontSize: 11, flexShrink: 0,
        fontFamily: "var(--font-mono)",
      }}>{isUser ? "YOU" : "ORS"}</div>
      <div style={{ flex: 1, minWidth: 0 }}>
        <div className="row gap-8" style={{ marginBottom: 6 }}>
          <span style={{ fontWeight: 600, fontSize: 13 }}>{name}</span>
          {model && <span className="mono text-xs muted">{model}</span>}
        </div>
        <div style={{
          fontSize: 13.5,
          lineHeight: 1.65,
          background: isUser ? "var(--navy-50)" : "white",
          border: isUser ? "1px solid var(--navy-100)" : "1px solid var(--ink-200)",
          borderRadius: 8,
          padding: "12px 14px",
          color: "var(--ink-900)",
          whiteSpace: isUser ? "pre-wrap" : undefined,
          wordWrap: "break-word",
          overflowX: "auto",
        }}>{children}</div>
      </div>
    </div>
  );
}

window.WorkspaceScreen = WorkspaceScreen;