"""Configuration models and validation for mcpbridge.

This module defines `BridgeConfig`, including:
- Provider model alias resolution
- Environment-variable API key fallbacks
- MCP transport/server configuration validation
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from typing import Any, Callable, Literal

from mcpbridge.exceptions import APIKeyMissingError, ConfigValidationError

logger = logging.getLogger("mcpbridge")

ProviderName = Literal[
    "anthropic",
    "openai",
    "gemini",
    "mistral",
    "cohere",
    "groq",
    "ollama",
    "together",
    "bedrock",
    "azure_openai",
    "openai_compatible",
]


MODEL_ALIASES: dict[str, dict[str, str]] = {
    "anthropic": {
        "opus": "claude-opus-4-20250514",
        "sonnet": "claude-sonnet-4-20250514",
        "haiku": "claude-haiku-4-5-20251001",
    },
    "openai": {
        "gpt4o": "gpt-4o",
        "gpt4o-mini": "gpt-4o-mini",
        "o1": "o1",
        "o3": "o3",
        "o3-mini": "o3-mini",
        "o4-mini": "o4-mini",
    },
    "gemini": {
        "pro": "gemini-1.5-pro",
        "flash": "gemini-1.5-flash",
        "flash2": "gemini-2.0-flash",
        "pro2": "gemini-2.0-pro",
    },
    "mistral": {
        "large": "mistral-large-latest",
        "small": "mistral-small-latest",
        "nemo": "open-mistral-nemo",
        "mixtral": "open-mixtral-8x22b",
    },
    "groq": {
        "llama3": "llama-3.3-70b-versatile",
        "llama3-small": "llama-3.1-8b-instant",
        "mixtral": "mixtral-8x7b-32768",
    },
    "cohere": {
        "command": "command-r-plus",
        "command-r": "command-r",
    },
    "together": {
        "llama3": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
        "mixtral": "mistralai/Mixtral-8x7B-Instruct-v0.1",
    },
}

ENV_KEY_MAP: dict[str, str | None] = {
    "anthropic": "ANTHROPIC_API_KEY",
    "openai": "OPENAI_API_KEY",
    "gemini": "GOOGLE_API_KEY",
    "mistral": "MISTRAL_API_KEY",
    "cohere": "COHERE_API_KEY",
    "groq": "GROQ_API_KEY",
    "together": "TOGETHER_API_KEY",
    "azure_openai": "AZURE_OPENAI_API_KEY",
    "bedrock": None,  # uses boto3 AWS credential chain
    "ollama": None,  # local, no key
    "openai_compatible": "OPENAI_API_KEY",
}


def _warn_unknown_keys(data: dict[str, Any], allowed: set[str], path: str) -> None:
    unknown = sorted(set(data.keys()) - allowed)
    if unknown:
        logger.warning("Unknown config keys at %s: %s", path, ", ".join(unknown))


def _resolve_model(provider: str, model: str | None) -> str:
    if not model:
        # Leave model resolution to adapters/provider defaults.
        return ""
    alias_map = MODEL_ALIASES.get(provider, {})
    return alias_map.get(model, model)


@dataclass(slots=True)
class ThinkingConfig:
    enabled: bool = False
    budget_tokens: int = 1024


@dataclass(slots=True)
class LLMConfig:
    provider: ProviderName
    model: str = ""
    api_key: str | None = None
    base_url: str | None = None

    temperature: float = 0.7
    max_tokens: int = 4096
    top_p: float | None = None
    top_k: int | None = None
    stop_sequences: list[str] = field(default_factory=list)
    stream: bool = False
    timeout: int = 60

    extra_params: dict[str, Any] = field(default_factory=dict)

    thinking: ThinkingConfig = field(default_factory=ThinkingConfig)

    azure_deployment: str = ""
    azure_api_version: str = "2024-02-01"

    aws_region: str = "us-east-1"
    aws_profile: str = ""


@dataclass(slots=True)
class MCPServerConfig:
    name: str
    url: str | None = None
    transport: str = "auto"
    headers: dict[str, str] = field(default_factory=dict)
    command: str | None = None
    args: list[str] = field(default_factory=list)
    env: dict[str, str] = field(default_factory=dict)


@dataclass(slots=True)
class MCPConfig:
    url: str | None = None
    transport: str = "auto"  # auto|http|sse|streamable_http|ws|stdio
    headers: dict[str, str] = field(default_factory=dict)
    timeout: int = 30

    command: str | None = None
    args: list[str] = field(default_factory=list)
    env: dict[str, str] = field(default_factory=dict)

    servers: list[MCPServerConfig] | None = None
    namespace_strategy: Literal["prefix", "error", "last_wins"] = "prefix"


@dataclass(slots=True)
class PromptConfig:
    system: str = ""
    interpolate: bool = True
    context_vars: dict[str, Any] = field(default_factory=dict)
    inject_tool_descriptions: bool = True
    inject_internal_instructions: bool = True
    user_suffix: str = ""
    user_prefix: str = ""


LoopCallable = Callable[..., Any]


@dataclass(slots=True)
class LoopConfig:
    max_iterations: int = 10
    max_tokens_total: int = 32000
    tool_timeout: int = 30
    parallel_tool_calls: bool = True

    on_tool_call: LoopCallable | None = None
    on_tool_result: LoopCallable | None = None
    on_iteration: LoopCallable | None = None

    retry_on_tool_error: bool = True
    error_strategy: Literal["raise", "return_error_to_llm", "skip"] = "return_error_to_llm"


@dataclass(slots=True)
class SessionConfig:
    persist_history: bool = True
    max_history_tokens: int = 32000
    history_trim_strategy: Literal["oldest_first", "summarize"] = "oldest_first"


@dataclass(slots=True)
class BridgeConfig:
    """Top-level validated configuration for the bridge."""

    llm: LLMConfig
    mcp: MCPConfig
    prompt: PromptConfig
    loop: LoopConfig
    session: SessionConfig

    @staticmethod
    def from_dict(config: dict[str, Any]) -> "BridgeConfig":
        """Validate and construct a `BridgeConfig` from an untrusted dict."""

        if not isinstance(config, dict):
            raise ConfigValidationError("Config must be a dict")

        allowed_top = {"llm", "mcp", "prompt", "loop", "session"}
        _warn_unknown_keys(config, allowed_top, path="root")

        llm_raw = config.get("llm")
        mcp_raw = config.get("mcp")
        prompt_raw = config.get("prompt", {})
        loop_raw = config.get("loop", {})
        session_raw = config.get("session", {})

        if not isinstance(llm_raw, dict):
            raise ConfigValidationError("Missing or invalid `llm` config (must be an object)")
        if not isinstance(mcp_raw, dict):
            raise ConfigValidationError("Missing or invalid `mcp` config (must be an object)")

        llm_allowed = {
            "provider",
            "model",
            "api_key",
            "base_url",
            "temperature",
            "max_tokens",
            "top_p",
            "top_k",
            "stop_sequences",
            "stream",
            "timeout",
            "extra_params",
            "thinking",
            "azure_deployment",
            "azure_api_version",
            "aws_region",
            "aws_profile",
        }
        _warn_unknown_keys(llm_raw, llm_allowed, path="llm")

        provider = llm_raw.get("provider")
        if not provider or not isinstance(provider, str):
            raise ConfigValidationError("`llm.provider` is required and must be a string")
        if provider not in MODEL_ALIASES and provider not in {"bedrock", "ollama", "azure_openai", "openai_compatible"}:
            # Provider list in this package is fixed; adapters can add logic later.
            raise ConfigValidationError(f"Unsupported llm.provider: {provider}")

        model_in = llm_raw.get("model")
        model = _resolve_model(provider, model_in) if isinstance(model_in, str) else ""

        api_key_raw = llm_raw.get("api_key")
        api_key: str | None
        if api_key_raw is None:
            api_key = None
        elif isinstance(api_key_raw, str) and api_key_raw.strip() == "":
            api_key = None
        elif isinstance(api_key_raw, str):
            api_key = api_key_raw
        else:
            raise ConfigValidationError("`llm.api_key` must be a string or null")

        # Env fallback (except for bedrock/ollama which use credential chains / local)
        env_key = ENV_KEY_MAP.get(provider)
        if api_key is None and env_key is not None:
            api_key = os.environ.get(env_key)
            if api_key is not None and api_key.strip() == "":
                api_key = None
        if api_key is None and env_key is not None:
            raise APIKeyMissingError(f"Missing API key for provider `{provider}` (set `{env_key}` or `llm.api_key`).")

        thinking_raw = llm_raw.get("thinking", {}) or {}
        if not isinstance(thinking_raw, dict):
            raise ConfigValidationError("`llm.thinking` must be an object")
        thinking = ThinkingConfig(
            enabled=bool(thinking_raw.get("enabled", False)),
            budget_tokens=int(thinking_raw.get("budget_tokens", 1024)),
        )

        llm_cfg = LLMConfig(
            provider=provider,  # type: ignore[arg-type]
            model=model,
            api_key=api_key,
            base_url=llm_raw.get("base_url"),
            temperature=float(llm_raw.get("temperature", 0.7)),
            max_tokens=int(llm_raw.get("max_tokens", 4096)),
            top_p=llm_raw.get("top_p"),
            top_k=llm_raw.get("top_k"),
            stop_sequences=list(llm_raw.get("stop_sequences", []) or []),
            stream=bool(llm_raw.get("stream", False)),
            timeout=int(llm_raw.get("timeout", 60)),
            extra_params=dict(llm_raw.get("extra_params", {}) or {}),
            thinking=thinking,
            azure_deployment=str(llm_raw.get("azure_deployment", "")),
            azure_api_version=str(llm_raw.get("azure_api_version", "2024-02-01")),
            aws_region=str(llm_raw.get("aws_region", "us-east-1")),
            aws_profile=str(llm_raw.get("aws_profile", "")),
        )

        mcp_allowed = {
            "url",
            "transport",
            "headers",
            "timeout",
            "command",
            "args",
            "env",
            "servers",
            "namespace_strategy",
        }
        _warn_unknown_keys(mcp_raw, mcp_allowed, path="mcp")

        # Enforce `servers` vs top-level config behavior.
        servers_raw = mcp_raw.get("servers")
        namespace_strategy = mcp_raw.get("namespace_strategy", "prefix")
        if namespace_strategy not in {"prefix", "error", "last_wins"}:
            raise ConfigValidationError("`mcp.namespace_strategy` must be one of: prefix|error|last_wins")

        mcp_cfg: MCPConfig
        if servers_raw is not None:
            if not isinstance(servers_raw, list):
                raise ConfigValidationError("`mcp.servers` must be a list")
            if (mcp_raw.get("url") is not None) or (mcp_raw.get("command") is not None):
                logger.warning("`mcp.servers` provided; ignoring top-level `mcp.url`/`mcp.command`.")

            servers: list[MCPServerConfig] = []
            for i, srv in enumerate(servers_raw):
                if not isinstance(srv, dict):
                    raise ConfigValidationError(f"`mcp.servers[{i}]` must be an object")
                if "name" not in srv:
                    raise ConfigValidationError(f"`mcp.servers[{i}].name` is required")
                name = str(srv["name"])
                transport = str(srv.get("transport", "auto"))
                headers = dict(srv.get("headers", {}) or {})
                url = srv.get("url")
                command = srv.get("command")
                args = list(srv.get("args", []) or [])
                env = dict(srv.get("env", {}) or {})
                if url and command:
                    raise ConfigValidationError(f"`mcp.servers[{i}]` cannot set both `url` and `command`.")
                servers.append(
                    MCPServerConfig(
                        name=name,
                        url=str(url) if url is not None else None,
                        command=str(command) if command is not None else None,
                        args=args,
                        env=env,
                        headers=headers,
                        transport=transport,
                    )
                )

            mcp_cfg = MCPConfig(
                servers=servers,
                namespace_strategy=namespace_strategy,  # type: ignore[arg-type]
                transport=str(mcp_raw.get("transport", "auto")),
                headers=dict(mcp_raw.get("headers", {}) or {}),
                timeout=int(mcp_raw.get("timeout", 30)),
            )
        else:
            url = mcp_raw.get("url")
            command = mcp_raw.get("command")
            if url is not None and command is not None:
                raise ConfigValidationError("`mcp.url` and `mcp.command` cannot both be set.")

            mcp_cfg = MCPConfig(
                url=str(url) if url is not None else None,
                command=str(command) if command is not None else None,
                args=list(mcp_raw.get("args", []) or []),
                env=dict(mcp_raw.get("env", {}) or {}),
                transport=str(mcp_raw.get("transport", "auto")),
                headers=dict(mcp_raw.get("headers", {}) or {}),
                timeout=int(mcp_raw.get("timeout", 30)),
                servers=None,
                namespace_strategy=namespace_strategy,  # type: ignore[arg-type]
            )

        prompt_allowed = {
            "system",
            "interpolate",
            "context_vars",
            "inject_tool_descriptions",
            "inject_internal_instructions",
            "user_suffix",
            "user_prefix",
        }
        _warn_unknown_keys(prompt_raw, prompt_allowed, path="prompt")

        prompt_cfg = PromptConfig(
            system=str(prompt_raw.get("system", "")),
            interpolate=bool(prompt_raw.get("interpolate", True)),
            context_vars=dict(prompt_raw.get("context_vars", {}) or {}),
            inject_tool_descriptions=bool(prompt_raw.get("inject_tool_descriptions", True)),
            inject_internal_instructions=bool(prompt_raw.get("inject_internal_instructions", True)),
            user_suffix=str(prompt_raw.get("user_suffix", "")),
            user_prefix=str(prompt_raw.get("user_prefix", "")),
        )

        loop_allowed = {
            "max_iterations",
            "max_tokens_total",
            "tool_timeout",
            "parallel_tool_calls",
            "on_tool_call",
            "on_tool_result",
            "on_iteration",
            "retry_on_tool_error",
            "error_strategy",
        }
        _warn_unknown_keys(loop_raw, loop_allowed, path="loop")

        error_strategy = str(loop_raw.get("error_strategy", "return_error_to_llm"))
        if error_strategy not in {"raise", "return_error_to_llm", "skip"}:
            raise ConfigValidationError("`loop.error_strategy` must be one of: raise|return_error_to_llm|skip")

        loop_cfg = LoopConfig(
            max_iterations=int(loop_raw.get("max_iterations", 10)),
            max_tokens_total=int(loop_raw.get("max_tokens_total", 32000)),
            tool_timeout=int(loop_raw.get("tool_timeout", 30)),
            parallel_tool_calls=bool(loop_raw.get("parallel_tool_calls", True)),
            on_tool_call=loop_raw.get("on_tool_call"),
            on_tool_result=loop_raw.get("on_tool_result"),
            on_iteration=loop_raw.get("on_iteration"),
            retry_on_tool_error=bool(loop_raw.get("retry_on_tool_error", True)),
            error_strategy=error_strategy,  # type: ignore[arg-type]
        )

        session_allowed = {"persist_history", "max_history_tokens", "history_trim_strategy"}
        _warn_unknown_keys(session_raw, session_allowed, path="session")
        trim_strategy = str(session_raw.get("history_trim_strategy", "oldest_first"))
        if trim_strategy not in {"oldest_first", "summarize"}:
            raise ConfigValidationError("`session.history_trim_strategy` must be one of: oldest_first|summarize")

        session_cfg = SessionConfig(
            persist_history=bool(session_raw.get("persist_history", True)),
            max_history_tokens=int(session_raw.get("max_history_tokens", 32000)),
            history_trim_strategy=trim_strategy,  # type: ignore[arg-type]
        )

        return BridgeConfig(llm=llm_cfg, mcp=mcp_cfg, prompt=prompt_cfg, loop=loop_cfg, session=session_cfg)

