import asyncio
import logging
import re
from collections.abc import AsyncIterator, Sequence
from contextlib import suppress
from typing import Any

from langchain.messages import AIMessage, AIMessageChunk, HumanMessage
from sqlalchemy.ext.asyncio import async_sessionmaker

from kbqa.agents.context import ChatRunContext
from kbqa.agents.factory import AgentBundle, create_agent_bundle
from kbqa.api.errors import AppError
from kbqa.chat.repository import ChatRepository, MessageSourceInput
from kbqa.chat.sse import encode_event, encode_ping
from kbqa.config import Settings
from kbqa.rag.retriever import KnowledgeRetriever

INSUFFICIENT_ANSWER = "知识库中没有足够资料回答这个问题。"
logger = logging.getLogger(__name__)


class ChatService:
    def __init__(
        self,
        settings: Settings,
        session_factory: async_sessionmaker,
        retriever: KnowledgeRetriever,
    ) -> None:
        self.settings = settings
        self.session_factory = session_factory
        self.retriever = retriever
        self._active_sessions: set[str] = set()
        self._agents: AgentBundle | None = None

    def reserve(self, session_id: str) -> None:
        if session_id in self._active_sessions:
            raise AppError("SESSION_BUSY", "当前会话正在生成回答", 409)
        self._active_sessions.add(session_id)

    def _get_agents(self) -> AgentBundle:
        if self._agents is None:
            self._agents = create_agent_bundle(self.settings)
        return self._agents

    async def stream(
        self,
        session_id: str,
        query: str,
        request_id: str,
    ) -> AsyncIterator[bytes]:
        assistant_id: str | None = None
        graph_stream = None
        try:
            async with self.session_factory() as session:
                repository = ChatRepository(session)
                if await repository.get_session(session_id) is None:
                    yield encode_event(
                        "error",
                        {
                            "code": "SESSION_NOT_FOUND",
                            "message": "会话不存在",
                            "retryable": False,
                        },
                    )
                    return
                _, assistant = await repository.begin_turn(session_id, query)
                assistant_id = assistant.id
                history = await repository.load_recent_messages(
                    session_id, self.settings.history_message_limit
                )

            messages = [
                HumanMessage(content=message.content)
                if message.role == "user"
                else AIMessage(content=message.content)
                for message in history
                if message.id != assistant_id
            ]
            sources: list[dict[str, Any]] = []
            answer_parts: list[str] = []
            answering = False
            no_evidence = False
            yield encode_event("status", {"phase": "researching"})

            graph_stream = self._get_agents().main_agent.astream(
                {"messages": messages},
                context=ChatRunContext(
                    retriever=self.retriever,
                    request_id=request_id,
                    session_id=session_id,
                    assistant_message_id=assistant_id,
                ),
                config={
                    "recursion_limit": 12,
                    "tags": ["main-agent"],
                    "metadata": {
                        "request_id": request_id,
                        "session_id": session_id,
                        "assistant_message_id": assistant_id,
                    },
                },
                stream_mode=["messages", "custom", "updates"],
                version="v2",
            )
            iterator = graph_stream.__aiter__()
            pending = asyncio.create_task(anext(iterator))
            while True:
                done, _ = await asyncio.wait({pending}, timeout=15.0)
                if not done:
                    yield encode_ping()
                    continue
                try:
                    chunk = pending.result()
                except StopAsyncIteration:
                    break
                pending = asyncio.create_task(anext(iterator))
                chunk_type = chunk.get("type")
                data = chunk.get("data")
                if chunk_type == "custom" and isinstance(data, dict):
                    if data.get("kind") == "no_evidence":
                        no_evidence = True
                        pending.cancel()
                        with suppress(asyncio.CancelledError):
                            await pending
                        await graph_stream.aclose()
                        graph_stream = None
                        break
                    if data.get("kind") == "sources":
                        sources = list(data.get("sources", []))
                        yield encode_event("sources", {"sources": sources})
                        yield encode_event("status", {"phase": "answering"})
                        answering = True
                    continue
                if chunk_type != "messages" or not answering:
                    continue
                token, _metadata = data
                if not isinstance(token, AIMessageChunk):
                    continue
                text = self._message_text(token.content)
                if text:
                    answer_parts.append(text)
                    yield encode_event("token", {"content": text})

            if no_evidence:
                answer = INSUFFICIENT_ANSWER
                yield encode_event("status", {"phase": "answering"})
                yield encode_event("token", {"content": answer})
            else:
                answer = "".join(answer_parts).strip()
                self._validate_citations(answer, sources)

            source_inputs = [
                MessageSourceInput(
                    document_id=str(source["document_id"]),
                    chunk_id=str(source["chunk_id"]),
                    citation_label=str(source["label"]),
                    score=float(source["score"]),
                )
                for source in sources
            ]
            async with self.session_factory() as session:
                await ChatRepository(session).complete_answer(assistant_id, answer, source_inputs)
            yield encode_event("done", {"message_id": assistant_id})
        except asyncio.CancelledError:
            if assistant_id is not None:
                await self._mark_failed(assistant_id, "CLIENT_DISCONNECTED")
            raise
        except AppError as exc:
            if assistant_id is not None:
                await self._mark_failed(assistant_id, exc.code)
            yield encode_event(
                "error",
                {"code": exc.code, "message": exc.message, "retryable": exc.retryable},
            )
        except Exception:
            logger.exception("Chat stream failed request_id=%s", request_id)
            if assistant_id is not None:
                await self._mark_failed(assistant_id, "MODEL_UPSTREAM_ERROR")
            yield encode_event(
                "error",
                {
                    "code": "MODEL_UPSTREAM_ERROR",
                    "message": "模型调用失败，请稍后重试",
                    "retryable": True,
                },
            )
        finally:
            if graph_stream is not None:
                await graph_stream.aclose()
            self._active_sessions.discard(session_id)

    async def _mark_failed(self, message_id: str, error_code: str) -> None:
        async with self.session_factory() as session:
            await ChatRepository(session).fail_answer(message_id, error_code)

    @staticmethod
    def _message_text(content: Any) -> str:
        if isinstance(content, str):
            return content
        if isinstance(content, Sequence):
            parts: list[str] = []
            for block in content:
                if isinstance(block, str):
                    parts.append(block)
                elif isinstance(block, dict) and block.get("type") == "text":
                    parts.append(str(block.get("text", "")))
            return "".join(parts)
        return ""

    @staticmethod
    def _validate_citations(answer: str, sources: list[dict[str, Any]]) -> None:
        if not answer:
            raise AppError("EMPTY_MODEL_ANSWER", "模型没有返回有效回答", 502, retryable=True)
        allowed = {str(source["label"]) for source in sources}
        used = set(re.findall(r"\[(S\d+)]", answer))
        if sources and not used:
            raise AppError("MISSING_CITATION", "模型回答缺少知识库引用", 502, retryable=True)
        if not used.issubset(allowed):
            raise AppError("INVALID_CITATION", "模型回答包含无效引用", 502, retryable=True)
