my_bookmate/backend/app/api/chat.py

import json
import logging
import re
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func
from pydantic import BaseModel

from app.core.database import get_db
from app.models.book import Book, Chapter, Segment
from app.models.chat import Conversation, Message
from app.llm.hub import resolve_provider
from app.rag.engine import (
    search_segments, build_chat_prompt, parse_citations,
    estimate_tokens, match_chapters, _trim_messages,
)
from app.mcp.agent import agent_stream, build_mcp_context, execute_mcp_tool
from app.assistant.service import stream_assistant

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/chat", tags=["chat"])


class CreateConversationRequest(BaseModel):
    # kind="book" 绑定单本书（book_id 必填）；kind="assistant" 为全局「第二大脑」对话
    kind: str = "book"
    book_id: str | None = None
    title: str | None = None
    model: str = "claude"


class SendMessageRequest(BaseModel):
    content: str
    model: str | None = None
    # 被重试的用户消息 id：提供时不新建 user 消息，新 AI 回答作为该问题的新版本
    # （parent_message_id 指向它），历史上下文只取该问题之前的对话
    reply_to: str | None = None


class ConversationOut(BaseModel):
    id: str
    kind: str
    book_id: str | None
    title: str | None
    model: str
    created_at: str

    class Config:
        from_attributes = True


class MessageOut(BaseModel):
    id: str
    conversation_id: str
    parent_message_id: str | None
    role: str
    content: str
    citations: list | None
    model: str | None
    created_at: str

    class Config:
        from_attributes = True


@router.get("/conversations", response_model=list[ConversationOut])
async def list_conversations(
    book_id: str | None = None,
    kind: str | None = None,
    db: AsyncSession = Depends(get_db),
):
    q = select(Conversation).order_by(Conversation.updated_at.desc())
    if book_id:
        q = q.where(Conversation.book_id == book_id)
    if kind:
        q = q.where(Conversation.kind == kind)
    result = await db.execute(q)
    convs = result.scalars().all()
    return [ConversationOut(
        id=c.id, kind=c.kind, book_id=c.book_id, title=c.title, model=c.model,
        created_at=c.created_at.isoformat() if c.created_at else "",
    ) for c in convs]


@router.post("/conversations", response_model=ConversationOut)
async def create_conversation(req: CreateConversationRequest, db: AsyncSession = Depends(get_db)):
    if req.kind == "assistant":
        # 全局「第二大脑」对话：不绑定书
        book_id = None
        title = req.title or "与 AI 助手的对话"
    else:
        book = await db.get(Book, req.book_id)
        if not book:
            raise HTTPException(404, "书不存在")
        book_id = req.book_id
        title = req.title or f"与《{book.title}》的对话"
    conv = Conversation(book_id=book_id, title=title, model=req.model, kind=req.kind)
    db.add(conv)
    await db.commit()
    await db.refresh(conv)
    return ConversationOut(
        id=conv.id, kind=conv.kind, book_id=conv.book_id, title=conv.title, model=conv.model,
        created_at=conv.created_at.isoformat() if conv.created_at else "",
    )


@router.get("/conversations/{conv_id}/messages", response_model=list[MessageOut])
async def get_messages(conv_id: str, db: AsyncSession = Depends(get_db)):
    conv = await db.get(Conversation, conv_id)
    if not conv:
        raise HTTPException(404, "会话不存在")
    result = await db.execute(
        select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
    )
    msgs = result.scalars().all()
    return [_msg_to_out(m) for m in msgs]


