from __future__ import annotations

"""YAML-driven agent and orchestrator configuration."""

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml

from .agents import Agent, LoopStrategy
from .core import AgentError, AgentOutput, AgentTrait, LlmProvider
from .llm import (
    AnthropicProvider,
    AzureOpenAiProvider,
    GeminiProvider,
    OllamaProvider,
    OpenAiProvider,
)
from .mcp import WebMcpOptions
from .memory import SlidingWindowMemory, TokenAwareMemory, VectorMemory
from .multi import MultiAgentOrchestrator, RoutingStrategy
from .tools import (
    CalculatorTool,
    FileReadTool,
    FileWriteTool,
    ToolRegistry,
    WebFetchTool,
    WebSearchTool,
)


@dataclass(slots=True)
class ConfiguredRuntime:
    """Runtime loaded from YAML."""

    agent: AgentTrait | None = None
    multi: MultiAgentOrchestrator | None = None

    async def run(self, input_text: str) -> AgentOutput:
        if self.agent is not None:
            return await self.agent.run(input_text)
        if self.multi is not None:
            return await self.multi.run(input_text)
        raise AgentError.invalid_configuration("runtime is empty")


async def load_runtime_from_yaml(path: str | Path) -> ConfiguredRuntime:
    return await load_runtime_from_yaml_str(Path(path).read_text(encoding="utf-8"))


async def load_runtime_from_yaml_str(yaml_text: str) -> ConfiguredRuntime:
    data = yaml.safe_load(yaml_text) or {}
    config = RuntimeConfig.from_dict(data)
    return await config.build()


async def load_agent_from_yaml(path: str | Path) -> AgentTrait:
    return await load_agent_from_yaml_str(Path(path).read_text(encoding="utf-8"))


async def load_agent_from_yaml_str(yaml_text: str) -> AgentTrait:
    data = yaml.safe_load(yaml_text) or {}
    return await AgentYamlConfig.from_dict(data).build_boxed()


async def load_multi_agent_from_yaml(path: str | Path) -> MultiAgentOrchestrator:
    return await load_multi_agent_from_yaml_str(Path(path).read_text(encoding="utf-8"))


async def load_multi_agent_from_yaml_str(yaml_text: str) -> MultiAgentOrchestrator:
    data = yaml.safe_load(yaml_text) or {}
    return await MultiAgentYamlConfig.from_dict(data).build()


@dataclass(slots=True)
class RuntimeConfig:
    """Top-level YAML runtime config."""

    kind: str
    agent: AgentYamlConfig | None = None
    multi_agent: MultiAgentYamlConfig | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "RuntimeConfig":
        kind = data.get("kind")
        if kind == "agent":
            return cls(kind=kind, agent=AgentYamlConfig.from_dict(data))
        if kind == "multi_agent":
            return cls(kind=kind, multi_agent=MultiAgentYamlConfig.from_dict(data))
        raise AgentError.invalid_configuration(
            "runtime kind must be 'agent' or 'multi_agent'"
        )

    async def build(self) -> ConfiguredRuntime:
        if self.agent is not None:
            return ConfiguredRuntime(agent=await self.agent.build_boxed())
        if self.multi_agent is not None:
            return ConfiguredRuntime(multi=await self.multi_agent.build())
        raise AgentError.invalid_configuration("runtime kind is missing a payload")


@dataclass(slots=True)
class MemoryYamlConfig:
    """YAML config for memory backends."""

    type: str = "in_memory"
    window_size: int | None = None
    max_tokens: int | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> "MemoryYamlConfig":
        if not data:
            return cls()
        return cls(
            type=str(data.get("type", "in_memory")),
            window_size=data.get("window_size"),
            max_tokens=data.get("max_tokens"),
        )

    def is_default(self) -> bool:
        return self.type == "in_memory"


@dataclass(slots=True)
class LoopStrategyYamlConfig:
    """YAML config for loop strategies."""

    type: str = "re_act"
    max_steps: int | None = None
    instruction: str | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> "LoopStrategyYamlConfig":
        if not data:
            return cls()
        return cls(
            type=str(data.get("type", "re_act")),
            max_steps=data.get("max_steps"),
            instruction=data.get("instruction"),
        )

    def is_default(self) -> bool:
        return self.type == "re_act" and self.max_steps is None

    def into_loop_strategy(self) -> LoopStrategy:
        if self.type == "re_act":
            return LoopStrategy.react(self.max_steps or 8)
        if self.type == "chain_of_thought":
            return LoopStrategy.cot()
        if self.type == "plan_and_execute":
            return LoopStrategy.plan_and_execute(self.max_steps or 8)
        if self.type == "custom":
            return LoopStrategy.custom(self.instruction or "")
        raise AgentError.invalid_configuration(f"unknown loop strategy: {self.type}")


