AIRegulation-DocAnalysis/frontend/src/pages/RagChat/RagChatPage.tsx

import { useRef, useEffect, useCallback, useState } from 'react';
import { Topbar } from '../../components/layout/Topbar';
import { Send, Download } from 'lucide-react';
import { usePageState } from '../../contexts';
import type { RagCitation } from '../../contexts';

const TOKEN_KEY = 'auth_token';
function authHeader(): Record<string, string> {
  const t = localStorage.getItem(TOKEN_KEY);
  return t ? { Authorization: `Bearer ${t}` } : {};
}

// Map a raw source doc from the backend "retrieved" event to our Citation shape.
function mapSource(s: Record<string, unknown>, idx: number): RagCitation {
  const rawScore = typeof s.score === 'number' ? s.score : 0;
  const displayScore = rawScore <= 1 ? Math.round(rawScore * 100) : Math.round(rawScore);
  return {
    index: idx,
    score: displayScore,
    name: String(s.doc_name ?? ''),
    clause: String(s.clause ?? s.section_title ?? ''),
    snippet: String(s.preview ?? s.text ?? ''),
    docId: s.doc_id ? String(s.doc_id) : undefined,
  };
}

// Parse message text and replace [N] with clickable <button class="cite-ref"> elements.
function renderWithCitations(
  text: string,
  onCiteClick: (n: number) => void,
): React.ReactNode[] {
  const parts = text.split(/(\[\d+\])/g);
  return parts.map((part, i) => {
    const m = part.match(/^\[(\d+)\]$/);
    if (m) {
      const n = parseInt(m[1], 10);
      return (
        <button
          key={i}
          className="cite-ref"
          onClick={() => onCiteClick(n)}
          title={`Jump to source [${n}]`}
        >
          {n}
        </button>
      );
    }
    return part;
  });
}

const MOCK_QUICK = [
  'What does EU AI Act Art. 9 require for risk management?',
  'Which documents need CSMS certification?',
  'Summarize MIIT training data rules',
  'What are high-risk AI categories under Annex III?',
];

