course-agent-od/web/src/components/chat/source-references.tsx

"use client";

import { useState, useMemo } from "react";
import { Database, Globe, ChevronDown, ChevronRight, Star } from "lucide-react";
import { Button } from "@/components/ui/button";
import { cn } from "@/lib/utils";

interface SourceReference {
  id?: number;
  title: string;
  filename?: string;
  page?: number;
  score?: number;
  preview: string;
  url?: string;
  source_type?: "web" | "rag";
}

interface SourceReferencesProps {
  sources: SourceReference[];
  maxSources?: number;
  answerContent?: string;
  messageId?: string;
}

/** 从 answer 中提取长度 >= minLen 的不重叠片段，用于在来源原文中高亮匹配 */
function findMatchedSpans(
  answer: string,
  sourceText: string,
  minLen = 15
): Array<{ start: number; end: number }> {
  if (!answer || !sourceText) return [];

  // 从 answer 中切分出所有有意义的中文/英文片段
  const segments: string[] = [];
  // 按句号、换行等断开
  const sentences = answer.split(/[。，；！？\n、：（）]/);
  for (const s of sentences) {
    const trimmed = s.trim();
    if (trimmed.length >= minLen) {
      segments.push(trimmed);
    }
  }

  // 在 sourceText 中查找每个片段
  const spans: Array<{ start: number; end: number }> = [];
  for (const seg of segments) {
    let pos = 0;
    while (pos < sourceText.length) {
      const idx = sourceText.indexOf(seg, pos);
      if (idx === -1) break;
      const end = idx + seg.length;
      // 检查是否与已有 span 重叠，有则合并
      const overlapping = spans.find(
        (s) => !(end <= s.start || idx >= s.end)
      );
      if (overlapping) {
        overlapping.start = Math.min(overlapping.start, idx);
        overlapping.end = Math.max(overlapping.end, end);
      } else {
        spans.push({ start: idx, end });
      }
      pos = end;
    }
  }

  return spans.sort((a, b) => a.start - b.start);
}

/** 将匹配的 span 用 <mark> 包裹 */
function highlightText(
  text: string,
  spans: Array<{ start: number; end: number }>
): React.ReactNode {
  if (!spans.length) return text;

  // 合并重叠/相邻的 span
  const merged: Array<{ start: number; end: number }> = [];
  for (const span of spans) {
    const last = merged[merged.length - 1];
    if (last && span.start <= last.end + 3) {
      last.end = Math.max(last.end, span.end);
    } else {
      merged.push({ ...span });
    }
  }

  const parts: React.ReactNode[] = [];
  let last = 0;
  for (const span of merged) {
    if (span.start > last) {
      parts.push(text.slice(last, span.start));
    }
    parts.push(
      <mark key={span.start} className="bg-yellow-200/70 dark:bg-yellow-500/30 rounded-sm px-0.5">
        {text.slice(span.start, span.end)}
      </mark>
    );
    last = span.end;
  }
  if (last < text.length) {
    parts.push(text.slice(last));
  }
  return <>{parts}</>;
}

