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

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

interface Message {
  id: string;
  role: 'user' | 'assistant';
  text: string;
  // citation indices mentioned in this assistant message (1-based, matching citations array)
  citationRefs?: number[];
}

interface Citation {
  index: number;   // 1-based, matches [N] markers in text
  score: number;   // 0–100 display percentage
  name: string;    // doc_name
  clause: string;  // section_title or clause
  snippet: string; // preview text
  docId?: string;
}

// Map a raw source doc from the backend "retrieved" event to our Citation shape.
// Backend fields: { id, score(0-1), preview, doc_name, clause, doc_id }
function mapSource(s: Record<string, unknown>, idx: number): Citation {
  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() {
  const [messages, setMessages] = useState<Message[]>([
    {
      id: 'init', role: 'assistant',
      text: 'Hello! I can answer questions about your indexed regulations and compliance documents. Try asking about EU AI Act requirements, MIIT rules, or ISO/SAE 21434 scope.',
    }
  ]);
  const [quickPrompts, setQuickPrompts] = useState<string[]>(MOCK_QUICK);
  const [input, setInput] = useState('');
  const [streaming, setStreaming] = useState(false);
  const [citations, setCitations] = useState<Citation[]>([]);
  const [highlightedCit, setHighlightedCit] = useState<number | null>(null);
  const [sessionId, setSessionId] = useState<string | null>(null);

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

  // Fetch quick questions from backend on mount
  useEffect(() => {
    fetch('/api/v1/rag/quick-questions')
      .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' });
    }
    // Clear highlight after 3s
    setTimeout(() => setHighlightedCit(h => h === n ? null : h), 3000);
  }, []);

  async function send(text?: string) {
    const q = (text ?? input).trim();
    if (!q || streaming) return;
    setInput('');

    const userMsg: Message = { id: Date.now().toString(), role: 'user', text: q };
    setMessages(m => [...m, userMsg]);

    const assistantId = (Date.now() + 1).toString();
    setMessages(m => [...m, { id: assistantId, role: 'assistant', text: '' }]);
    setStreaming(true);
    setCitations([]);
    setHighlightedCit(null);

    const ctrl = new AbortController();
    abortRef.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' },
        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: Citation[] = [];

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

        // SSE blocks separated by double newline
        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') {
              // Backend assigned a session_id — persist for next request
              if (j.session_id) setSessionId(j.session_id);

            } else if (j.type === 'retrieved' && Array.isArray(j.docs)) {
              // Sources arrive before the answer starts
              const mapped = j.docs.map((d: Record<string, unknown>, i: number) => mapSource(d, i + 1));
              newCitations.push(...mapped);
              setCitations([...mapped]);

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

            } else if (j.type === 'status') {
              // Status message (e.g. "找到N条相关法规…") — could show in UI if desired
              // For now we ignore it to keep the bubble clean

            } else if (j.type === 'done') {
              // Extract which citation numbers appear in the final answer
              setMessages(m => m.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') {
              setMessages(m => m.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') {
        setMessages(m => m.map(msg =>
          msg.id === assistantId
            ? { ...msg, text: 'Could not reach the RAG API. Please check the backend.' }
            : msg
        ));
      }
    } finally {
      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={input}
                onChange={e => setInput(e.target.value)}
                onKeyDown={e => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send(); } }}
                rows={2}
              />
              <button
                className="btn primary"
                onClick={() => send()}
                disabled={!input.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>
  );
}