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

import React, { useCallback, useEffect, useRef, useState } from 'react';
import { useTheme } from '../../contexts';
import type { ChatMessage, RetrievalData } from '../../types';
import { getQuickQuestions, ragChat } from '../../api/rag';
import { CitedAnswer } from './CitedAnswer';

const ragQuickQuestionsDefault = [
  '电动自行车上路需要什么条件？',
  '驾驶证如何申请？',
  '超速行驶如何处罚？',
  '车辆年检有哪些规定？',
  '电动汽车电池安全标准？',
  '正面碰撞测试要求？',
  'AEB系统测试标准？',
  '高速公路安全距离？',
];

export const RagChatPage: React.FC = () => {
  const { theme } = useTheme();
  const nextMessageIdRef = useRef(1);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  // retrievals: right-panel shows sources of the most recent assistant reply
  const [retrievals, setRetrievals] = useState<RetrievalData[]>([]);
  const [input, setInput] = useState<string>('');
  const [loading, setLoading] = useState<boolean>(false);
  const [showClearConfirm, setShowClearConfirm] = useState<boolean>(false);
  const [selectedRetrieval, setSelectedRetrieval] = useState<RetrievalData | null>(null);
  const [quickQuestions, setQuickQuestions] = useState<string[]>(ragQuickQuestionsDefault);
  const [filterRegulationType, setFilterRegulationType] = useState<string>('');
  const [highlightedSourceIdx, setHighlightedSourceIdx] = useState<number | null>(null);
  const [sessionId, setSessionId] = useState<string | undefined>();

  // Auto-scroll ref
  const messagesEndRef = useRef<HTMLDivElement>(null);
  // AbortController for cancelling in-flight requests
  const abortRef = useRef<AbortController | null>(null);

  function nextMessageId() {
    const id = nextMessageIdRef.current;
    nextMessageIdRef.current += 1;
    return id;
  }

  // Scroll to bottom whenever messages change
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);

  async function loadQuickQuestions() {
    try {
      const response = await getQuickQuestions();
      setQuickQuestions(response.questions.map(q => q.question));
    } catch {
      // keep defaults
    }
  }

  useEffect(() => {
    const timerId = window.setTimeout(() => { void loadQuickQuestions(); }, 0);
    return () => window.clearTimeout(timerId);
  }, []);

  /**
   * Core query executor — shared by sendMessage and regenerateLastAnswer.
   * Manages session_id, AbortController, SSE parsing, and state updates.
   */
  const executeQuery = useCallback((text: string) => {
    // Cancel any in-flight request
    abortRef.current?.abort();
    abortRef.current = new AbortController();

    const activeFilters = filterRegulationType.trim() || undefined;
    let currentResponse = '';
    // Capture the assistant message id so we can attach sources later
    let assistantMsgId: number | null = null;

    void ragChat(
      text,
      5,
      (data) => {
        if (data.type === 'session' && data.session_id) {
          setSessionId(data.session_id);
        } else if (data.type === 'retrieved' && data.docs) {
          const docs: RetrievalData[] = data.docs.map(d => ({
            id: parseInt(d.id.replace('chunk-', ''), 10) || 1,
            file: d.doc_name,
            clause: d.clause,
            score: d.score,
            content: d.preview,
            docId: d.doc_id,
            downloadUrl: d.download_url,
          }));
          setRetrievals(docs);
          // Attach sources to the assistant message once we know its id
          if (assistantMsgId !== null) {
            setMessages(prev => prev.map(m =>
              m.id === assistantMsgId ? { ...m, sources: docs } : m
            ));
          }
        } else if (data.type === 'chunk' && data.text) {
          currentResponse += data.text;
          setMessages(prev => {
            const last = prev[prev.length - 1];
            if (last?.role === 'assistant' && last.id === assistantMsgId) {
              return [...prev.slice(0, -1), { ...last, content: currentResponse }];
            }
            // First chunk: create assistant message
            const newId = nextMessageId();
            assistantMsgId = newId;
            return [...prev, { id: newId, role: 'assistant' as const, content: currentResponse }];
          });
        } else if (data.type === 'done') {
          if (data.session_id) setSessionId(data.session_id);
          setLoading(false);
        } else if (data.type === 'error') {
          setLoading(false);
          setMessages(prev => [
            ...prev,
            { id: nextMessageId(), role: 'assistant' as const, content: '抱歉，生成回答时出错，请稍后再试。' },
          ]);
        }
      },
      (error) => {
        console.error('RAG chat error:', error);
        setLoading(false);
        setMessages(prev => [
          ...prev,
          { id: nextMessageId(), role: 'assistant' as const, content: '抱歉，连接服务器时出错，请稍后再试。' },
        ]);
      },
      () => { setLoading(false); },
      activeFilters,
      sessionId,
      abortRef.current.signal,
    );
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [filterRegulationType, sessionId]);

  const sendMessage = (text: string) => {
    if (!text.trim() || loading) return;
    setMessages(prev => [...prev, { id: nextMessageId(), role: 'user' as const, content: text }]);
    setInput('');
    setLoading(true);
    setRetrievals([]);
    setHighlightedSourceIdx(null);
    executeQuery(text);
  };

  const regenerateLastAnswer = () => {
    if (loading) return;
    const lastUserMsg = [...messages].reverse().find(m => m.role === 'user');
    if (!lastUserMsg) return;
    // Remove the last assistant message
    setMessages(prev => {
      const lastAssistantIdx = [...prev].reverse().findIndex(m => m.role === 'assistant');
      if (lastAssistantIdx === -1) return prev;
      const idx = prev.length - 1 - lastAssistantIdx;
      return [...prev.slice(0, idx)];
    });
    setLoading(true);
    setRetrievals([]);
    setHighlightedSourceIdx(null);
    executeQuery(lastUserMsg.content);
  };

  const clearMessages = () => {
    abortRef.current?.abort();
    setMessages([]);
    setRetrievals([]);
    setSessionId(undefined);
    setShowClearConfirm(false);
    setLoading(false);
  };

  return (
    <div style={{ flex: 1, display: 'flex', height: 'calc(100vh - 128px)' }}>
      {/* ── Left: chat panel ─────────────────────────────────── */}
      <div style={{
        flex: '0 0 60%',
        display: 'flex',
        flexDirection: 'column',
        borderRight: `1px solid ${theme.border}`,
      }}>
        {/* Message list */}
        <div style={{
          flex: 1,
          overflowY: 'auto',
          padding: '24px 32px',
          display: 'flex',
          flexDirection: 'column',
          gap: 20,
        }}>
          {messages.length === 0 ? (
            <div style={{ textAlign: 'center', padding: 60, color: theme.text3 }}>
              <div style={{
                width: 72, height: 72, borderRadius: 16,
                background: theme.bgCard, display: 'flex', alignItems: 'center',
                justifyContent: 'center', margin: '0 auto 20px',
                border: `1px solid ${theme.border}`,
              }}>
                <svg width="28" height="28" viewBox="0 0 24 24" fill="none">
                  <path d="M21 11.5a8.38 8.38 0 0 1-.9 3.8 8.5 8.5 0 0 1-7.6 4.7 8.38 8.38 0 0 1-3.8-.9L3 21l1.9-5.7a8.38 8.38 0 0 1-.9-3.8 8.5 8.5 0 0 1 4.7-7.6 8.38 8.38 0 0 1 3.8-.9h.5a8.48 8.48 0 0 1 8 8v.5z" stroke={theme.accent} strokeWidth="1.5"/>
                </svg>
              </div>
              <div style={{ fontSize: 15, fontWeight: 500, marginBottom: 6, color: theme.text }}>开始法规对话</div>
              <div className="mono" style={{ fontSize: 11 }}>选择快捷问题或输入您的问题</div>
            </div>
          ) : (
            messages.map(msg => (
              <div key={msg.id} style={{
                display: 'flex', gap: 12,
                flexDirection: msg.role === 'user' ? 'row-reverse' : 'row',
              }}>
                {msg.role === 'assistant' && (
                  <div style={{
                    width: 32, height: 32, borderRadius: 8,
                    background: theme.gradientAccent, display: 'flex',
                    alignItems: 'center', justifyContent: 'center', flexShrink: 0,
                  }}>
                    <svg width="16" height="16" viewBox="0 0 24 24" fill="none">
                      <circle cx="12" cy="12" r="6" fill="#fff"/>
                    </svg>
                  </div>
                )}
                <div style={{
                  maxWidth: '80%',
                  padding: msg.role === 'user' ? '12px 18px' : '14px 18px',
                  background: msg.role === 'user' ? theme.gradientAccent : theme.bgCard,
                  borderRadius: 12,
                  color: msg.role === 'user' ? '#fff' : theme.text,
                  fontSize: 14, lineHeight: 1.6, whiteSpace: 'pre-wrap',
                  border: msg.role === 'assistant' ? `1px solid ${theme.border}` : 'none',
                }}>
                  {msg.role === 'assistant' ? (
                    <CitedAnswer
                      text={msg.content}
                      sources={msg.sources ?? retrievals}
                      onCiteClick={(idx) => {
                        const msgSources = msg.sources ?? retrievals;
                        setRetrievals(msgSources);
                        setHighlightedSourceIdx(idx);
                        const el = document.getElementById(`source-${idx}`);
                        if (el) el.scrollIntoView({ behavior: 'smooth', block: 'center' });
                      }}
                    />
                  ) : msg.content}
                  {msg.role === 'assistant' && msg.retrievalIds && msg.retrievalIds.length > 0 && (
                    <div style={{
                      marginTop: 10, paddingTop: 10,
                      borderTop: `1px solid ${theme.border}`,
                      display: 'flex', alignItems: 'center', gap: 6,
                    }}>
                      <svg width="12" height="12" viewBox="0 0 24 24" fill="none">
                        <path d="M14 2H6C5 2 4 3 4 4V20C4 21 5 22 6 22H18C19 22 20 21 20 20V8L14 2Z" stroke={theme.accent} strokeWidth="1.5"/>
                      </svg>
                      <span className="mono" style={{ fontSize: 11, color: theme.accent }}>
                        {msg.retrievalIds.length} 个法规引用
                      </span>
                    </div>
                  )}
                </div>
              </div>
            ))
          )}
          {loading && (
            <div style={{ display: 'flex', gap: 12 }}>
              <div style={{
                width: 32, height: 32, borderRadius: 8,
                background: theme.gradientAccent, display: 'flex',
                alignItems: 'center', justifyContent: 'center',
              }}>
                <svg width="16" height="16" viewBox="0 0 24 24" fill="none">
                  <circle cx="12" cy="12" r="6" fill="#fff"/>
                </svg>
              </div>
              <div style={{
                padding: '14px 18px', background: theme.bgCard,
                borderRadius: 12, border: `1px solid ${theme.border}`,
                display: 'flex', alignItems: 'center', gap: 8,
              }}>
                <div style={{
                  width: 6, height: 6, borderRadius: '50%',
                  background: theme.accent, animation: 'pulse 1s infinite',
                }} />
                <span style={{ fontSize: 13, color: theme.text2 }}>检索中...</span>
              </div>
            </div>
          )}
          {/* Scroll anchor */}
          <div ref={messagesEndRef} />
        </div>

        {/* Input area */}
        <div style={{ padding: '16px 32px 20px', background: theme.bg, borderTop: `1px solid ${theme.border}` }}>
          {/* Filter row */}
          <div style={{ display: 'flex', gap: 8, marginBottom: 10, alignItems: 'center' }}>
            <span className="mono" style={{ fontSize: 11, color: theme.text3, whiteSpace: 'nowrap' }}>法规类型</span>
            <input
              value={filterRegulationType}
              onChange={e => setFilterRegulationType(e.target.value)}
              placeholder="如: GB / UN-ECE / IATF（留空不过滤）"
              style={{
                flex: 1, maxWidth: 280, padding: '5px 10px', fontSize: 12,
                background: theme.bgHover, border: `1px solid ${theme.border}`,
                borderRadius: 6, color: theme.text, outline: 'none',
              }}
            />
          </div>

          {/* Quick questions */}
          <div style={{ display: 'flex', gap: 8, marginBottom: 12, flexWrap: 'wrap' }}>
            {quickQuestions.map(q => (
              <button
                key={q}
                onClick={() => sendMessage(q)}
                disabled={loading}
                style={{
                  padding: '6px 14px', fontSize: 12, background: theme.bgCard,
                  border: `1px solid ${theme.border}`, borderRadius: 6,
                  color: theme.text2, cursor: loading ? 'not-allowed' : 'pointer',
                }}
              >{q}</button>
            ))}
          </div>

          {/* Send row */}
          <div style={{ display: 'flex', gap: 10 }}>
            <input
              value={input}
              onChange={e => setInput(e.target.value)}
              onKeyDown={e => e.key === 'Enter' && !e.shiftKey && sendMessage(input)}
              placeholder="输入法规问题..."
              style={{
                flex: 1, padding: 12, fontSize: 14,
                background: theme.bgCard, border: `1px solid ${theme.border}`,
                borderRadius: 8, color: theme.text, outline: 'none',
              }}
            />
            <button
              onClick={() => sendMessage(input)}
              disabled={loading || !input.trim()}
              style={{
                padding: '12px 24px', fontSize: 14, fontWeight: 600,
                background: loading || !input.trim() ? theme.bgHover : theme.gradientAccent,
                color: loading || !input.trim() ? theme.text3 : '#fff',
                border: 'none', borderRadius: 8,
                cursor: loading || !input.trim() ? 'not-allowed' : 'pointer',
              }}
            >发送</button>
            {loading && (
              <button
                onClick={() => { abortRef.current?.abort(); setLoading(false); }}
                style={{
                  padding: '12px 16px', fontSize: 13, background: theme.bgCard,
                  border: `1px solid ${theme.border}`, borderRadius: 8,
                  color: theme.text2, cursor: 'pointer',
                }}
              >停止</button>
            )}
            {!loading && messages.length > 0 && (
              <button
                onClick={() => setShowClearConfirm(true)}
                style={{
                  padding: '12px 16px', fontSize: 13, background: theme.bgCard,
                  border: `1px solid ${theme.border}`, borderRadius: 8,
                  color: theme.text2, cursor: 'pointer',
                }}
              >清空</button>
            )}
            {!loading && messages.filter(m => m.role === 'assistant').length > 0 && (
              <button
                onClick={regenerateLastAnswer}
                style={{
                  padding: '12px 16px', fontSize: 13, background: theme.bgCard,
                  border: `1px solid ${theme.border}`, borderRadius: 8,
                  color: theme.text2, cursor: 'pointer',
                }}
              >重生成</button>
            )}
          </div>
        </div>
      </div>

      {/* ── Right: retrieved sources panel ───────────────────── */}
      <div style={{ flex: '0 0 40%', display: 'flex', flexDirection: 'column', background: theme.bgCard }}>
        <div style={{
          padding: '20px 24px', borderBottom: `1px solid ${theme.border}`,
          display: 'flex', alignItems: 'center', gap: 10,
        }}>
          <div style={{
            width: 28, height: 28, borderRadius: 6,
            background: theme.gradientAccent, display: 'flex',
            alignItems: 'center', justifyContent: 'center',
          }}>
            <svg width="14" height="14" viewBox="0 0 24 24" fill="none">
              <path d="M14 2H6C5 2 4 3 4 4V20C4 21 5 22 6 22H18C19 22 20 21 20 20V8L14 2Z" stroke="#fff" strokeWidth="1.5"/>
            </svg>
          </div>
          <span className="mono" style={{ fontSize: 12, fontWeight: 600, color: theme.accent, letterSpacing: '1px' }}>
            RETRIEVED FRAGMENTS
          </span>
          {retrievals.length > 0 && (
            <span className="mono" style={{
              fontSize: 11, padding: '4px 10px',
              background: theme.bgHover, borderRadius: 4, color: theme.text3,
            }}>{retrievals.length}</span>
          )}
        </div>

        <div style={{ flex: 1, overflowY: 'auto', padding: '16px 24px' }}>
          {retrievals.length > 0 ? (
            <div style={{ display: 'flex', flexDirection: 'column', gap: 12 }}>
              {retrievals.map((r, i) => (
                <div
                  key={r.id}
                  id={`source-${i + 1}`}
                  onClick={() => setSelectedRetrieval(r)}
                  style={{
                    padding: 16, background: highlightedSourceIdx === i + 1 ? theme.bgElevated : theme.bgHover,
                    borderRadius: 10, border: `1px solid ${highlightedSourceIdx === i + 1 ? theme.accent : theme.border}`,
                    cursor: 'pointer', position: 'relative',
                    transition: 'border-color 0.2s, background 0.2s',
                  }}
                >
                  <div style={{
                    position: 'absolute', left: 0, top: 16, bottom: 16,
                    width: 3, background: theme.gradientAccent, borderRadius: 2,
                  }} />
                  <div style={{ paddingLeft: 8 }}>
                    <div style={{
                      display: 'flex', alignItems: 'center',
                      justifyContent: 'space-between', marginBottom: 8,
                    }}>
                      <span className="mono" style={{ fontSize: 11, fontWeight: 700, color: theme.accent }}>#{i + 1}</span>
                      <div style={{ display: 'flex', alignItems: 'center', gap: 10 }}>
                        {r.downloadUrl && (
                          <a
                            href={r.downloadUrl}
                            target="_blank"
                            rel="noreferrer"
                            onClick={e => e.stopPropagation()}
                            style={{ fontSize: 11, color: theme.accent, textDecoration: 'none' }}
                          >下载文档</a>
                        )}
                        <span className="mono" style={{ fontSize: 11, fontWeight: 600, color: theme.accent }}>
                          {(r.score * 100).toFixed(0)}%
                        </span>
                      </div>
                    </div>
                    <div style={{ fontSize: 13, fontWeight: 500, marginBottom: 4, color: theme.text }}>{r.file}</div>
                    <div className="mono" style={{ fontSize: 11, color: theme.text3, marginBottom: 8 }}>
                      {r.clause}{r.docId ? ` · ${r.docId}` : ''}
                    </div>
                    <div style={{
                      fontSize: 12, color: theme.text2, lineHeight: 1.5,
                      display: '-webkit-box', WebkitLineClamp: 3,
                      WebkitBoxOrient: 'vertical', overflow: 'hidden',
                    }}>{r.content}</div>
                  </div>
                </div>
              ))}
            </div>
          ) : (
            <div style={{ textAlign: 'center', padding: 40, color: theme.text3 }}>
              <div style={{
                width: 48, height: 48, borderRadius: 10,
                background: theme.bgHover, display: 'flex', alignItems: 'center',
                justifyContent: 'center', margin: '0 auto 16px',
              }}>
                <svg width="20" height="20" viewBox="0 0 24 24" fill="none">
                  <path d="M14 2H6C5 2 4 3 4 4V20C4 21 5 22 6 22H18C19 22 20 21 20 20V8L14 2Z" stroke={theme.text3} strokeWidth="1.5"/>
                </svg>
              </div>
              <div className="mono" style={{ fontSize: 11 }}>对话后显示相关法规</div>
            </div>
          )}
        </div>
      </div>

      {/* ── Clear confirm modal ───────────────────────────────── */}
      {showClearConfirm && (
        <div style={{
          position: 'fixed', top: 0, left: 0, right: 0, bottom: 0,
          background: 'rgba(0,0,0,0.6)', display: 'flex',
          alignItems: 'center', justifyContent: 'center', zIndex: 1000,
        }}>
          <div style={{
            padding: 24, background: theme.bgCard, borderRadius: 16,
            maxWidth: 400, border: `1px solid ${theme.border}`,
          }}>
            <div style={{ fontSize: 15, fontWeight: 600, marginBottom: 12, color: theme.text }}>确定清空对话？</div>
            <div style={{ fontSize: 13, color: theme.text2, marginBottom: 20 }}>此操作不可恢复，会话历史将被重置</div>
            <div style={{ display: 'flex', gap: 10, justifyContent: 'flex-end' }}>
              <button
                onClick={() => setShowClearConfirm(false)}
                style={{
                  padding: '10px 18px', fontSize: 13, background: theme.bgHover,
                  border: 'none', borderRadius: 8, color: theme.text2, cursor: 'pointer',
                }}
              >取消</button>
              <button
                onClick={clearMessages}
                style={{
                  padding: '10px 18px', fontSize: 13, fontWeight: 600,
                  background: theme.accent, border: 'none', borderRadius: 8,
                  color: '#fff', cursor: 'pointer',
                }}
              >确认</button>
            </div>
          </div>
        </div>
      )}

      {/* ── Source detail modal ───────────────────────────────── */}
      {selectedRetrieval && (
        <div
          onClick={() => setSelectedRetrieval(null)}
          style={{
            position: 'fixed', top: 0, left: 0, right: 0, bottom: 0,
            background: 'rgba(0,0,0,0.6)', display: 'flex',
            alignItems: 'center', justifyContent: 'center', zIndex: 1000,
          }}
        >
          <div
            onClick={e => e.stopPropagation()}
            style={{
              width: 520, maxWidth: '90%', maxHeight: '80%',
              overflowY: 'auto', padding: 24, background: theme.bgCard,
              borderRadius: 16, border: `1px solid ${theme.accent}`,
              boxShadow: '0 8px 32px rgba(0,0,0,0.3)',
            }}
          >
            <div style={{
              display: 'flex', alignItems: 'center',
              justifyContent: 'space-between', marginBottom: 16,
            }}>
              <div style={{ display: 'flex', alignItems: 'center', gap: 10, flexWrap: 'wrap' }}>
                <div style={{ padding: '4px 10px', background: theme.gradientAccent, borderRadius: 6 }}>
                  <span className="mono" style={{ fontSize: 11, fontWeight: 600, color: '#fff' }}>
                    {(selectedRetrieval.score * 100).toFixed(0)}%
                  </span>
                </div>
                <span style={{ fontSize: 14, fontWeight: 600, color: theme.text }}>{selectedRetrieval.file}</span>
                {selectedRetrieval.downloadUrl && (
                  <a
                    href={selectedRetrieval.downloadUrl}
                    target="_blank"
                    rel="noreferrer"
                    style={{ fontSize: 12, color: theme.accent, textDecoration: 'none' }}
                  >下载关联文档</a>
                )}
              </div>
              <button
                onClick={() => setSelectedRetrieval(null)}
                style={{
                  width: 28, height: 28, background: theme.bgHover,
                  border: 'none', borderRadius: 6, cursor: 'pointer',
                  display: 'flex', alignItems: 'center', justifyContent: 'center',
                }}
              >
                <svg width="14" height="14" viewBox="0 0 24 24" fill="none">
                  <path d="M18 6L6 18M6 6L18 18" stroke={theme.text3} strokeWidth="2" strokeLinecap="round"/>
                </svg>
              </button>
            </div>
            <div style={{ padding: '10px 14px', background: theme.bgHover, borderRadius: 8, marginBottom: 16 }}>
              <span className="mono" style={{ fontSize: 12, color: theme.accent }}>{selectedRetrieval.clause}</span>
              {selectedRetrieval.docId && (
                <span className="mono" style={{ fontSize: 11, color: theme.text3, marginLeft: 8 }}>{selectedRetrieval.docId}</span>
              )}
            </div>
            <div style={{ fontSize: 14, lineHeight: 1.7, color: theme.text2, whiteSpace: 'pre-wrap' }}>
              {selectedRetrieval.content}
            </div>
          </div>
        </div>
      )}
    </div>
  );
};