from collections.abc import AsyncIterator, Awaitable, Callable

from sqlalchemy.orm import Session

from agent.deepseek_client import DeepSeekClient, should_use_web_search
import config
from memory.session_store import SessionStore


class ChatAgent:
    def __init__(self, db: Session, deepseek: DeepSeekClient | None = None):
        self.db = db
        self.store = SessionStore(db)
        self.deepseek = deepseek or DeepSeekClient()

    def _system_prompt(self) -> str:
        prompt = config.settings.system_prompt
        if config.settings.web_search_enabled:
            prompt += (
                "\n\n【联网已开启】你可以使用 web_search 工具搜索互联网。"
                "遇到新闻、天气、价格、实时事件、最新动态等问题时，"
                "请先搜索再回答，并注明信息来源。"
            )
        return prompt

    def _build_messages(
        self, session_id: str, user_input: str, include_unsaved: bool = True
    ) -> list[dict[str, str]]:
        messages: list[dict[str, str]] = [
            {"role": "system", "content": self._system_prompt()}
        ]
        history = self.store.get_recent_messages(
            session_id, config.settings.max_history_turns
        )
        for msg in history:
            messages.append({"role": msg.role, "content": msg.content})
        if include_unsaved:
            messages.append({"role": "user", "content": user_input})
        return messages

    async def stream_reply(
        self,
        session_id: str,
        user_input: str,
        model: str | None = None,
        on_search: Callable[[str], Awaitable[None]] | None = None,
    ) -> AsyncIterator[str]:
        session = self.store.get_session(session_id)
        if not session:
            raise LookupError(f"会话不存在: {session_id}")

        messages = self._build_messages(session_id, user_input)
        self.store.add_message(session_id, "user", user_input)
        use_tools = should_use_web_search(user_input)

        full_content = ""
        async for token in self.deepseek.stream_chat(
            messages,
            model=model,
            use_tools=use_tools,
            on_search=on_search,
        ):
            if token.startswith("\1SEARCH:"):
                continue
            full_content += token
            yield token

        assistant_msg = self.store.add_message(session_id, "assistant", full_content)
        yield f"\0{assistant_msg.id}"

    async def reply(
        self,
        session_id: str,
        user_input: str,
        model: str | None = None,
    ) -> tuple[str, str, str]:
        session = self.store.get_session(session_id)
        if not session:
            raise LookupError(f"会话不存在: {session_id}")

        messages = self._build_messages(session_id, user_input)
        user_msg = self.store.add_message(session_id, "user", user_input)
        content = await self.deepseek.chat(
            messages, model=model, use_tools=should_use_web_search(user_input)
        )
        assistant_msg = self.store.add_message(session_id, "assistant", content)
        return user_msg.id, assistant_msg.id, content