@router.post("/conversations/{conv_id}/messages")
async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession = Depends(get_db)):
    conv = await db.get(Conversation, conv_id)
    if not conv:
        raise HTTPException(404, "会话不存在")

    model_name = req.model or conv.model

    # 重试：复用被重试的用户消息（不新建），新回答作为它的新版本
    reply_user: Message | None = None
    if req.reply_to:
        reply_user = await db.get(Message, req.reply_to)
        if not reply_user or reply_user.role != "user" or reply_user.conversation_id != conv_id:
            raise HTTPException(400, "reply_to 必须是本会话内的一条用户消息")
        user_msg = reply_user
    else:
        user_msg = Message(conversation_id=conv_id, role="user", content=req.content, model=model_name)
        db.add(user_msg)
        await db.commit()

    # 全局「第二大脑」对话：book_id 为空，走独立管线（跨书检索 + 知识库概览 + 原生工具 + MCP）
    if conv.kind == "assistant":
        provider = await resolve_provider(db, model_name)
        if provider is None:
            raise HTTPException(400, f"模型 {model_name} 未配置，请先在页面「管理模型」中添加或选择可用模型")
        return await stream_assistant(conv, user_msg, req.content, db, provider)

    book = await db.get(Book, conv.book_id)
    if not book:
        raise HTTPException(404, "书不存在")

    chapters = (await db.execute(
        select(Chapter).where(Chapter.book_id == book.id).order_by(Chapter.order_idx)
    )).scalars().all()

    # 章节标题匹配：命中章节时把整章原文注入 prompt（解决「第X章讲了什么」这类问题）
    matched = match_chapters(list(chapters), req.content)
    chapter_segments = []
    if matched:
        result = await db.execute(
            select(Segment)
            .where(Segment.chapter_id.in_([c.id for c in matched]))
            .order_by(Segment.order_idx)
        )
        # 先按章节目录顺序、再按章内顺序排列，否则各章的开头段会交错排在一起
        chapter_rank = {c.id: c.order_idx for c in matched}
        chapter_segments = sorted(
            result.scalars().all(),
            key=lambda s: (chapter_rank[s.chapter_id], s.order_idx),
        )

    segments = await search_segments(
        db, book.id, req.content, top_k=15, chapters=list(chapters)
    )

    history_rows = (await db.execute(
        select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
    )).scalars().all()

    # 历史上下文：重试时只取被重试问题之前的对话（旧版本回答不进 prompt）；
    # 正常发送时排除当前这轮 user 消息即可。
    if reply_user is not None:
        row_ids = {m.id for m in history_rows}
        if reply_user.id in row_ids:
            idx = next(i for i, m in enumerate(history_rows) if m.id == reply_user.id)
            history = [{"role": m.role, "content": m.content} for m in history_rows[:idx]]
        else:
            # 数据异常（如 SRU），退回为全量历史
            history = [{"role": m.role, "content": m.content} for m in history_rows]
    else:
        history = [{"role": m.role, "content": m.content} for m in history_rows[:-1]]

    provider = await resolve_provider(db, model_name)
    if provider is None:
        raise HTTPException(400, f"模型 {model_name} 未配置，请先在页面「管理模型」中添加或选择可用模型")

    messages = build_chat_prompt(
        book, list(chapters), segments, history, req.content,
        chapter_segments=chapter_segments,
        # 整章原文预算随模型上下文放大：一整章常有上万 token，固定 6000 会截掉后面的小节
        chapter_budget=max(6000, int(provider.max_context * 0.4)),
    )

    # MCP 外部工具：仅当模型原生支持 function calling 时接入工具循环；
    # 无 MCP 配置 / 单服务器失败时自动跳过，行为与改造前完全一致
    mcp_ctx = await build_mcp_context(db) if provider.supports_tools else None
    if mcp_ctx is not None:
        system = messages[0]["content"]
        system += (
            "\n\n## 外部信息工具\n"
            "当问题需要书中没有的外部信息（最新动态、时事、网络资料、行业数据等）时，"
            f"可调用以下外部 MCP 工具获取：{', '.join(mcp_ctx.names)}\n"
            "- 外部信息要说明来源与时效；回答主体仍以书中内容为准，并保留 [[ref:N]] 引用标注。\n"
        )
        messages[0]["content"] = system

    total_tokens = estimate_tokens("".join(m["content"] for m in messages))
    if total_tokens > provider.max_context * 0.8:
        # 内容感知裁剪：保留 system（含检索原文 + 整章注入），从最旧历史开始丢
        messages = _trim_messages(messages, int(provider.max_context * 0.8))

    async def stream_response():
        full_content = ""
        try:
            if mcp_ctx is not None:
                async for kind, payload in agent_stream(
                    messages,
                    provider,
                    mcp_ctx.ai_tools,
                    execute=lambda name, args: execute_mcp_tool(mcp_ctx, name, args),
                    describe=mcp_ctx.display_name,
                ):
                    if kind == "token":
                        full_content += payload
                        yield f"data: {json.dumps({'type': 'token', 'content': payload}, ensure_ascii=False)}\n\n"
                    elif kind == "tool":
                        yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'start'}, ensure_ascii=False)}\n\n"
                    elif kind == "tool_done":
                        yield f"data: {json.dumps({'type': 'tool', 'name': payload, 'status': 'done'}, ensure_ascii=False)}\n\n"
            else:
                async for token in provider.chat_stream(messages, temperature=0.7):
                    full_content += token
                    yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
        except Exception as e:
            logger.exception("流式生成失败（conversation=%s）", conv_id)
            yield f"data: {json.dumps({'type': 'error', 'content': str(e)}, ensure_ascii=False)}\n\n"
            return

        parsed = parse_citations(full_content, segments)

        assistant_msg = Message(
            conversation_id=conv_id,
            parent_message_id=user_msg.id,
            role="assistant",
            content=parsed["text"],
            citations=json.dumps(parsed["citations"], ensure_ascii=False),
            model=model_name,
        )
        db.add(assistant_msg)
        await db.commit()
        await db.refresh(assistant_msg)

        yield f"data: {json.dumps({'type': 'done', 'message_id': assistant_msg.id, 'citations': parsed['citations']}, ensure_ascii=False)}\n\n"

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


