"""Thryve – main entry class that wires all subsystems together."""

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator
from typing import Any, Callable

from thryve.config import ThryveConfig
from thryve.core.agent.agent import Agent
from thryve.core.agent.models import LoopConfig, TurnResult
from thryve.core.agent.multi_agent import MultiAgentOrchestrator
from thryve.core.backends.adapter import BackendAdapter
from thryve.core.context.manager import ContextManager
from thryve.core.context.models import BackendInfo, Message
from thryve.core.graph.executor import GraphExecutor
from thryve.core.graph.graph import Graph
from thryve.core.memory.manager import MemoryManager
from thryve.core.memory.types import MemoryChunk
from thryve.core.tools.models import Tool
from thryve.core.tools.registry import ToolRegistry
from thryve.embedding import EmbeddingModel, create_embedding
from thryve.llm import create_llm
from thryve.memory import create_memory
from thryve.utils import get_logger

logger = get_logger("thryve")


class Thryve:
    """Top-level facade for the Thryve framework.

    Initializes LLM, embedding, memory, context, agent, and graph subsystems
    in the correct dependency order and exposes a unified public API.

    All public methods default to **synchronous** (blocking) for ease of use
    in REPLs and scripts.  Async variants are available with an ``_async``
    suffix.

    Usage::

        thryve = Thryve(ThryveConfig.from_env())
        reply = thryve.chat("Hello!")
    """

    def __init__(self, config: ThryveConfig | None = None) -> None:
        self.config = config or ThryveConfig()

        # Strict initialization order (per DESIGN.md §5.2):
        # 1. LLM (no deps)
        self.llm: BackendAdapter = self._init_llm()
        # 2. Embedding (depends on backend config)
        self.embedding: EmbeddingModel = self._init_embedding()
        # 3. Memory (depends on embedding)
        self.memory: MemoryManager = self._init_memory()
        # 4. Context (depends on memory – for archive strategy)
        self.context: ContextManager = self._init_context()
        # 5. Tools
        self._tool_registry = ToolRegistry()
        # 6. Agent orchestrator
        self.agent_orchestrator = MultiAgentOrchestrator()
        # 7. Graph executor
        self.graph_executor = GraphExecutor()

        logger.info("Thryve initialised (backend=%s)", self.llm.get_info().provider)

    # ------------------------------------------------------------------
    # Subsystem initialisation
    # ------------------------------------------------------------------

    def _init_llm(self) -> BackendAdapter:
        return create_llm(self.config.llm)

    def _init_embedding(self) -> EmbeddingModel:
        try:
            return create_embedding(self.config.embedding)
        except Exception:
            logger.warning("Embedding model could not be initialised; memory will use keyword-only search")
            return None  # type: ignore[return-value]

    def _init_memory(self) -> MemoryManager:
        emb_fn = self.embedding.embed if self.embedding else None
        return create_memory(self.config.memory, embedding_fn=emb_fn)

    def _init_context(self) -> ContextManager:
        from thryve.core.context.compression import get_strategy

        strategy = get_strategy(
            self.config.context.compression_strategy,
            llm=self.llm,
        ) if self.config.context.compression_strategy == "summarize" else get_strategy("truncate")

        return ContextManager(
            max_tokens=self.config.context.max_tokens,
            strategy=strategy,
        )

    # ==================================================================
    # Chat interface
    # ==================================================================

    # ---- async (core implementation) ---------------------------------

    async def chat_async(self, message: str, session_id: str | None = None) -> str:
        """Async LLM chat with optional memory retrieval.

        1. Search memory for relevant context.
        2. Build the message list.
        3. Call the LLM.
        4. Store the exchange in short-term memory.
        """
        memories: list[MemoryChunk] = []
        try:
            memories = self.memory.search(message, top_k=5)
        except Exception:
            pass

        messages = self._build_messages(message, memories)
        response = await self.llm.generate_async(messages)

        try:
            self.memory.append_to_short_term(f"User: {message}")
            self.memory.append_to_short_term(f"Assistant: {response.content}")
        except Exception:
            pass

        return response.content

    async def chat_stream_async(
        self, message: str, session_id: str | None = None,
    ) -> AsyncIterator[str]:
        """Async streaming chat – yields text chunks as they arrive.

        Memory is retrieved before the call and the full reply is
        persisted to short-term memory after the stream ends.
        """
        memories: list[MemoryChunk] = []
        try:
            memories = self.memory.search(message, top_k=5)
        except Exception:
            pass
        messages = self._build_messages(message, memories)
        full_content: list[str] = []
        async for chunk in self.llm.generate_stream(messages):
            full_content.append(chunk)
            yield chunk
        full_text = "".join(full_content)
        try:
            self.memory.append_to_short_term(f"User: {message}")
            self.memory.append_to_short_term(f"Assistant: {full_text}")
        except Exception:
            pass

    async def chat_with_agent_async(
        self,
        message: str,
        agent_name: str = "default",
        **kwargs: Any,
    ) -> TurnResult:
        """Async agent chat – the agent can call registered tools."""
        agent = self._get_or_create_agent(agent_name)
        return await agent.run(message, context=self.context, **kwargs)

    async def execute_graph_async(
        self,
        graph: Graph,
        inputs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Async DAG workflow execution."""
        return await self.graph_executor.execute(graph, inputs)

    # ---- sync (default, no suffix) -----------------------------------

    def chat(self, message: str, session_id: str | None = None) -> str:
        """Send a message and get a reply (synchronous, blocking).

        This is the default entry point.  For async usage, see
        :meth:`chat_async`.
        """
        return asyncio.run(self.chat_async(message, session_id=session_id))

    def chat_stream(
        self,
        message: str,
        session_id: str | None = None,
        callback: Callable[[str], None] | None = None,
    ) -> str:
        """Streaming chat (synchronous).

        Calls *callback* for each text chunk as it arrives, then returns the
        full reply as a string.  For async usage, see :meth:`chat_stream_async`.

        Example::

            reply = thryve.chat_stream(
                "什么是大语言模型",
                callback=lambda c: print(c, end="", flush=True),
            )
        """
        full: list[str] = []

        async def _consume() -> str:
            async for chunk in self.chat_stream_async(message, session_id=session_id):
                full.append(chunk)
                if callback:
                    callback(chunk)
            return "".join(full)

        return asyncio.run(_consume())

    def chat_with_agent(
        self,
        message: str,
        agent_name: str = "default",
        **kwargs: Any,
    ) -> TurnResult:
        """Run an agent that can use registered tools (synchronous).

        For async usage, see :meth:`chat_with_agent_async`.
        """
        return asyncio.run(self.chat_with_agent_async(message, agent_name=agent_name, **kwargs))

    def execute_graph(
        self,
        graph: Graph,
        inputs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Execute a DAG workflow (synchronous).

        For async usage, see :meth:`execute_graph_async`.
        """
        return asyncio.run(self.execute_graph_async(graph, inputs))

    # ==================================================================
    # Memory interface
    # ==================================================================

    def add_to_memory(self, content: str, permanent: bool = False) -> None:
        if permanent:
            self.memory.promote_to_long_term(content)
        else:
            self.memory.append_to_short_term(content)

    def search_memory(self, query: str, top_k: int = 5) -> list[MemoryChunk]:
        return self.memory.search(query, top_k=top_k)

    # ==================================================================
    # Tool interface
    # ==================================================================

    def register_tool(self, tool: Tool) -> None:
        self._tool_registry.register(tool)

    def list_tools(self) -> list[Tool]:
        return self._tool_registry.list_tools()

    # ==================================================================
    # Agent interface
    # ==================================================================

    def register_agent(self, agent: Agent) -> None:
        self.agent_orchestrator.register_agent(agent)

    # ==================================================================
    # Introspection
    # ==================================================================

    def get_llm_info(self) -> BackendInfo:
        return self.llm.get_info()

    # ==================================================================
    # Internals
    # ==================================================================

    def _build_messages(
        self, user_message: str, memories: list[MemoryChunk],
    ) -> list[Message]:
        """Assemble a message list incorporating memory context."""
        messages: list[Message] = []

        if memories:
            ctx_text = "\n".join(f"- {m.text}" for m in memories)
            messages.append(
                Message(
                    role="system",
                    content=f"Relevant context from memory:\n{ctx_text}",
                )
            )

        messages.append(Message(role="user", content=user_message))
        return messages

    def _get_or_create_agent(self, name: str) -> Agent:
        agent = self.agent_orchestrator.labor_market.get(name)
        if agent is not None:
            return agent

        agent = Agent(
            name=name,
            llm=self.llm,
            tools=self._tool_registry.list_tools(),
            config=LoopConfig(
                max_turns=self.config.agent.max_turns,
                max_steps_per_turn=self.config.agent.max_steps_per_turn,
                loop_detection_window=self.config.agent.loop_detection_window,
                enable_auto_continue=self.config.agent.enable_auto_continue,
            ),
        )
        self.agent_orchestrator.register_agent(agent)
        return agent