function SourceRow({
  source,
  answerContent,
  messageId,
}: {
  source: SourceReference;
  answerContent?: string;
  messageId?: string;
}) {
  const [showPreview, setShowPreview] = useState(false);

  const matchedSpans = useMemo(() => {
    if (!answerContent || !source.preview) return [];
    return findMatchedSpans(answerContent, source.preview);
  }, [answerContent, source.preview]);

  const previewContent = useMemo(() => {
    if (!matchedSpans.length) return source.preview;
    return highlightText(source.preview, matchedSpans);
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [showPreview, matchedSpans]);

  return (
    <div
      id={source.id != null && messageId ? `source-${messageId}-${source.id}` : undefined}
      className="scroll-mt-20"
    >
      <button
        className={cn(
          "w-full flex items-center gap-1.5 px-2 py-1 -mx-2 rounded text-left",
          "hover:bg-muted/60 transition-colors group text-xs"
        )}
        onClick={() => setShowPreview(!showPreview)}
      >
        <ChevronRight
          className={cn(
            "h-3 w-3 flex-shrink-0 text-muted-foreground/60 transition-transform",
            showPreview && "rotate-90"
          )}
        />
        {source.id != null && (
          <span className="inline-flex items-center justify-center min-w-[18px] h-[18px] rounded-full bg-primary/10 text-primary text-[10px] font-bold flex-shrink-0">
            {source.id}
          </span>
        )}
        <span className="font-medium truncate">{source.title}</span>
        {source.filename && (
          <span className="text-muted-foreground/70 truncate hidden sm:inline">
            {source.filename.replace(/\.(pdf|docx?|txt|md)$/i, "")}
          </span>
        )}
        {source.score != null && source.score > 0 && source.score < 1 && (
          <span className="flex items-center gap-0.5 text-[10px] text-muted-foreground flex-shrink-0 ml-auto">
            <Star className="h-2.5 w-2.5 text-yellow-500" />
            {(source.score * 100).toFixed(0)}%
          </span>
        )}
      </button>

      {showPreview && (
        <div className="ml-7 pl-3 pr-2 py-1.5 mb-0.5 border-l-2 border-primary/20 bg-muted/30 rounded-r text-xs text-muted-foreground leading-relaxed max-h-48 overflow-y-auto">
          {previewContent}
        </div>
      )}
    </div>
  );
}

export default function SourceReferences({
  sources,
  maxSources = 20,
  answerContent,
  messageId,
}: SourceReferencesProps) {
  const [expanded, setExpanded] = useState(false);

  if (!sources || sources.length === 0) {
    return null;
  }

  const displaySources = sources.slice(0, maxSources);
  const showExpandButton = sources.length > maxSources;
  const visibleSources = expanded ? sources : displaySources;

  const ragSources = visibleSources.filter((s) => s.source_type !== "web");
  const webSources = visibleSources.filter((s) => s.source_type === "web");

  const sourceSection = (
    <div className="space-y-1">
      {ragSources.length > 0 && (
        <div className="space-y-1">
          <div className="flex items-center gap-1.5 text-[11px] font-medium text-muted-foreground px-0.5">
            <Database className="h-3 w-3 text-blue-500" />
            <span>
              知识库来源 ({ragSources.length}
              {expanded && showExpandButton
                ? `/${sources.filter((s) => s.source_type !== "web").length}`
                : ""}
              )
            </span>
          </div>
          <div className="divide-y divide-border/30">
            {ragSources.map((source, i) => (
              <SourceRow
                key={i}
                source={source}
                answerContent={answerContent}
                messageId={messageId}
              />
            ))}
          </div>
        </div>
      )}

      {webSources.length > 0 && (
        <div className="space-y-1">
          <div className="flex items-center gap-1.5 text-[11px] font-medium text-muted-foreground px-0.5">
            <Globe className="h-3 w-3 text-green-500" />
            <span>
              网络来源 ({webSources.length}
              {expanded && showExpandButton
                ? `/${sources.filter((s) => s.source_type === "web").length}`
                : ""}
              )
            </span>
          </div>
          <div className="divide-y divide-border/30">
            {webSources.map((source, i) => (
              <SourceRow
                key={i}
                source={source}
                answerContent={answerContent}
                messageId={messageId}
              />
            ))}
          </div>
        </div>
      )}
    </div>
  );

  return (
    <div className="mt-3 space-y-2">
      {expanded ? (
        <div className="max-h-96 overflow-y-auto pr-1">
          {sourceSection}
        </div>
      ) : (
        sourceSection
      )}

      {showExpandButton && (
        <Button
          variant="ghost"
          size="sm"
          className="w-full h-7 text-[11px] text-muted-foreground"
          onClick={() => setExpanded(!expanded)}
        >
          {expanded
            ? `收起（共 ${sources.length} 个）`
            : `展开全部 ${sources.length} 个来源`}
          <ChevronDown
            className={cn(
              "ml-1 h-3 w-3 transition-transform",
              expanded && "rotate-180"
            )}
          />
        </Button>
      )}
    </div>
  );
}