@router.post("/messages/{msg_id}/retry")
async def retry_message(msg_id: str, db: AsyncSession = Depends(get_db)):
    msg = await db.get(Message, msg_id)
    if not msg or msg.role != "user":
        raise HTTPException(400, "只能对用户消息重试")
    # 走 reply_to 语义：不新建 user 消息，而是给该问题追加一个新版本回答
    return await send_message(
        msg.conversation_id,
        SendMessageRequest(content=msg.content, model=msg.model, reply_to=msg.id),
        db,
    )


@router.get("/segments/{segment_id}/context")
async def segment_context(
    segment_id: str,
    before: int = Query(2, ge=0, le=6),
    after: int = Query(2, ge=0, le=6),
    db: AsyncSession = Depends(get_db),
):
    """引用段落的阅读上下文：章节标题 + 前后段落，供前端「阅读原文」弹层。"""
    seg = await db.get(Segment, segment_id)
    if not seg:
        raise HTTPException(404, "段落不存在")
    ch = await db.get(Chapter, seg.chapter_id)

    prev_rows = await db.execute(
        select(Segment.id, Segment.text, Segment.order_idx, Segment.cfi)
        .where(Segment.chapter_id == seg.chapter_id, Segment.order_idx < seg.order_idx)
        .order_by(Segment.order_idx.desc())
        .limit(before)
    )
    next_rows = await db.execute(
        select(Segment.id, Segment.text, Segment.order_idx, Segment.cfi)
        .where(Segment.chapter_id == seg.chapter_id, Segment.order_idx > seg.order_idx)
        .order_by(Segment.order_idx)
        .limit(after)
    )
    prev_list = [{"id": r[0], "text": r[1], "order_idx": r[2], "cfi": r[3]} for r in prev_rows.all()]
    next_list = [{"id": r[0], "text": r[1], "order_idx": r[2], "cfi": r[3]} for r in next_rows.all()]
    prev_list.reverse()  # 保持阅读顺序

    return {
        "segment": {"id": seg.id, "text": seg.text, "order_idx": seg.order_idx, "cfi": seg.cfi},
        "chapter": {"id": ch.id, "title": ch.title, "level": ch.level} if ch else None,
        "before": prev_list,
        "after": next_list,
    }


