"""Agent runtime orchestration — ties config, providers, and runtime together.

The AgentEngine is the top-level entry point for running an agent. It:
1. Loads and validates the agent config
2. Creates the model provider
3. Starts all MCP servers
4. Collects and filters tools
5. Runs the tool call loop
6. Shuts down everything cleanly
"""

from __future__ import annotations

from pathlib import Path

from hatchdx import HdxError
from hatchdx.agent.config import AgentConfig, load_agent_config
from hatchdx.agent.providers.base import ModelProvider, ProviderAuthError
from hatchdx.agent.runtime.loop import LoopResult, run_agent_loop
from hatchdx.agent.runtime.servers import ManagedServer, ServerConnectionError, ServerManager
from hatchdx.agent.runtime.tools import CollectedTools, collect_and_filter_tools


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------


class AgentError(HdxError):
    """Base class for agent engine errors."""


# ---------------------------------------------------------------------------
# Provider factory
# ---------------------------------------------------------------------------


def create_provider(config: AgentConfig) -> ModelProvider:
    """Create a model provider from agent config.

    Supports:
    - anthropic: Anthropic SDK (requires ANTHROPIC_API_KEY)
    - openai: OpenAI SDK (requires OPENAI_API_KEY)
    - local: Ollama / LM Studio via OpenAI-compatible API (no key needed)

    Raises:
        AgentError: If the provider is unknown.
    """
    provider = config.model.provider

    if provider == "anthropic":
        from hatchdx.agent.providers.anthropic import AnthropicProvider

        return AnthropicProvider(model=config.model.model)

    elif provider == "openai":
        from hatchdx.agent.providers.openai import OpenAIProvider

        return OpenAIProvider(
            model=config.model.model,
            base_url=config.model.base_url,
        )

    elif provider == "local":
        from hatchdx.agent.providers.local import LocalProvider

        return LocalProvider(
            model=config.model.model,
            base_url=config.model.base_url,
        )

    else:
        raise AgentError(
            f"Unknown provider: '{provider}'\n"
            "Supported providers: anthropic, openai, local"
        )


# ---------------------------------------------------------------------------
# Agent Engine
# ---------------------------------------------------------------------------


def estimate_cost_for_provider(
    provider: str, model: str, input_tokens: int, output_tokens: int
) -> float:
    """Estimate cost in USD using the correct provider pricing.

    Args:
        provider: Provider name (anthropic, openai, local).
        model: Model identifier string.
        input_tokens: Number of input tokens.
        output_tokens: Number of output tokens.

    Returns:
        Estimated cost in USD.
    """
    if provider == "anthropic":
        from hatchdx.agent.providers.anthropic import estimate_cost

        return estimate_cost(model, input_tokens, output_tokens)
    elif provider == "openai":
        from hatchdx.agent.providers.openai import estimate_cost

        return estimate_cost(model, input_tokens, output_tokens)
    elif provider == "local":
        return 0.0
    else:
        # Fallback to Anthropic pricing for unknown providers
        from hatchdx.agent.providers.anthropic import estimate_cost

        return estimate_cost(model, input_tokens, output_tokens)


class AgentEngine:
    """Top-level agent runtime. Manages the full lifecycle of an agent session.

    Usage::

        engine = AgentEngine.from_config_path(Path("agents/my-agent/agent.yaml"))
        async with engine:
            result = await engine.chat("What's the weather?")
    """

    def __init__(
        self,
        config: AgentConfig,
        config_dir: Path | None = None,
    ) -> None:
        self.config = config
        self.config_dir = config_dir or Path.cwd()
        self._provider: ModelProvider | None = None
        self._server_manager: ServerManager | None = None
        self._collected_tools: CollectedTools | None = None
        self._servers: list[ManagedServer] = []
        self._started = False

    @classmethod
    def from_config_path(cls, config_path: Path) -> AgentEngine:
        """Create an AgentEngine from an agent.yaml file path."""
        config = load_agent_config(config_path)
        return cls(config, config_dir=config_path.parent)

    # -- Context manager ------------------------------------------------------

    async def __aenter__(self) -> AgentEngine:
        await self.start()
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.shutdown()

    # -- Lifecycle ------------------------------------------------------------

    async def start(self) -> list[ManagedServer]:
        """Start the agent: create provider, start servers, collect tools.

        Returns:
            List of started ManagedServer instances.

        Raises:
            ProviderAuthError: If API key is missing.
            ServerConnectionError: If a server fails to start.
        """
        # Create provider (validates API key)
        self._provider = create_provider(self.config)

        # Start servers
        self._server_manager = ServerManager()
        await self._server_manager.__aenter__()

        self._servers = await self._server_manager.start_all(
            self.config.servers,
            base_dir=self.config_dir,
        )

        # Discover and collect tools
        server_tools = await self._server_manager.discover_all_tools()
        self._collected_tools = collect_and_filter_tools(server_tools, self.config)

        self._started = True
        return self._servers

    async def shutdown(self) -> None:
        """Shut down all servers and clean up resources."""
        if self._server_manager is not None:
            await self._server_manager.__aexit__(None, None, None)
            self._server_manager = None
        self._started = False

    # -- Properties -----------------------------------------------------------

    @property
    def provider(self) -> ModelProvider:
        if self._provider is None:
            raise AgentError("Engine not started. Call start() first.")
        return self._provider

    @property
    def server_manager(self) -> ServerManager:
        if self._server_manager is None:
            raise AgentError("Engine not started. Call start() first.")
        return self._server_manager

    @property
    def collected_tools(self) -> CollectedTools:
        if self._collected_tools is None:
            raise AgentError("Engine not started. Call start() first.")
        return self._collected_tools

    @property
    def servers(self) -> list[ManagedServer]:
        return list(self._servers)

    # -- Chat -----------------------------------------------------------------

    async def chat(
        self,
        message: str,
        on_turn: object | None = None,
    ) -> LoopResult:
        """Send a message to the agent and get the full loop result.

        Args:
            message: User message.
            on_turn: Optional async callback(LoopTurn) for streaming progress.

        Returns:
            LoopResult with all turns, final text, and stats.
        """
        if not self._started:
            raise AgentError("Engine not started. Call start() first.")

        return await run_agent_loop(
            provider=self.provider,
            server_manager=self.server_manager,
            collected_tools=self.collected_tools,
            system_prompt=self.config.system_prompt,
            user_message=message,
            max_tokens=self.config.model.max_tokens,
            temperature=self.config.model.temperature,
            max_tool_calls=self.config.settings.max_tool_calls,
            tool_call_timeout=self.config.settings.tool_call_timeout,
            retry_on_error=self.config.settings.retry_on_error,
            max_retries=self.config.settings.max_retries,
            on_turn=on_turn,
        )
