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

import React, { 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[]>([]);
  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);

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

  async function loadQuickQuestions() {
    try {
      const response = await getQuickQuestions();
      setQuickQuestions(response.questions.map(q => q.question));
    } catch (error) {
      console.error('Failed to load quick questions:', error);
    }
  }

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

  const sendMessage = (text: string) => {
    if (!text.trim()) return;

    const userMsg = { id: nextMessageId(), role: 'user' as const, content: text };
    setMessages((prev) => [...prev, userMsg]);
    setInput('');
    setLoading(true);
    setRetrievals([]);
    setHighlightedSourceIdx(null);

    let currentResponse = '';
    const activeFilters = filterRegulationType.trim() || undefined;

    void ragChat(
      text,
      5,
      (data: unknown) => {
        const sseData = data as {
          type: string;
          text?: string;
          docs?: Array<{
            id: string;
            score: number;
            preview: string;
            doc_name: string;
            clause: string;
            doc_id?: string;
            download_url?: string;
          }>;
        };

        if (sseData.type === 'retrieved' && sseData.docs) {
          const retrievedDocs: RetrievalData[] = sseData.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(retrievedDocs);
        } else if (sseData.type === 'chunk' && sseData.text) {
          currentResponse += sseData.text;
          setMessages((prev) => {
            const lastMsg = prev[prev.length - 1];
            if (lastMsg?.role === 'assistant') {
              return [...prev.slice(0, -1), { ...lastMsg, content: currentResponse }];
            }
            return [...prev, { id: nextMessageId(), role: 'assistant' as const, content: currentResponse }];
          });
        } else if (sseData.type === 'done') {
          setLoading(false);
        } else if (sseData.type === 'error') {
          setLoading(false);
        }
      },
      (error: Error) => {
        console.error('RAG chat error:', error);
        setLoading(false);
        setMessages((prev) => [
          ...prev,
          { id: nextMessageId(), role: 'assistant' as const, content: '抱歉，连接服务器时出错，请稍后再试。' }
        ]);
      },
      () => {
        setLoading(false);
      },
      activeFilters
    );
  };

  const clearMessages = () => {
    setMessages([]);
    setRetrievals([]);
    setShowClearConfirm(false);
  };

  const regenerateLastAnswer = () => {
    if (messages.length < 2) return;
    const lastUserMsg = messages.filter((m) => m.role === 'user').pop();
    if (!lastUserMsg) return;

    setLoading(true);
    setMessages((prev) => [...prev.slice(0, -1)]);
    setRetrievals([]);
    setHighlightedSourceIdx(null);

    let currentResponse = '';
    const activeFilters = filterRegulationType.trim() || undefined;

    void ragChat(
      lastUserMsg.content,
      5,
      (data: unknown) => {
        const sseData = data as {
          type: string;
          text?: string;
          docs?: Array<{
            id: string;
            score: number;
            preview: string;
            doc_name: string;
            clause: string;
            doc_id?: string;
            download_url?: string;
          }>;
        };

        if (sseData.type === 'retrieved' && sseData.docs) {
          const retrievedDocs: RetrievalData[] = sseData.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(retrievedDocs);
        } else if (sseData.type === 'chunk' && sseData.text) {
          currentResponse += sseData.text;
          setMessages((prev) => {
            const lastMsg = prev[prev.length - 1];
            if (lastMsg?.role === 'assistant') {
              return [...prev.slice(0, -1), { ...lastMsg, content: currentResponse }];
            }
            return [...prev, { id: nextMessageId(), role: 'assistant' as const, content: currentResponse }];
          });
        } else if (sseData.type === 'done') {
          setLoading(false);
        }
      },
      (error: Error) => {
        console.error('RAG chat error:', error);
        setLoading(false);
        setMessages((prev) => [
          ...prev,
          { id: nextMessageId(), role: 'assistant' as const, content: '抱歉，连接服务器时出错，请稍后再试。' }
        ]);
      },
      () => {
        setLoading(false);
      },
      activeFilters
    );
  };

  return (
    <div style={{
      flex: 1,
      display: 'flex',
      height: 'calc(100vh - 128px)',
    }}>
      <div style={{
        flex: '0 0 60%',
        display: 'flex',
        flexDirection: 'column',
        borderRight: `1px solid ${theme.border}`,
      }}>
        <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={retrievals}
                      onCiteClick={(idx) => {
                        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>
          )}
        </div>

        <div style={{
          padding: '16px 32px 20px',
          background: theme.bg,
          borderTop: `1px solid ${theme.border}`,
        }}>
          <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>

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

          <div style={{ display: 'flex', gap: 10 }}>
            <input
              value={input}
              onChange={(e) => setInput(e.target.value)}
              onKeyDown={(e) => e.key === 'Enter' && 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>
            {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>
            )}
            {messages.filter(m => m.role === 'assistant').length > 0 && (
              <button
                onClick={regenerateLastAnswer}
                disabled={loading}
                style={{
                  padding: '12px 16px',
                  fontSize: 13,
                  background: theme.bgCard,
                  border: `1px solid ${theme.border}`,
                  borderRadius: 8,
                  color: theme.text2,
                  cursor: loading ? 'not-allowed' : 'pointer',
                }}
              >重生成</button>
            )}
          </div>
        </div>
      </div>

      <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={(event) => event.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,
                      overflow: 'hidden',
                      textOverflow: 'ellipsis',
                      whiteSpace: 'nowrap',
                    }}>{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>

      {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>
      )}

      {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%',
              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>
  );
};