export function RagChatPage() {
  // All persistent state lives in PageStateContext — survives route changes
  const { ragState, setRagState, ragStreamingRef, ragAbortRef } = usePageState();
  const { messages, citations, sessionId, inputDraft } = ragState;

  // Local-only UI state: highlighted citation and streaming indicator
  // These are fine to reset on navigation since they're transient UI feedback
  const [highlightedCit, setHighlightedCit] = useState<number | null>(null);
  const [streaming, setStreaming] = useState(ragStreamingRef.current);
  const [quickPrompts, setQuickPrompts] = useState<string[]>(MOCK_QUICK);

  const bottomRef = useRef<HTMLDivElement>(null);
  const citRailRef = useRef<HTMLDivElement>(null);
  const citItemRefs = useRef<Record<number, HTMLDivElement | null>>({});

  // Fetch quick questions from backend on mount (only once per session)
  useEffect(() => {
    fetch('/api/v1/rag/quick-questions', { headers: authHeader() })
      .then(r => r.json())
      .then(d => {
        if (Array.isArray(d?.questions) && d.questions.length > 0) {
          setQuickPrompts(d.questions.map((q: { question: string }) => q.question));
        }
      })
      .catch(() => { /* keep mock */ });
  }, []);

  // Auto-scroll to latest message
  useEffect(() => {
    bottomRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  // Jump to citation N in the rail and highlight it
  const jumpToCitation = useCallback((n: number) => {
    setHighlightedCit(n);
    const el = citItemRefs.current[n];
    if (el) {
      el.scrollIntoView({ behavior: 'smooth', block: 'nearest' });
    }
    setTimeout(() => setHighlightedCit(h => h === n ? null : h), 3000);
  }, []);

  async function send(text?: string) {
    const q = (text ?? inputDraft).trim();
    if (!q || ragStreamingRef.current) return;
    setRagState(s => ({ ...s, inputDraft: '' }));

    const userMsgId = Date.now().toString();
    const assistantId = (Date.now() + 1).toString();

    setRagState(s => ({
      ...s,
      messages: [
        ...s.messages,
        { id: userMsgId, role: 'user', text: q },
        { id: assistantId, role: 'assistant', text: '' },
      ],
      citations: [],
    }));

    ragStreamingRef.current = true;
    setStreaming(true);
    setHighlightedCit(null);

    const ctrl = new AbortController();
    ragAbortRef.current = ctrl;

    try {
      const body: Record<string, unknown> = { query: q, top_k: 5 };
      if (sessionId) body.session_id = sessionId;

      const res = await fetch('/api/v1/rag/chat', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json', ...authHeader() },
        body: JSON.stringify(body),
        signal: ctrl.signal,
      });

      if (!res.body) throw new Error('No stream');
      const reader = res.body.getReader();
      const dec = new TextDecoder();
      let buffer = '';
      const newCitations: RagCitation[] = [];

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += dec.decode(value, { stream: true });

        const blocks = buffer.split('\n\n');
        buffer = blocks.pop() ?? '';

        for (const block of blocks) {
          const dataLine = block.split('\n').find(l => l.startsWith('data: '));
          if (!dataLine) continue;
          const raw = dataLine.slice(6).trim();
          if (!raw) continue;
          try {
            const j = JSON.parse(raw);

            if (j.type === 'session') {
              if (j.session_id) setRagState(s => ({ ...s, sessionId: j.session_id }));

            } else if (j.type === 'retrieved' && Array.isArray(j.docs)) {
              const mapped = j.docs.map((d: Record<string, unknown>, i: number) => mapSource(d, i + 1));
              newCitations.push(...mapped);
              setRagState(s => ({ ...s, citations: [...mapped] }));

            } else if (j.type === 'chunk' && j.text) {
              setRagState(s => ({
                ...s,
                messages: s.messages.map(msg =>
                  msg.id === assistantId
                    ? { ...msg, text: msg.text + (j.text as string) }
                    : msg
                ),
              }));

            } else if (j.type === 'done') {
              setRagState(s => ({
                ...s,
                messages: s.messages.map(msg => {
                  if (msg.id !== assistantId) return msg;
                  const refs = [...new Set(
                    [...msg.text.matchAll(/\[(\d+)\]/g)].map(r => parseInt(r[1], 10))
                  )].filter(n => n >= 1 && n <= newCitations.length);
                  return { ...msg, citationRefs: refs };
                }),
              }));
              break;

            } else if (j.type === 'error') {
              setRagState(s => ({
                ...s,
                messages: s.messages.map(msg =>
                  msg.id === assistantId
                    ? { ...msg, text: `Error: ${j.text ?? 'Unknown error'}` }
                    : msg
                ),
              }));
            }
          } catch { /* malformed JSON chunk, skip */ }
        }
      }
    } catch (e: unknown) {
      if (e instanceof Error && e.name !== 'AbortError') {
        setRagState(s => ({
          ...s,
          messages: s.messages.map(msg =>
            msg.id === assistantId
              ? { ...msg, text: 'Could not reach the RAG API. Please check the backend.' }
              : msg
          ),
        }));
      }
    } finally {
      ragStreamingRef.current = false;
      setStreaming(false);
    }
  }

  const lastAssistantId = [...messages].reverse().find(m => m.role === 'assistant')?.id;

  return (
    <div className="chat-page">
      <Topbar
        title="Regulation Q&A"
        actions={
          <button
            className="btn sm"
            onClick={() => {
              const text = messages.map(m => `${m.role === 'user' ? 'Q' : 'A'}: ${m.text}`).join('\n\n');
              const blob = new Blob([text], { type: 'text/plain' });
              const url = URL.createObjectURL(blob);
              const a = document.createElement('a'); a.href = url; a.download = 'chat-export.txt'; a.click();
              URL.revokeObjectURL(url);
            }}
          >
            <Download size={13} />Export chat
          </button>
        }
      />

      <div className="chat-body">
        {/* ── History pane ── */}
        <div className="history-pane">
          <div className="history-header">Quick prompts</div>
          {quickPrompts.map(q => (
            <button key={q} className="quick-item" onClick={() => send(q)}>
              {q}
            </button>
          ))}
        </div>

        {/* ── Chat main ── */}
        <div className="chat-main">
          <div className="messages">
            {messages.map(msg => (
              <div key={msg.id} className={`message msg-${msg.role}`}>
                {msg.role === 'assistant' && <div className="msg-avatar">AI</div>}
                <div className="msg-bubble">
                  {msg.role === 'assistant'
                    ? renderWithCitations(msg.text, jumpToCitation)
                    : msg.text
                  }
                  {streaming && msg.id === lastAssistantId && (
                    <span className="blink-cursor">▋</span>
                  )}
                </div>
                {msg.role === 'user' && <div className="msg-avatar user-av">You</div>}
              </div>
            ))}
            <div ref={bottomRef} />
          </div>

          <div className="composer">
            <div className="quick-chips">
              {quickPrompts.slice(0, 3).map(q => (
                <button key={q} className="chip" onClick={() => send(q)}>
                  {q.length > 42 ? q.slice(0, 42) + '…' : q}
                </button>
              ))}
            </div>
            <div className="composer-row">
              <textarea
                className="composer-input"
                placeholder="Ask about your regulations…"
                value={inputDraft}
                onChange={e => setRagState(s => ({ ...s, inputDraft: e.target.value }))}
                onKeyDown={e => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send(); } }}
                rows={2}
              />
              <button
                className="btn primary"
                onClick={() => send()}
                disabled={!inputDraft.trim() || streaming}
              >
                <Send size={14} />
              </button>
            </div>
          </div>
        </div>

        {/* ── Citation rail ── */}
        <div className="citation-rail" ref={citRailRef}>
          <div className="citation-header">
            Sources {citations.length > 0 && `(${citations.length})`}
          </div>
          {citations.length === 0 && (
            <p style={{ padding: '12px 16px', fontSize: 12, color: 'var(--muted)', lineHeight: 1.5 }}>
              Citations will appear here after a response is generated.
            </p>
          )}
          {citations.map(c => (
            <div
              key={c.index}
              ref={el => { citItemRefs.current[c.index] = el; }}
              className={`citation-item${highlightedCit === c.index ? ' highlighted' : ''}`}
            >
              <div className="cit-index">[{c.index}]</div>
              <div style={{ flex: 1, minWidth: 0 }}>
                <div style={{ display: 'flex', alignItems: 'baseline', gap: 4, marginBottom: 3 }}>
                  <div className="cit-name">{c.name}</div>
                  {c.clause && <span className="cit-clause">{c.clause}</span>}
                  <span className="cit-score" style={{ marginLeft: 'auto' }}>{c.score}%</span>
                </div>
                <div className="cit-snippet">{c.snippet}</div>
              </div>
            </div>
          ))}
        </div>
      </div>
    </div>
  );
}