@dataclass(slots=True)
class RoutingYamlConfig:
    """YAML config for routing strategies."""

    type: str = "sequential"
    order: list[str] | None = None
    agents: list[str] | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> "RoutingYamlConfig":
        if not data:
            return cls()
        return cls(
            type=str(data.get("type", "sequential")),
            order=data.get("order"),
            agents=data.get("agents"),
        )

    def is_default(self) -> bool:
        return self.type == "sequential" and self.order is None

    def into_routing_strategy(self, default_order: list[str]) -> RoutingStrategy:
        if self.type == "sequential":
            return RoutingStrategy.sequential(self.order or default_order)
        if self.type == "parallel":
            return RoutingStrategy.parallel(self.agents or [])
        raise AgentError.invalid_configuration(
            f"unsupported routing strategy in YAML: {self.type}"
        )


@dataclass(slots=True)
class LlmYamlConfig:
    """YAML config for provider creation."""

    provider: str
    api_key: str | None = None
    base_url: str | None = None
    model: str | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "LlmYamlConfig":
        return cls(
            provider=str(data.get("provider")),
            api_key=data.get("api_key"),
            base_url=data.get("base_url"),
            model=data.get("model"),
        )


@dataclass(slots=True)
class ToolYamlConfig:
    """YAML config for tools."""

    type: str
    root: str | None = None
    target: str | None = None
    api_key: str | None = None
    api_key_header: str | None = None
    api_key_prefix: str | None = None
    headers: dict[str, str] | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ToolYamlConfig":
        return cls(
            type=str(data.get("type")),
            root=data.get("root"),
            target=data.get("target"),
            api_key=data.get("api_key"),
            api_key_header=data.get("api_key_header"),
            api_key_prefix=data.get("api_key_prefix"),
            headers=data.get("headers"),
        )


@dataclass(slots=True)
class AgentYamlConfig:
    """YAML config for a single agent."""

    name: str | None
    llm: LlmYamlConfig
    system: str | None = None
    memory: MemoryYamlConfig = field(default_factory=MemoryYamlConfig)
    tools: list[ToolYamlConfig] = field(default_factory=list)
    loop_strategy: LoopStrategyYamlConfig = field(
        default_factory=LoopStrategyYamlConfig
    )
    model: str | None = None
    temperature: float | None = None
    max_tokens: int | None = None
    max_steps: int | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "AgentYamlConfig":
        return cls(
            name=data.get("name"),
            llm=LlmYamlConfig.from_dict(data["llm"]),
            system=data.get("system"),
            memory=MemoryYamlConfig.from_dict(data.get("memory")),
            tools=[ToolYamlConfig.from_dict(item) for item in data.get("tools", [])],
            loop_strategy=LoopStrategyYamlConfig.from_dict(data.get("loop_strategy")),
            model=data.get("model"),
            temperature=data.get("temperature"),
            max_tokens=data.get("max_tokens"),
            max_steps=data.get("max_steps"),
        )

    async def build_boxed(self) -> AgentTrait:
        llm = _build_llm(self.llm)
        tools = await _build_tools(self.tools)
        loop_strategy = self.loop_strategy.into_loop_strategy()

        builder = Agent.builder().llm_arc(llm).tools(tools)
        if self.memory.type == "sliding_window":
            builder = builder.memory(
                SlidingWindowMemory.new(self.memory.window_size or 0)
            )
        elif self.memory.type == "token_aware":
            builder = builder.memory(TokenAwareMemory.new(self.memory.max_tokens or 0))
        elif self.memory.type == "vector":
            builder = builder.memory(VectorMemory.new())

        if self.system is not None:
            builder = builder.system(self.system)
        if self.model is not None:
            builder = builder.model(self.model)
        if self.temperature is not None:
            builder = builder.temperature(self.temperature)
        if self.max_tokens is not None:
            builder = builder.max_tokens(self.max_tokens)
        if self.max_steps is not None:
            builder = builder.max_steps(self.max_steps)
        builder = builder.loop_strategy(loop_strategy)
        return builder.build()


