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

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

interface Message {
  id: string;
  role: 'user' | 'assistant';
  text: string;
}

interface Citation {
  score: number;
  name: string;
  clause: string;
  snippet: string;
}

const HISTORY = [
  { id: 'h1', title: 'EU AI Act Article 9 scope', date: '2025-11-18' },
  { id: 'h2', title: 'MIIT training data requirements', date: '2025-11-15' },
  { id: 'h3', title: 'ISO 21434 CSMS audit scope', date: '2025-11-10' },
];

const 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?',
];

const MOCK_CITATIONS: Citation[] = [
  {
    score: 94, name: 'EU AI Act', clause: 'Art. 9(1)',
    snippet: 'Providers of high-risk AI systems shall establish a risk management system consisting of a continuous iterative process run throughout the entire lifecycle.'
  },
  {
    score: 87, name: 'Vehicle AI Safety Manual', clause: '§4.2.1',
    snippet: 'All AI systems classified as high-risk must maintain a documented risk register with quarterly review cadence.'
  },
  {
    score: 72, name: 'ISO/SAE 21434', clause: 'Clause 9.3',
    snippet: 'The cybersecurity management system shall include AI model update governance procedures and audit log retention policy.'
  },
];

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 [input, setInput] = useState('');
  const [streaming, setStreaming] = useState(false);
  const [citations, setCitations] = useState<Citation[]>(MOCK_CITATIONS);
  const bottomRef = useRef<HTMLDivElement>(null);
  const abortRef = useRef<AbortController | null>(null);

  useEffect(() => {
    bottomRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  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);

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

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

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

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += dec.decode(value);
        const lines = buffer.split('\n');
        buffer = lines.pop() ?? '';
        for (const line of lines) {
          if (line.startsWith('data: ')) {
            const data = line.slice(6);
            if (data === '[DONE]') break;
            try {
              const j = JSON.parse(data);
              if (j.text) setMessages(m => m.map(msg =>
                msg.id === assistantId ? { ...msg, text: msg.text + j.text } : msg
              ));
              if (j.citations) setCitations(j.citations);
            } catch {
              setMessages(m => m.map(msg =>
                msg.id === assistantId ? { ...msg, text: msg.text + data } : msg
              ));
            }
          }
        }
      }
    } 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"><Download size={13} />Export chat</button>}
      />
      <div className="chat-body">
        <div className="history-pane">
          <div className="history-header">Chat history</div>
          {HISTORY.map(h => (
            <div key={h.id} className="history-item">
              <div className="history-title">{h.title}</div>
              <div className="history-date">{h.date}</div>
            </div>
          ))}
          <div className="quick-header">Quick prompts</div>
          {QUICK.map(q => (
            <button key={q} className="quick-item" onClick={() => send(q)}>{q}</button>
          ))}
        </div>

        <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.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">
              {QUICK.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>

        <div className="citation-rail">
          <div className="citation-header">Sources</div>
          {citations.map(c => (
            <div key={`${c.name}-${c.clause}`} className="citation-item">
              <div className="cit-score">{c.score}%</div>
              <div>
                <div className="cit-name">{c.name} <span className="cit-clause">{c.clause}</span></div>
                <div className="cit-snippet">{c.snippet}</div>
              </div>
            </div>
          ))}
        </div>
      </div>
    </div>
  );
}