"""
知识库对话相关 API
"""
import asyncio
import json
import logging
import re
import time
from typing import Optional, List, Dict, Any

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, delete

from app.core.database import get_db
from app.core.deps import get_current_user
from app.models.llm_model_config import LLMModelConfig
from app.models.project import Project
from app.models.user import User
from app.models.chat_session import ChatSession, ChatMessage
from app.schemas.response import success_response
from app.services.project_service import (
    get_project_or_404,
    require_project_read_access,
    require_project_write_access,
)
from app.services.rag_service import rag_service
from app.services.zvec_service import zvec_service
from app.services.project_vectorization_task_service import project_vectorization_task_service
from pydantic import BaseModel


router = APIRouter()

logger = logging.getLogger(__name__)

CHAT_HISTORY_MESSAGE_LIMIT = 20

# 旧版本会在被中断的助手消息内容末尾追加 "interrupt" 标记。
# 自 status 字段上线后新消息不再写该标记，判断中断只以 status 为准；
# 下面两个函数仅用于读取历史消息时的只读兼容（绝不修改数据库）。
_LEGACY_INTERRUPTED_MARKER = "interrupt"


def _legacy_interrupted_content(content: str) -> bool:
    """只读兼容旧数据：判断内容是否带旧版中断标记。"""
    if not content:
        return False
    return content == _LEGACY_INTERRUPTED_MARKER or content.endswith(
        f"\n\n{_LEGACY_INTERRUPTED_MARKER}"
    )


def _strip_legacy_interrupt_marker(content: str) -> str:
    """只读兼容旧数据：返回去掉旧版中断标记的展示内容（不改数据库）。"""
    if content == _LEGACY_INTERRUPTED_MARKER:
        return ""
    suffix = f"\n\n{_LEGACY_INTERRUPTED_MARKER}"
    if content.endswith(suffix):
        return content[: -len(suffix)].rstrip()
    return content


def _is_interrupted_message(message) -> bool:
    """判断助手消息是否已中断。

    新消息以 status 字段为准；旧数据（status 缺失或仍为 pending）仅做
    只读兼容：按旧版内容标记判断，不修改数据库。
    """
    if message.role != "assistant":
        return False
    status = (message.status or "").strip() or "pending"
    if status == "interrupted":
        return True
    return status == "pending" and _legacy_interrupted_content(message.content or "")

class VectorizeRequest(BaseModel):
    """批量向量化请求"""
    force: bool = False  # False=增量（跳过未变更文件）, True=全量重建