@router.get("/search")
async def search_book(book_id: str, q: str = Query(min_length=1), db: AsyncSession = Depends(get_db)):
    segments = await search_segments(db, book_id, q, top_k=10)
    return [
        {"id": s.id, "text": s.text[:300], "cfi": s.cfi, "char_count": s.char_count}
        for s in segments
    ]


@router.delete("/messages/{msg_id}")
async def delete_message(msg_id: str, db: AsyncSession = Depends(get_db)):
    msg = await db.get(Message, msg_id)
    if not msg:
        raise HTTPException(404, "消息不存在")
    # 级联删除：该消息名下的所有回答版本（parent 指向它），避免残留孤儿版本
    versions = (await db.execute(
        select(Message).where(Message.parent_message_id == msg_id)
    )).scalars().all()
    for v in versions:
        await db.delete(v)
    await db.delete(msg)
    await db.commit()
    return {"deleted": msg_id, "cascade": len(versions)}


@router.delete("/conversations/{conv_id}")
async def delete_conversation(conv_id: str, db: AsyncSession = Depends(get_db)):
    conv = await db.get(Conversation, conv_id)
    if not conv:
        raise HTTPException(404, "会话不存在")
    # ORM cascade="all, delete-orphan" 会一并删除该会话下的所有消息
    await db.delete(conv)
    await db.commit()
    return {"deleted": conv_id}


@router.post("/conversations/{conv_id}/branch")
async def branch_conversation(conv_id: str, from_message_id: str = Query(...), db: AsyncSession = Depends(get_db)):
    conv = await db.get(Conversation, conv_id)
    if not conv:
        raise HTTPException(404, "会话不存在")
    branch = Conversation(book_id=conv.book_id, title=f"{conv.title or '对话'} (分支)", model=conv.model, kind=conv.kind)
    db.add(branch)
    await db.flush()

    result = await db.execute(
        select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
    )
    msgs = result.scalars().all()
    for m in msgs:
        new_msg = Message(
            conversation_id=branch.id,
            role=m.role,
            content=m.content,
            citations=m.citations,
            model=m.model,
        )
        db.add(new_msg)
        if m.id == from_message_id:
            break
    await db.commit()
    await db.refresh(branch)
    return {
        "id": branch.id,
        "book_id": branch.book_id,
        "title": branch.title,
        "model": branch.model,
        "created_at": branch.created_at.isoformat() if branch.created_at else "",
    }


@router.get("/conversations/{conv_id}/export")
async def export_conversation(conv_id: str, fmt: str = "markdown", db: AsyncSession = Depends(get_db)):
    conv = await db.get(Conversation, conv_id)
    if not conv:
        raise HTTPException(404, "会话不存在")
    result = await db.execute(
        select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
    )
    msgs = result.scalars().all()
    book = await db.get(Book, conv.book_id) if conv.book_id else None

    book_label = "全局助手" if conv.kind == "assistant" else (book.title if book else "未知")
    lines = [f"# {conv.title or '对话'}", f"书籍: {book_label}\n"]
    for m in msgs:
        role_label = "🧑 用户" if m.role == "user" else "🤖 AI"
        # 把 mermaid mindmap 转成 flowchart LR：老版阅读器渲染 mindmap 是放射状粗线，很难读
        content = _transform_mindmaps(m.content)
        lines.append(f"## {role_label}\n\n{content}\n")
        if m.citations:
            try:
                citations = json.loads(m.citations)
                for i, c in enumerate(citations):
                    lines.append(f"> 引用 {i+1}: {c.get('excerpt', '')}\n")
            except (json.JSONDecodeError, TypeError):
                pass
    return {"content": "\n".join(lines), "format": fmt, "title": conv.title}


_FENCE_RE = re.compile(r"```(mermaid)?\s*\n(.*?)\n```", re.DOTALL)