@dataclass(slots=True)
class NamedAgentYamlConfig:
    """Named agent entry inside a multi-agent YAML config."""

    name: str
    agent: AgentYamlConfig

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "NamedAgentYamlConfig":
        name = str(data.get("name"))
        agent = AgentYamlConfig.from_dict(data)
        return cls(name=name, agent=agent)


@dataclass(slots=True)
class MultiAgentYamlConfig:
    """YAML config for multi-agent orchestration."""

    agents: list[NamedAgentYamlConfig]
    routing: RoutingYamlConfig = field(default_factory=RoutingYamlConfig)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "MultiAgentYamlConfig":
        return cls(
            agents=[
                NamedAgentYamlConfig.from_dict(item) for item in data.get("agents", [])
            ],
            routing=RoutingYamlConfig.from_dict(data.get("routing")),
        )

    async def build(self) -> MultiAgentOrchestrator:
        builder = MultiAgentOrchestrator.builder()
        default_order = [agent.name for agent in self.agents]
        for agent in self.agents:
            builder = builder.add_agent_boxed(
                agent.name, await agent.agent.build_boxed()
            )
        builder = builder.routing(self.routing.into_routing_strategy(default_order))
        return builder.build()


def _build_llm(config: LlmYamlConfig) -> LlmProvider:
    if config.provider == "open_ai":
        builder = OpenAiProvider.from_env()
        if config.api_key is not None:
            builder = builder.api_key(config.api_key)
        if config.base_url is not None:
            builder = builder.base_url(config.base_url)
        if config.model is not None:
            builder = builder.model(config.model)
        return builder.build()
    if config.provider == "azure_open_ai":
        builder = AzureOpenAiProvider.from_env()
        if config.api_key is not None:
            builder = builder.api_key(config.api_key)
        if config.base_url is not None:
            builder = builder.base_url(config.base_url)
        if config.model is not None:
            builder = builder.model(config.model)
        return builder.build()
    if config.provider == "anthropic":
        builder = AnthropicProvider.from_env()
        if config.api_key is not None:
            builder = builder.api_key(config.api_key)
        if config.base_url is not None:
            builder = builder.base_url(config.base_url)
        if config.model is not None:
            builder = builder.model(config.model)
        return builder.build()
    if config.provider == "gemini":
        builder = GeminiProvider.from_env()
        if config.api_key is not None:
            builder = builder.api_key(config.api_key)
        if config.base_url is not None:
            builder = builder.base_url(config.base_url)
        if config.model is not None:
            builder = builder.model(config.model)
        return builder.build()
    if config.provider == "ollama":
        builder = OllamaProvider.builder()
        if config.base_url is not None:
            builder = builder.base_url(config.base_url)
        if config.model is not None:
            builder = builder.model(config.model)
        return builder.build()
    raise AgentError.invalid_configuration(
        "requested LLM provider feature is not enabled in this build"
    )


async def _build_tools(configs: list[ToolYamlConfig]) -> ToolRegistry:
    registry = ToolRegistry.new()
    for config in configs:
        if config.type == "calculator":
            registry = registry.register(CalculatorTool())
        elif config.type == "web_fetch":
            registry = registry.register(WebFetchTool())
        elif config.type == "web_search":
            registry = registry.register(WebSearchTool())
        elif config.type == "file_read":
            registry = registry.register(FileReadTool(config.root))
        elif config.type == "file_write":
            registry = registry.register(FileWriteTool(config.root))
        elif config.type == "mcp":
            target = config.target or ""
            if target.startswith(("http://", "https://")):
                options = WebMcpOptions.new()
                if config.headers:
                    options = options.headers(config.headers)
                if config.api_key is not None:
                    options = options.api_key(config.api_key)
                if config.api_key_header is not None:
                    options = options.api_key_header(config.api_key_header)
                if config.api_key_prefix is not None:
                    options = options.api_key_prefix(config.api_key_prefix)
                registry = await registry.register_mcp_http_with_options(
                    target, options
                )
            else:
                registry = await registry.register_mcp(target)
        else:
            raise AgentError.invalid_configuration(f"unknown tool type: {config.type}")
    return registry