@router.get("/projects/{project_id}/vectorize/progress", response_model=dict)
async def get_vectorize_progress(
    project_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """查询项目向量化进度"""
    project = await get_project_or_404(db, project_id)
    await require_project_read_access(db, project_id, current_user)

    progress = await zvec_service.get_progress(db, project_id, project.storage_key)
    return success_response(data=progress)


@router.post("/projects/{project_id}/vectorize", response_model=dict)
async def vectorize_project(
    project_id: int,
    req: VectorizeRequest,
    background_tasks: BackgroundTasks,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """创建项目文件夹向量化后台任务（增量或全量）"""
    await get_project_or_404(db, project_id)
    await require_project_write_access(db, project_id, current_user)

    if await zvec_service.get_embedding_config(db) is None:
        raise HTTPException(
            status_code=400,
            detail="尚未配置可用的向量模型，请先在「模型配置 - 向量模型」中添加并启用。",
        )

    running_task = await project_vectorization_task_service.get_running_task(db, project_id)
    if running_task:
        task = running_task
    else:
        task = await project_vectorization_task_service.create_task(
            db,
            project_id,
            current_user.id,
            force=req.force,
        )
        background_tasks.add_task(project_vectorization_task_service.run_task, task.task_id)

    return success_response(
        data=project_vectorization_task_service.serialize_task(task),
        message="向量化任务已提交",
    )


@router.get("/projects/{project_id}/vectorize/tasks/latest", response_model=dict)
async def get_latest_vectorize_task(
    project_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """查询项目最近一次向量化任务"""
    await require_project_read_access(db, project_id, current_user)
    task = await project_vectorization_task_service.get_latest_task(db, project_id)
    return success_response(
        data=project_vectorization_task_service.serialize_task(task) if task else None
    )


@router.get("/projects/{project_id}/vectorize/tasks/{task_id}", response_model=dict)
async def get_vectorize_task(
    project_id: int,
    task_id: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """查询项目向量化任务状态"""
    await require_project_read_access(db, project_id, current_user)
    task = await project_vectorization_task_service.get_task(db, project_id, task_id)
    if not task:
        raise HTTPException(status_code=404, detail="向量化任务不存在")
    return success_response(data=project_vectorization_task_service.serialize_task(task))


class ChatCreateRequest(BaseModel):
    """创建对话会话请求"""
    project_id: int
    llm_config_id: int
    title: str = "新对话"


class ChatMessageRequest(BaseModel):
    """发送聊天消息请求"""
    session_id: int
    message: str
    # 是否同时写入一条新的用户消息。重新生成旧回复时应传 false，避免用户提问重复入库。
    insert_user_message: bool = True


class ChatSessionUpdateRequest(BaseModel):
    """更新对话会话请求"""
    title: str


class ChatResponse(BaseModel):
    """对话响应"""
    session_id: int
    user_message: str
    assistant_message: str


def _parse_refs(value):
    if not value:
        return []
    try:
        return json.loads(value)
    except (ValueError, TypeError):
        return []


def _normalize_refs(refs):
    """将存储的引用归一化为 [{citation_id, file_path, anchor_text, excerpt}] 形式。

    兼容旧格式（纯文件路径数组）与新格式（对象数组）。
    """
    normalized = []
    for idx, ref in enumerate(refs or [], 1):
        if isinstance(ref, dict):
            file_path = ref.get("file_path")
            citation_id = ref.get("citation_id", idx)
            anchor_text = ref.get("anchor_text") or ""
            excerpt = ref.get("excerpt") or ""
            content = ref.get("content") or ""
            chunk_index = ref.get("chunk_index")
            hit_terms = ref.get("hit_terms") or []
            quotes = ref.get("quotes") or []
            quote_occurrences = ref.get("quote_occurrences") or []
        else:
            file_path = ref
            citation_id = idx
            anchor_text = ""
            excerpt = ""
            content = ""
            chunk_index = None
            hit_terms = []
            quotes = []
            quote_occurrences = []
        if file_path:
            normalized.append({
                "citation_id": citation_id,
                "file_path": file_path,
                "anchor_text": anchor_text,
                "excerpt": excerpt,
                "content": content,
                "chunk_index": chunk_index,
                "hit_terms": hit_terms,
                "quotes": quotes,
                "quote_occurrences": quote_occurrences,
            })
    return normalized


def _build_reference_items(refs, project_id=None):
    return [
        {
            "citation_id": ref["citation_id"],
            "file_path": ref["file_path"],
            "file_name": ref["file_path"].rsplit("/", 1)[-1],
            "anchor_text": ref.get("anchor_text") or "",
            "excerpt": ref.get("excerpt") or "",
            "content": ref.get("content") or "",
            "chunk_index": ref.get("chunk_index"),
            "hit_terms": ref.get("hit_terms") or [],
            "quotes": ref.get("quotes") or [],
            "quote_occurrences": ref.get("quote_occurrences") or [],
            "project_id": project_id,
        }
        for ref in _normalize_refs(refs)
    ]


def _canonicalize_message_citations(content: str, refs):
    """按文件合并引用并将正文引用重新编号，兼容历史重复分块数据。"""
    normalized = _normalize_refs(refs)
    if not normalized:
        return content, []

    canonical_refs = []
    file_to_id = {}
    old_to_new = {}
    for ref in normalized:
        file_path = ref["file_path"]
        old_id = int(ref["citation_id"])
        new_id = file_to_id.get(file_path)
        if new_id is None:
            new_id = len(canonical_refs) + 1
            file_to_id[file_path] = new_id
            canonical_refs.append({**ref, "citation_id": new_id})
        else:
            canonical_ref = canonical_refs[new_id - 1]
            current_excerpt = canonical_ref.get("excerpt", "").strip()
            incoming_excerpt = ref.get("excerpt", "").strip()
            if incoming_excerpt and incoming_excerpt not in current_excerpt:
                canonical_ref["excerpt"] = (
                    f"{current_excerpt}\n\n{incoming_excerpt}"
                    if current_excerpt else incoming_excerpt
                )
        old_to_new[old_id] = new_id

    def replace_citation(match):
        old_id = int(match.group(1))
        return f"[{old_to_new.get(old_id, old_id)}]"

    normalized_content = _CITATION_PATTERN.sub(replace_citation, content or "")
    normalized_content = re.sub(r"\[(\d+)\](?:\s*\[\1\])+", r"[\1]", normalized_content)
    return normalized_content, canonical_refs


def _stream_event(event: str, data) -> str:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


_CITATION_PATTERN = re.compile(r"\[(\d+)\]")


def _compact_cited_refs(answer: str, retrieved_docs):
    """仅保留实际使用的文档，并按正文首次引用顺序连续编号。"""
    docs_by_id = {
        int(doc.get("citation_id", index)): doc
        for index, doc in enumerate(retrieved_docs, 1)
    }
    old_to_new = {}
    file_to_new = {}
    refs = []

    for match in _CITATION_PATTERN.finditer(answer or ""):
        old_id = int(match.group(1))
        doc = docs_by_id.get(old_id)
        if doc is None or old_id in old_to_new:
            continue
        file_path = doc.get("file_path")
        if not file_path:
            continue

        new_id = file_to_new.get(file_path)
        if new_id is None:
            new_id = len(refs) + 1
            file_to_new[file_path] = new_id
            refs.append({
                "citation_id": new_id,
                "file_path": file_path,
                "anchor_text": doc.get("anchor_text") or "",
                # excerpt 保留精确命中的分块文本；content 保留命中点前后的
                # 上下文窗口（前端据此标出"命中区间"并支持跳转原文）
                "excerpt": doc.get("excerpt") or doc.get("anchor_text") or "",
                "content": doc.get("content") or "",
                "chunk_index": doc.get("chunk_index"),
                "hit_terms": doc.get("hit_terms") or [],
            })
        old_to_new[old_id] = new_id

    def replace_citation(match):
        old_id = int(match.group(1))
        return f"[{old_to_new.get(old_id, old_id)}]"

    normalized_answer = _CITATION_PATTERN.sub(replace_citation, answer or "")
    normalized_answer = re.sub(
        r"\[(\d+)\](?:\s*\[\1\])+", r"[\1]", normalized_answer
    )
    return normalized_answer, refs


MAX_SESSION_TITLE_LENGTH = 60


def _normalize_session_title(title: Optional[str]) -> str:
    """会话标题统一处理：取首行、去掉多余引号与空白，超长截断，空值回退。"""
    title = (title or "").strip().strip('"').strip("「」").strip()
    title = title.splitlines()[0].strip() if title else ""
    if not title:
        return "新对话"
    return title[:MAX_SESSION_TITLE_LENGTH]


@router.post("/sessions", response_model=dict)
async def create_chat_session(
    req: ChatCreateRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """创建新的对话会话"""
    await get_project_or_404(db, req.project_id)
    await require_project_read_access(db, req.project_id, current_user)
    model_result = await db.execute(
        select(LLMModelConfig).where(
            LLMModelConfig.config_id == req.llm_config_id,
            LLMModelConfig.model_type == "chat",
            LLMModelConfig.is_active == True,
        )
    )
    if model_result.scalar_one_or_none() is None:
        raise HTTPException(status_code=400, detail="请选择有效的对话模型")

    session = ChatSession(
        user_id=current_user.id,
        project_id=req.project_id,
        llm_config_id=req.llm_config_id,
        title=_normalize_session_title(req.title),
    )
    db.add(session)
    await db.commit()
    await db.refresh(session)

    return success_response(data={
        "session_id": session.id,
        "project_id": session.project_id,
        "llm_config_id": session.llm_config_id,
        "title": session.title,
        "created_at": session.created_at.isoformat(),
        "updated_at": session.updated_at.isoformat(),
    })


@router.get("/sessions", response_model=dict)
async def list_chat_sessions(
    project_id: Optional[int] = None,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """列出当前用户的对话会话，可按知识库筛选"""
    stmt = select(ChatSession).where(ChatSession.user_id == current_user.id)
    if project_id is not None:
        await require_project_read_access(db, project_id, current_user)
        stmt = stmt.where(ChatSession.project_id == project_id)
    stmt = stmt.order_by(ChatSession.updated_at.desc(), ChatSession.created_at.desc())
    result = await db.execute(stmt)
    sessions = result.scalars().all()

    return success_response(data=[{
        "session_id": s.id,
        "project_id": s.project_id,
        "llm_config_id": s.llm_config_id,
        "title": s.title,
        "created_at": s.created_at.isoformat(),
        "updated_at": s.updated_at.isoformat(),
    } for s in sessions])


@router.get("/sessions/{session_id}/messages", response_model=dict)
async def get_session_messages(
    session_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """获取对话会话的所有消息"""
    stmt = select(ChatSession).where(ChatSession.id == session_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=404, detail="对话会话不存在")

    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权访问该对话")

    await require_project_read_access(db, session.project_id, current_user)

    msg_stmt = select(ChatMessage).where(
        ChatMessage.session_id == session_id
    ).order_by(ChatMessage.created_at.asc())
    msg_result = await db.execute(msg_stmt)
    messages = msg_result.scalars().all()

    data = []
    for message in messages:
        stored_refs = _parse_refs(message.referenced_files)
        content = message.content
        refs = stored_refs
        status = message.status or "pending"
        if message.role == "assistant":
            content, refs = _canonicalize_message_citations(content, stored_refs)
            # 旧数据只读兼容：status 仍为 pending 且内容带旧版中断标记时，
            # 按“已中断”返回并在展示内容中移除标记文本（不修改数据库）。
            if status in ("", "pending") and _legacy_interrupted_content(content):
                status = "interrupted"
                content = _strip_legacy_interrupt_marker(content)
        thinking_log = []
        if message.thinking_log:
            try:
                thinking_log = json.loads(message.thinking_log) or []
            except (TypeError, ValueError):
                thinking_log = []
        data.append({
            "id": message.id,
            "role": message.role,
            "content": content,
            "status": status,
            "duration_ms": message.duration_ms,
            "thinking_log": thinking_log,
            "referenced_files": refs,
            "references": _build_reference_items(refs, session.project_id) if message.role == "assistant" else [],
            "created_at": message.created_at.isoformat(),
        })

    return success_response(data=data)


@router.get("/search", response_model=dict)
async def search_chat_messages(
    keyword: str,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """搜索聊天消息内容"""
    keyword = (keyword or "").strip()
    if not keyword:
        return success_response(data=[])

    escaped_keyword = (
        keyword.replace("\\", "\\\\")
        .replace("%", "\\%")
        .replace("_", "\\_")
    )
    stmt = (
        select(
            ChatMessage.id,
            ChatMessage.session_id,
            ChatMessage.role,
            ChatMessage.content,
            ChatMessage.created_at,
            ChatSession.title,
            ChatSession.project_id,
            ChatSession.llm_config_id,
            Project.name.label("project_name"),
            LLMModelConfig.model_name.label("model_name"),
        )
        .join(ChatSession, ChatMessage.session_id == ChatSession.id)
        .join(Project, Project.id == ChatSession.project_id)
        .join(LLMModelConfig, LLMModelConfig.config_id == ChatSession.llm_config_id)
        .where(
            ChatSession.user_id == current_user.id,
            ChatMessage.content.ilike(f"%{escaped_keyword}%", escape="\\"),
        )
        .order_by(ChatMessage.created_at.desc())
        .limit(50)
    )
    result = await db.execute(stmt)
    rows = result.all()

    def _build_snippet(content: str, term: str) -> str:
        text = content or ""
        if not text:
            return ""
        lowered = text.lower()
        needle = term.lower()
        idx = lowered.find(needle)
        if idx < 0:
            return text[:120]
        start = max(0, idx - 24)
        end = min(len(text), idx + len(term) + 48)
        prefix = "..." if start > 0 else ""
        suffix = "..." if end < len(text) else ""
        return f"{prefix}{text[start:end].strip()}{suffix}"

    return success_response(data=[
        {
            "result_id": row.id,
            "session_id": row.session_id,
            "message_id": row.id,
            "role": row.role,
            "content": row.content,
            "snippet": _build_snippet(row.content, keyword),
            "session_title": row.title,
            "project_id": row.project_id,
            "project_name": row.project_name,
            "llm_config_id": row.llm_config_id,
            "model_name": row.model_name,
            "created_at": row.created_at.isoformat() if row.created_at else None,
        }
        for row in rows
    ])


@router.post("/send", response_model=dict)
async def send_chat_message(
    req: ChatMessageRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """发送对话消息并获取回复"""
    question = (req.message or "").strip()
    if not question:
        raise HTTPException(status_code=400, detail="消息内容不能为空")

    stmt = select(ChatSession).where(ChatSession.id == req.session_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=404, detail="对话会话不存在")

    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权访问该对话")

    await require_project_read_access(db, session.project_id, current_user)

    query_message = ChatMessage(
        session_id=req.session_id,
        role="user",
        content=question,
        status="completed",
    )
    db.add(query_message)
    await db.flush()

    try:
        msg_stmt = (
            select(ChatMessage)
            .where(
                ChatMessage.session_id == req.session_id,
                ChatMessage.id != query_message.id,
            )
            .order_by(ChatMessage.created_at.desc())
            .limit(CHAT_HISTORY_MESSAGE_LIMIT)
        )
        msg_result = await db.execute(msg_stmt)
        prev_messages = list(reversed(msg_result.scalars().all()))

        conversation_history = [
            {"role": m.role, "content": m.content}
            for m in prev_messages
            # 被中断的助手消息不参与上下文，避免污染后续模型输入
            if not _is_interrupted_message(m)
        ]

        start_time = time.monotonic()
        thinking_log = [
            {
                "stage": "retrieval",
                "message": "正在检索知识库…",
                "elapsed_ms": 0,
            }
        ]
        retrieved_docs = await rag_service.retrieve_documents(
            db,
            session.project_id,
            question,
            top_k=5,
        )
        thinking_log.append({
            "stage": "retrieved",
            "message": f"已检索到 {len(retrieved_docs)} 个相关文档",
            "count": len(retrieved_docs),
            "elapsed_ms": int((time.monotonic() - start_time) * 1000),
        })
        thinking_log.append({
            "stage": "generate",
            "message": "正在生成回答…",
            "elapsed_ms": int((time.monotonic() - start_time) * 1000),
        })

        assistant_response = await rag_service.generate_response(
            db,
            question,
            session.project_id,
            session.llm_config_id,
            retrieved_docs,
            conversation_history,
        )
        thinking_log.append({
            "stage": "align",
            "message": "正在整理引用…",
            "elapsed_ms": int((time.monotonic() - start_time) * 1000),
        })

        # 仅保留实际引用的文档，并把引用编号压缩为从 1 开始的连续序列。
        assistant_response, cited_refs = _compact_cited_refs(
            assistant_response, retrieved_docs
        )
        if cited_refs:
            # 路线 B：embedding 对齐，为每个引用回填原文支撑句
            await rag_service.align_citation_quotes(
                db, assistant_response, cited_refs
            )

        assistant_message = ChatMessage(
            session_id=req.session_id,
            role="assistant",
            content=assistant_response,
            status="completed",
            duration_ms=int((time.monotonic() - start_time) * 1000),
            thinking_log=json.dumps(thinking_log, ensure_ascii=False),
            referenced_files=json.dumps(cited_refs, ensure_ascii=False) if cited_refs else None,
        )
        db.add(assistant_message)

        # 累加会话消息计数（用户 + 助手 共 2 条）
        session.message_count = (session.message_count or 0) + 2

        await db.commit()

        return success_response(data={
            "session_id": req.session_id,
            "user_message": question,
            "assistant_message": assistant_response,
            "referenced_files": [ref["file_path"] for ref in cited_refs],
            "references": _build_reference_items(cited_refs, session.project_id),
        })

    except Exception as e:
        await db.rollback()
        raise HTTPException(
            status_code=500,
            detail=f"生成对话回复失败: {str(e)}"
        )


@router.post("/send/stream", response_model=dict)
async def send_chat_message_stream(
    req: ChatMessageRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """发送对话消息并流式返回回复"""
    question = (req.message or "").strip()
    if not question:
        raise HTTPException(status_code=400, detail="消息内容不能为空")

    stmt = select(ChatSession).where(ChatSession.id == req.session_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=404, detail="对话会话不存在")

    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权访问该对话")

    await require_project_read_access(db, session.project_id, current_user)

    # 先读取历史消息（不含本轮提问），用于多轮上下文与判断是否首轮
    msg_stmt = (
        select(ChatMessage)
        .where(ChatMessage.session_id == req.session_id)
        .order_by(ChatMessage.created_at.desc())
        .limit(CHAT_HISTORY_MESSAGE_LIMIT)
    )
    msg_result = await db.execute(msg_stmt)
    prev_messages = list(reversed(msg_result.scalars().all()))

    conversation_history = [
        {"role": m.role, "content": m.content}
        for m in prev_messages
        # 被中断的助手消息不参与上下文，避免污染后续模型输入
        if not _is_interrupted_message(m)
    ]

    # 先持久化用户提问和空的助手消息占位行，再执行向量检索。
    # 占位行先入库后，即使客户端在检索阶段断开（停止生成），
    # 前端也能拿到真实消息 id 并正确标记“已停止”，而不会留下悬空占位。
    assistant_message = ChatMessage(
        session_id=req.session_id,
        role="assistant",
        content="",
    )
    if req.insert_user_message:
        query_message = ChatMessage(
            session_id=req.session_id,
            role="user",
            content=question,
            status="completed",
        )
        db.add(query_message)
    db.add(assistant_message)
    session.message_count = (session.message_count or 0) + (2 if req.insert_user_message else 1)
    await db.commit()

    user_message_id = query_message.id if req.insert_user_message else None
    assistant_message_id = assistant_message.id

    llm_config_id = session.llm_config_id
    project_id = session.project_id
    # 生成总耗时起点：从消息占位行入库后开始计时，用于“思考过程”展示
    start_time = time.monotonic()
    # 每累计这么多字符就回写一次占位行，平衡「刷新可见性」与「写库频率」
    FLUSH_EVERY_CHARS = 120

    async def event_generator():
        assistant_response_parts = []
        thinking_log: List[Dict[str, Any]] = []
        chars_since_flush = 0
        finished = False

        def _log(
            stage: str,
            message: str,
            *,
            duration_ms: Optional[int] = None,
            count: Optional[int] = None,
        ) -> Dict[str, Any]:
            """记录一个思考过程条目。

            duration_ms 表示该阶段已完成的实际耗时；进行中的阶段不填，
            由前端显示为进度状态，避免“开始时间点”造成的误导。
            """
            entry = {
                "stage": stage,
                "message": message,
            }
            if duration_ms is not None:
                entry["duration_ms"] = duration_ms
            if count is not None:
                entry["count"] = count
            thinking_log.append(entry)
            return entry

        async def _persist_assistant(parts, *, completed):
            """把已生成内容回写到占位的助手消息行。

            completed=True 时附带引用解析与最终状态；False 表示流式中途的增量回写。
            """
            assistant_response = "".join(parts)
            cited_refs = []
            if completed:
                assistant_response, cited_refs = _compact_cited_refs(
                    assistant_response, retrieved_docs
                )
                if cited_refs:
                    # 路线 B：embedding 对齐，为每个引用回填原文支撑句
                    await rag_service.align_citation_quotes(
                        db, assistant_response, cited_refs
                    )
            references = _build_reference_items(cited_refs, project_id)

            msg = await db.get(ChatMessage, assistant_message_id)
            if msg is not None:
                msg.content = assistant_response
                if completed:
                    msg.status = "completed"
                    msg.duration_ms = int((time.monotonic() - start_time) * 1000)
                    msg.thinking_log = (
                        json.dumps(thinking_log, ensure_ascii=False)
                        if thinking_log else None
                    )
                    msg.referenced_files = (
                        json.dumps(cited_refs, ensure_ascii=False) if cited_refs else None
                    )
            await db.commit()
            return assistant_response, references

        try:
            # 先回传两条消息的真实 id，使前端无需刷新即可获得删除入口等能力
            yield _stream_event("ids", {
                "session_id": req.session_id,
                "user_message_id": user_message_id,
                "assistant_message_id": assistant_message_id,
            })

            # 检索阶段放在流内执行，客户端可以在检索过程中实时看到思考进度
            yield _stream_event("thinking", _log("retrieval", "正在检索知识库…"))
            retrieval_start = time.monotonic()
            retrieved_docs = await rag_service.retrieve_documents(
                db,
                session.project_id,
                question,
                top_k=5,
            )
            yield _stream_event(
                "thinking",
                _log(
                    "retrieved",
                    f"已检索到 {len(retrieved_docs)} 个相关文档",
                    duration_ms=int((time.monotonic() - retrieval_start) * 1000),
                    count=len(retrieved_docs),
                ),
            )

            yield _stream_event("thinking", _log("generate", "正在生成回答…"))
            generate_start = time.monotonic()
            async for chunk in rag_service.generate_response_stream(
                db,
                question,
                project_id,
                llm_config_id,
                retrieved_docs,
                conversation_history,
            ):
                assistant_response_parts.append(chunk)
                chars_since_flush += len(chunk)
                yield _stream_event("chunk", {"content": chunk})

                # 边生成边落库：达到阈值就把当前进度回写到占位行
                if chars_since_flush >= FLUSH_EVERY_CHARS:
                    chars_since_flush = 0
                    await _persist_assistant(assistant_response_parts, completed=False)

            yield _stream_event(
                "thinking",
                _log(
                    "generated",
                    "回答生成完成",
                    duration_ms=int((time.monotonic() - generate_start) * 1000),
                ),
            )

            align_start = time.monotonic()
            yield _stream_event("thinking", _log("align", "正在整理引用…"))
            assistant_response, references = await _persist_assistant(
                assistant_response_parts, completed=True
            )
            finished = True
            yield _stream_event(
                "thinking",
                _log(
                    "aligned",
                    "引用整理完成",
                    duration_ms=int((time.monotonic() - align_start) * 1000),
                ),
            )
            duration_ms = int((time.monotonic() - start_time) * 1000)

            yield _stream_event("references", references or [])
            yield _stream_event("done", {
                "session_id": req.session_id,
                "user_message_id": user_message_id,
                "assistant_message_id": assistant_message_id,
                "content": assistant_response,
                "references": references or [],
                "status": "completed",
                "duration_ms": duration_ms,
                "thinking_log": thinking_log,
            })
        except (asyncio.CancelledError, GeneratorExit):
            # 客户端中途断开（停止生成/刷新页面）：把已生成内容与中断状态回写占位行。
            if not finished:
                try:
                    await db.rollback()
                    msg = await db.get(ChatMessage, assistant_message_id)
                    if msg is not None:
                        msg.content = "".join(assistant_response_parts)
                        msg.status = "interrupted"
                        msg.duration_ms = int((time.monotonic() - start_time) * 1000)
                        msg.thinking_log = (
                            json.dumps(thinking_log, ensure_ascii=False)
                            if thinking_log else None
                        )
                        await db.commit()
                except Exception as save_exc:  # noqa: BLE001
                    logger.warning("保存中断的助手回复失败: %s", save_exc)
            raise
        except Exception as exc:
            await db.rollback()
            if not assistant_response_parts:
                assistant_response_parts.append("回答生成失败，请稍后重试。")
            try:
                msg = await db.get(ChatMessage, assistant_message_id)
                if msg is not None:
                    msg.content = "".join(assistant_response_parts)
                    msg.status = "error"
                    msg.duration_ms = int((time.monotonic() - start_time) * 1000)
                    msg.thinking_log = (
                        json.dumps(thinking_log, ensure_ascii=False)
                        if thinking_log else None
                    )
                    await db.commit()
            except Exception as save_exc:  # noqa: BLE001
                logger.warning("保存失败的助手回复状态失败: %s", save_exc)
            yield _stream_event("error", {"detail": f"生成对话回复失败: {str(exc)}"})

    return StreamingResponse(event_generator(), media_type="text/event-stream")


@router.delete("/sessions/{session_id}", response_model=dict)
async def delete_chat_session(
    session_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """删除对话会话"""
    stmt = select(ChatSession).where(ChatSession.id == session_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=404, detail="对话会话不存在")

    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权删除该对话")

    await db.delete(session)
    msg_stmt = delete(ChatMessage).where(ChatMessage.session_id == session_id)
    await db.execute(msg_stmt)
    await db.commit()

    return success_response(message="对话会话已删除")


@router.delete("/messages/{message_id}", response_model=dict)
async def delete_chat_message(
    message_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """删除单条对话消息"""
    stmt = (
        select(ChatMessage, ChatSession)
        .join(ChatSession, ChatMessage.session_id == ChatSession.id)
        .where(ChatMessage.id == message_id)
    )
    result = await db.execute(stmt)
    row = result.first()

    if not row:
        raise HTTPException(status_code=404, detail="消息不存在")

    msg, session = row
    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权删除该消息")

    await db.delete(msg)
    if session.message_count and session.message_count > 0:
        session.message_count = max(0, session.message_count - 1)
    await db.commit()

    return success_response(message="消息已删除")


@router.post("/messages/{message_id}/interrupt", response_model=dict)
async def mark_message_interrupted(
    message_id: int,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """把生成被中断的助手消息标记为已停止。

    前端在点击“停止生成”后立即调用本接口，把统一中断标记写入数据库，
    避免依赖连接断开后生成器的兜底写入（存在时间差），保证切换会话或刷新后
    仍能识别该消息为“已停止生成”。
    """
    stmt = (
        select(ChatMessage, ChatSession)
        .join(ChatSession, ChatMessage.session_id == ChatSession.id)
        .where(ChatMessage.id == message_id)
    )
    result = await db.execute(stmt)
    row = result.first()

    if not row:
        raise HTTPException(status_code=404, detail="消息不存在")

    msg, session = row
    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权操作该消息")
    if msg.role != "assistant":
        raise HTTPException(status_code=400, detail="仅助手消息可标记为已停止")

    # 只更新状态标志，不再向回复内容里追加标记文本
    msg.status = "interrupted"
    await db.commit()

    return success_response(data={
        "message_id": msg.id,
        "status": msg.status,
        "duration_ms": msg.duration_ms,
    })


@router.put("/sessions/{session_id}", response_model=dict)
async def update_chat_session(
    session_id: int,
    req: ChatSessionUpdateRequest,
    current_user: User = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """更新对话会话标题"""
    stmt = select(ChatSession).where(ChatSession.id == session_id)
    result = await db.execute(stmt)
    session = result.scalar_one_or_none()

    if not session:
        raise HTTPException(status_code=404, detail="对话会话不存在")

    if session.user_id != current_user.id:
        raise HTTPException(status_code=403, detail="无权修改该对话")

    title = (req.title or "").strip()
    if not title:
        raise HTTPException(status_code=400, detail="标题不能为空")

    session.title = title
    await db.commit()

    return success_response(message="标题已更新", data={"session_id": session_id, "title": title})