def _transform_mindmaps(text: str) -> str:
    """把 markdown 里的 ```mermaid mindmap 块转换成 flowchart LR（清晰左→右树状）。"""
    def _repl(m: re.Match) -> str:
        body = m.group(2)
        if body.lstrip().startswith("mindmap"):
            return f"```mermaid\n{_mindmap_to_flowchart(body)}\n```"
        return m.group(0)

    return _FENCE_RE.sub(_repl, text)


_LABEL_QUOTES = {"d": True, "s": True}


def _safe_label(label: str) -> str:
    """把会破坏老版 mermaid flowchart 语法的字符替换为可安全渲染的等价字符。

    - ASCII 双/单引号 → 弯引号（“”/‘’）：老版 mermaid 不认 ["...\\"..."] 转义，
      引号会把字符串截断导致 Parse error；弯引号是普通 Unicode，可正常渲染且保留语义。
    - 反斜杠 → 斜杠；# → 全角＃（老版 mermaid 中 # 有特殊含义）。
    """
    states = dict(_LABEL_QUOTES)

    def _rep(m: re.Match) -> str:
        ch = m.group(0)
        key = "d" if ch == '"' else "s"
        pair = ("“", "”") if key == "d" else ("‘", "’")
        idx = 0 if states[key] else 1
        states[key] = not states[key]
        return pair[idx]

    label = re.sub(r'["\']', _rep, label)
    return label.replace("\\", "/").replace("#", "＃")


def _mindmap_to_flowchart(body: str) -> str:
    """mindmap（缩进式）→ flowchart LR（节点 + 边）。

    mindmap:
      root((主题))
        一级分支
          二级分支
    → flowchart LR:
      n1["主题"]
      n1 --> n2["一级分支"]
      n2 --> n3["二级分支"]
    """
    lines = []
    for raw in body.splitlines():
        s = raw.strip()
        if not s or s == "mindmap":
            continue
        lines.append(raw)
    if not lines:
        return "flowchart LR\n  n1[空思维导图]"

    out = ["flowchart LR"]
    stack: list[tuple[int, str]] = []  # (缩进长度, 节点 id)
    seq = 0

    def _new_id() -> str:
        nonlocal seq
        seq += 1
        return f"n{seq}"

    for raw in lines:
        m = re.match(r"[ \t]*", raw)
        indent_len = len(m.group(0).replace("\t", "  "))
        text = raw[m.end():].strip()
        # 去掉 ::icon(...) 与行尾 # 注释
        text = re.sub(r"::icon\([^)]*\)", "", text)
        text = re.split(r"\s+#\s*", text)[0].strip()
        # 去掉 mindmap 的 root 关键字：root((主题)) 里的 root 只是根标记，标签在形状符号内
        if text.startswith("root") and len(text) > 4:
            rest = text[4:].lstrip()
            if rest and rest[0] in "([{":
                text = rest
        # 去掉 mindmap 形状符号包裹：(()) / [[ ]] / [ ] / ( ) / { }
        for open_, close in (("((", "))"), ("[[", "]]"), ("[", "]"), ("(", ")"), ("{", "}")):
            if text.startswith(open_) and text.endswith(close):
                text = text[len(open_):-len(close)].strip()
                break
        label = text.strip() or _new_id()

        while stack and stack[-1][0] >= indent_len:
            stack.pop()
        nid = _new_id()
        safe = _safe_label(label).replace("\n", " ")
        out.append(f'  {nid}["{safe}"]')
        if stack:
            out.append(f"  {stack[-1][1]} --> {nid}")
        stack.append((indent_len, nid))

    return "\n".join(out)


def _msg_to_out(m: Message) -> MessageOut:
    citations = None
    if m.citations:
        try:
            citations = json.loads(m.citations)
        except (json.JSONDecodeError, TypeError):
            citations = None
    return MessageOut(
        id=m.id, conversation_id=m.conversation_id, parent_message_id=m.parent_message_id,
        role=m.role, content=m.content, citations=citations, model=m.model,
        created_at=m.created_at.isoformat() if m.created_at else "",
    )