from __future__ import annotations

from enum import StrEnum, auto
import os
from pathlib import Path
import re
import shlex
import tomllib
from typing import Annotated, Any, Literal

from dotenv import dotenv_values
from pydantic import BaseModel, Field, TypeAdapter, field_validator, model_validator
from pydantic.fields import FieldInfo
from pydantic_core import to_jsonable_python
from pydantic_settings import (
    BaseSettings,
    PydanticBaseSettingsSource,
    SettingsConfigDict,
)
import tomli_w

from revibe.core.model_config import DEFAULT_MODELS, ModelConfig
from revibe.core.paths.config_paths import (
    AGENT_DIR,
    CONFIG_DIR,
    CONFIG_FILE,
    PROMPT_DIR,
)
from revibe.core.paths.global_paths import GLOBAL_ENV_FILE, SESSION_LOG_DIR
from revibe.core.prompts import SystemPrompt
from revibe.core.tools.base import BaseToolConfig

PROJECT_DOC_FILENAMES = ["AGENTS.md", "REVIBE.md", ".revibe.md"]


def load_api_keys_from_env() -> None:
    if GLOBAL_ENV_FILE.path.is_file():
        env_vars = dotenv_values(GLOBAL_ENV_FILE.path)
        for key, value in env_vars.items():
            if value:
                os.environ.setdefault(key, value)


class MissingAPIKeyError(RuntimeError):
    def __init__(self, env_key: str, provider_name: str) -> None:
        super().__init__(
            f"Missing {env_key} environment variable for {provider_name} provider"
        )
        self.env_key = env_key
        self.provider_name = provider_name


class MissingPromptFileError(RuntimeError):
    def __init__(self, system_prompt_id: str, prompt_dir: str) -> None:
        super().__init__(
            f"Invalid system_prompt_id value: '{system_prompt_id}'. "
            f"Must be one of the available prompts ({', '.join(f'{p.name.lower()}' for p in SystemPrompt)}), "
            f"or correspond to a .md file in {prompt_dir}"
        )
        self.system_prompt_id = system_prompt_id
        self.prompt_dir = prompt_dir


class WrongBackendError(RuntimeError):
    def __init__(self, backend: Backend, is_mistral_api: bool) -> None:
        super().__init__(
            f"Wrong backend '{backend}' for mistral API. Use '{Backend.MISTRAL}' for mistral API."
        )
        self.backend = backend
        self.is_mistral_api = is_mistral_api


class TomlFileSettingsSource(PydanticBaseSettingsSource):
    def __init__(self, settings_cls: type[BaseSettings]) -> None:
        super().__init__(settings_cls)
        self.toml_data = self._load_toml()

    def _load_toml(self) -> dict[str, Any]:
        file = CONFIG_FILE.path
        try:
            with file.open("rb") as f:
                return tomllib.load(f)
        except FileNotFoundError:
            return {}
        except tomllib.TOMLDecodeError as e:
            raise RuntimeError(f"Invalid TOML in {file}: {e}") from e
        except OSError as e:
            raise RuntimeError(f"Cannot read {file}: {e}") from e

    def get_field_value(
        self, field: FieldInfo, field_name: str
    ) -> tuple[Any, str, bool]:
        return self.toml_data.get(field_name), field_name, False

    def __call__(self) -> dict[str, Any]:
        return self.toml_data


class ProjectContextConfig(BaseSettings):
    max_chars: int = 40_000
    default_commit_count: int = 5
    max_doc_bytes: int = 32 * 1024
    truncation_buffer: int = 1_000
    max_depth: int = 3
    max_files: int = 1000
    max_dirs_per_level: int = 20
    timeout_seconds: float = 2.0


class SessionLoggingConfig(BaseSettings):
    save_dir: str = ""
    session_prefix: str = "session"
    enabled: bool = True

    @field_validator("save_dir", mode="before")
    @classmethod
    def set_default_save_dir(cls, v: str) -> str:
        if not v:
            return str(SESSION_LOG_DIR.path)
        return v

    @field_validator("save_dir", mode="after")
    @classmethod
    def expand_save_dir(cls, v: str) -> str:
        return str(Path(v).expanduser().resolve())


class Backend(StrEnum):
    MISTRAL = auto()
    GENERIC = auto()
    OPENAI = auto()
    HUGGINGFACE = auto()
    GROQ = auto()
    CEREBRAS = auto()
    OLLAMA = auto()
    LLAMACPP = auto()
    QWEN = auto()


class ToolFormat(StrEnum):
    """Tool calling format for LLM interactions.

    NATIVE: Use the API's native function/tool calling mechanism
    XML: Use XML-based tool calling embedded in prompts (for models without native support)
    """

    NATIVE = auto()
    XML = auto()


class _ProviderBase(BaseModel):
    name: str
    api_base: str
    api_key_env_var: str = ""
    api_style: str = "openai"


class MistralProviderConfig(_ProviderBase):
    backend: Literal[Backend.MISTRAL] = Backend.MISTRAL


class OpenAIProviderConfig(_ProviderBase):
    backend: Literal[Backend.OPENAI] = Backend.OPENAI


class GroqProviderConfig(_ProviderBase):
    backend: Literal[Backend.GROQ] = Backend.GROQ


class HuggingFaceProviderConfig(_ProviderBase):
    backend: Literal[Backend.HUGGINGFACE] = Backend.HUGGINGFACE


class OllamaProviderConfig(_ProviderBase):
    backend: Literal[Backend.OLLAMA] = Backend.OLLAMA


class LlamaCppProviderConfig(_ProviderBase):
    backend: Literal[Backend.LLAMACPP] = Backend.LLAMACPP


class CerebrasProviderConfig(_ProviderBase):
    backend: Literal[Backend.CEREBRAS] = Backend.CEREBRAS


class GenericProviderConfig(_ProviderBase):
    backend: Literal[Backend.GENERIC] = Backend.GENERIC


class QwenProviderConfig(_ProviderBase):
    backend: Literal[Backend.QWEN] = Backend.QWEN


ProviderConfigUnion = Annotated[
    MistralProviderConfig
    | OpenAIProviderConfig
    | GroqProviderConfig
    | HuggingFaceProviderConfig
    | OllamaProviderConfig
    | LlamaCppProviderConfig
    | CerebrasProviderConfig
    | GenericProviderConfig
    | QwenProviderConfig,
    Field(discriminator="backend"),
]


type ProviderConfig = ProviderConfigUnion
PROVIDER_CONFIG_ADAPTER = TypeAdapter(ProviderConfigUnion)


class _MCPBase(BaseModel):
    name: str = Field(description="Short alias used to prefix tool names")
    prompt: str | None = Field(
        default=None, description="Optional usage hint appended to tool descriptions"
    )

    @field_validator("name", mode="after")
    @classmethod
    def normalize_name(cls, v: str) -> str:
        normalized = re.sub(r"[^a-zA-Z0-9_-]", "_", v)
        normalized = normalized.strip("_-")
        return normalized[:256]


class _MCPHttpFields(BaseModel):
    url: str = Field(description="Base URL of the MCP HTTP server")
    headers: dict[str, str] = Field(
        default_factory=dict,
        description=(
            "Additional HTTP headers when using 'http' transport (e.g., Authorization or X-API-Key)."
        ),
    )
    api_key_env: str = Field(
        default="",
        description=(
            "Environment variable name containing an API token to send for HTTP transport."
        ),
    )
    api_key_header: str = Field(
        default="Authorization",
        description=(
            "HTTP header name to carry the token when 'api_key_env' is set (e.g., 'Authorization' or 'X-API-Key')."
        ),
    )
    api_key_format: str = Field(
        default="Bearer {token}",
        description=(
            "Format string for the header value when 'api_key_env' is set. Use '{token}' placeholder."
        ),
    )

    def http_headers(self) -> dict[str, str]:
        hdrs = dict(self.headers or {})
        env_var = (self.api_key_env or "").strip()
        if env_var and (token := os.getenv(env_var)):
            target = (self.api_key_header or "").strip() or "Authorization"
            if not any(h.lower() == target.lower() for h in hdrs):
                try:
                    value = (self.api_key_format or "{token}").format(token=token)
                except Exception:
                    value = token
                hdrs[target] = value
        return hdrs


class MCPHttp(_MCPBase, _MCPHttpFields):
    transport: Literal["http"]


class MCPStreamableHttp(_MCPBase, _MCPHttpFields):
    transport: Literal["streamable-http"]


class MCPStdio(_MCPBase):
    transport: Literal["stdio"]
    command: str | list[str]
    args: list[str] = Field(default_factory=list)

    def argv(self) -> list[str]:
        base = (
            shlex.split(self.command)
            if isinstance(self.command, str)
            else list(self.command or [])
        )
        return [*base, *self.args] if self.args else base


MCPServer = Annotated[
    MCPHttp | MCPStreamableHttp | MCPStdio, Field(discriminator="transport")
]


DEFAULT_PROVIDERS: list[ProviderConfigUnion] = [
    MistralProviderConfig(
        name="mistral",
        api_base="https://api.mistral.ai/v1",
        api_key_env_var="MISTRAL_API_KEY",
    ),
    OpenAIProviderConfig(
        name="openai",
        api_base="https://api.openai.com/v1",
        api_key_env_var="OPENAI_API_KEY",
    ),
    HuggingFaceProviderConfig(
        name="huggingface",
        # Use the Hugging Face router base URL for inference and model listing
        api_base="https://router.huggingface.co/v1",
        api_key_env_var="HUGGINGFACE_API_KEY",
    ),
    GroqProviderConfig(
        name="groq",
        api_base="https://api.groq.com/openai/v1",
        api_key_env_var="GROQ_API_KEY",
    ),
    OllamaProviderConfig(
        name="ollama", api_base="http://127.0.0.1:11434/v1", api_key_env_var=""
    ),
    LlamaCppProviderConfig(
        name="llamacpp", api_base="http://127.0.0.1:8080/v1", api_key_env_var=""
    ),
    CerebrasProviderConfig(
        name="cerebras",
        api_base="https://api.cerebras.ai/v1",
        api_key_env_var="CEREBRAS_API_KEY",
    ),
    QwenProviderConfig(
        name="qwencode",
        api_base="",  # Uses OAuth base URL from qwen backend
        api_key_env_var="",
    ),
]


class VibeConfig(BaseSettings):
    active_model: str = "devstral-2"
    textual_theme: str = "terminal"
    vim_keybindings: bool = False
    disable_welcome_banner_animation: bool = False
    displayed_workdir: str = ""
    auto_compact_threshold: int = 200_000
    context_warnings: bool = False
    instructions: str = ""
    workdir: Path | None = Field(default=None, exclude=True)
    system_prompt_id: str = "cli"
    include_commit_signature: bool = True
    include_model_info: bool = True
    include_project_context: bool = True
    include_prompt_detail: bool = True
    enable_update_checks: bool = True
    api_timeout: float = 720.0
    providers: list[ProviderConfigUnion] = Field(
        default_factory=lambda: list(DEFAULT_PROVIDERS)
    )
    models: list[ModelConfig] = Field(default_factory=lambda: list(DEFAULT_MODELS))

    project_context: ProjectContextConfig = Field(default_factory=ProjectContextConfig)
    session_logging: SessionLoggingConfig = Field(default_factory=SessionLoggingConfig)
    tools: dict[str, BaseToolConfig] = Field(default_factory=dict)
    tool_paths: list[Path] = Field(
        default_factory=list,
        description=(
            "Additional directories to search for custom tools. "
            "Each path may be absolute or relative to the current working directory."
        ),
    )

    mcp_servers: list[MCPServer] = Field(
        default_factory=list, description="Preferred MCP server configuration entries."
    )

    enabled_tools: list[str] = Field(
        default_factory=list,
        description=(
            "An explicit list of tool names/patterns to enable. If set, only these"
            " tools will be active. Supports exact names, glob patterns (e.g.,"
            " 'serena_*'), and regex with 're:' prefix or regex-like patterns (e.g.,"
            " 're:^serena_.*' or 'serena.*')."
        ),
    )
    disabled_tools: list[str] = Field(
        default_factory=list,
        description=(
            "A list of tool names/patterns to disable. Ignored if 'enabled_tools'"
            " is set. Supports exact names, glob patterns (e.g., 'bash*'), and"
            " regex with 're:' prefix or regex-like patterns."
        ),
    )

    skill_paths: list[Path] = Field(
        default_factory=list,
        description=(
            "Additional directories to search for skills. "
            "Each path may be absolute or relative to the current working directory."
        ),
    )

    tool_format: ToolFormat = Field(
        default=ToolFormat.NATIVE,
        description=(
            "Tool calling format: 'native' uses the API's function calling mechanism, "
            "'xml' embeds tool definitions in the system prompt for models without native support."
        ),
    )

    model_config = SettingsConfigDict(
        env_prefix="REVIBE_", case_sensitive=False, extra="ignore"
    )

    @property
    def effective_workdir(self) -> Path:
        return self.workdir if self.workdir is not None else Path.cwd()

    @property
    def system_prompt(self) -> str:
        try:
            return SystemPrompt[self.system_prompt_id.upper()].read()
        except KeyError:
            pass

        custom_sp_path = (PROMPT_DIR.path / self.system_prompt_id).with_suffix(".md")
        if not custom_sp_path.is_file():
            raise MissingPromptFileError(self.system_prompt_id, str(PROMPT_DIR.path))
        return custom_sp_path.read_text()

    def get_active_model(self) -> ModelConfig:
        for model in self.models:
            # Handle both ModelConfig objects and dicts (can happen with model_construct)
            m_alias = (
                model.alias if isinstance(model, ModelConfig) else model.get("alias")
            )
            if m_alias == self.active_model:
                return (
                    model
                    if isinstance(model, ModelConfig)
                    else ModelConfig.model_validate(model)
                )
        raise ValueError(
            f"Active model '{self.active_model}' not found in configuration."
        )

    def get_provider_for_model(self, model: ModelConfig) -> ProviderConfigUnion:
        # Merge DEFAULT_PROVIDERS with configured providers
        providers_map: dict[str, Any] = {}
        for p in DEFAULT_PROVIDERS:
            providers_map[p.name] = p
        for p in self.providers:
            p_name = p.name if not isinstance(p, dict) else p.get("name")
            providers_map[p_name] = p

        m_provider = (
            model.provider if isinstance(model, ModelConfig) else model.get("provider")
        )
        provider = providers_map.get(m_provider)

        if provider is None:
            m_name = getattr(model, "name", None) or (
                model.get("name") if isinstance(model, dict) else "unknown"
            )
            raise ValueError(
                f"Provider '{m_provider}' for model '{m_name}' not found in configuration."
            )

        return (
            provider
            if not isinstance(provider, dict)
            else PROVIDER_CONFIG_ADAPTER.validate_python(provider)
        )

    @classmethod
    def settings_customise_sources(
        cls,
        settings_cls: type[BaseSettings],
        init_settings: PydanticBaseSettingsSource,
        env_settings: PydanticBaseSettingsSource,
        dotenv_settings: PydanticBaseSettingsSource,
        file_secret_settings: PydanticBaseSettingsSource,
    ) -> tuple[PydanticBaseSettingsSource, ...]:
        """Define the priority of settings sources.

        Note: dotenv_settings is intentionally excluded. API keys and other
        non-config environment variables are stored in .env but loaded manually
        into os.environ for use by providers. Only REVIBE_* prefixed environment
        variables (via env_settings) and TOML config are used for Pydantic settings.
        """
        return (
            init_settings,
            env_settings,
            TomlFileSettingsSource(settings_cls),
            file_secret_settings,
        )

    @model_validator(mode="after")
    def _check_api_key(self) -> VibeConfig:
        try:
            active_model = self.get_active_model()
            provider = self.get_provider_for_model(active_model)
            api_key_env = provider.api_key_env_var
            if api_key_env and not os.getenv(api_key_env):
                raise MissingAPIKeyError(api_key_env, provider.name)
        except ValueError:
            pass
        return self

    @model_validator(mode="after")
    def _check_api_backend_compatibility(self) -> VibeConfig:
        try:
            active_model = self.get_active_model()
            provider = self.get_provider_for_model(active_model)
            MISTRAL_API_BASES = [
                "https://codestral.mistral.ai",
                "https://api.mistral.ai",
            ]
            is_mistral_api = any(
                provider.api_base.startswith(api_base) for api_base in MISTRAL_API_BASES
            )
            if is_mistral_api and provider.backend != Backend.MISTRAL:
                raise WrongBackendError(provider.backend, is_mistral_api)

        except ValueError:
            pass
        return self

    @field_validator("tool_paths", mode="before")
    @classmethod
    def _expand_tool_paths(cls, v: Any) -> list[Path]:
        if not v:
            return []
        return [Path(p).expanduser().resolve() for p in v]

    @field_validator("skill_paths", mode="before")
    @classmethod
    def _expand_skill_paths(cls, v: Any) -> list[Path]:
        if not v:
            return []
        return [Path(p).expanduser().resolve() for p in v]

    @field_validator("models", mode="before")
    @classmethod
    def _validate_models(cls, v: Any) -> list[ModelConfig]:
        if not isinstance(v, list):
            return list(DEFAULT_MODELS)
        return [ModelConfig.model_validate(item) for item in v]

    @field_validator("tools", mode="before")
    @classmethod
    def _normalize_tool_configs(cls, v: Any) -> dict[str, BaseToolConfig]:
        if not isinstance(v, dict):
            return {}

        normalized: dict[str, BaseToolConfig] = {}
        for tool_name, tool_config in v.items():
            if isinstance(tool_config, BaseToolConfig):
                normalized[tool_name] = tool_config
            elif isinstance(tool_config, dict):
                normalized[tool_name] = BaseToolConfig.model_validate(tool_config)
            else:
                normalized[tool_name] = BaseToolConfig()

        return normalized

    @model_validator(mode="after")
    def _validate_model_uniqueness(self) -> VibeConfig:
        provider_aliases: dict[str, set[str]] = {}
        for model in self.models:
            if model.provider not in provider_aliases:
                provider_aliases[model.provider] = set()
            if model.alias in provider_aliases[model.provider]:
                raise ValueError(
                    f"Duplicate model alias found for provider '{model.provider}': '{model.alias}'. Aliases must be unique within each provider."
                )
            provider_aliases[model.provider].add(model.alias)
        return self

    @model_validator(mode="after")
    def _merge_default_models(self) -> VibeConfig:
        from revibe.core.model_config import DEFAULT_MODELS

        existing_keys = {(m.provider, m.name) for m in self.models}
        for m in DEFAULT_MODELS:
            if (m.provider, m.name) not in existing_keys:
                self.models.append(m)
        return self

    @model_validator(mode="after")
    def _check_system_prompt(self) -> VibeConfig:
        _ = self.system_prompt
        return self

    @classmethod
    def save_updates(cls, updates: dict[str, Any]) -> None:
        CONFIG_DIR.path.mkdir(parents=True, exist_ok=True)
        current_config = TomlFileSettingsSource(cls).toml_data

        def deep_merge(target: dict, source: dict) -> None:
            for key, value in source.items():
                if (
                    key in target
                    and isinstance(target.get(key), dict)
                    and isinstance(value, dict)
                ):
                    deep_merge(target[key], value)
                elif (
                    key in target
                    and isinstance(target.get(key), list)
                    and isinstance(value, list)
                ):
                    if key in {"providers", "models"}:
                        target[key] = value
                    else:
                        target[key] = list(set(value + target[key]))
                else:
                    target[key] = value

        deep_merge(current_config, updates)
        cls.dump_config(
            to_jsonable_python(current_config, exclude_none=True, fallback=str)
        )

    @classmethod
    def dump_config(cls, config: dict[str, Any]) -> None:
        with CONFIG_FILE.path.open("wb") as f:
            tomli_w.dump(config, f)

    @classmethod
    def _get_agent_config(cls, agent: str | None) -> dict[str, Any] | None:
        if agent is None:
            return None

        agent_config_path = (AGENT_DIR.path / agent).with_suffix(".toml")
        try:
            return tomllib.load(agent_config_path.open("rb"))
        except FileNotFoundError:
            raise ValueError(
                f"Config '{agent}.toml' for agent not found in {AGENT_DIR.path}"
            )

    @classmethod
    def _migrate(cls) -> None:
        pass

    @classmethod
    def load(cls, agent: str | None = None, **overrides: Any) -> VibeConfig:
        cls._migrate()
        agent_config = cls._get_agent_config(agent)
        init_data = {**(agent_config or {}), **overrides}
        return cls(**init_data)

    @classmethod
    def create_default(cls) -> dict[str, Any]:
        try:
            config = cls()
        except MissingAPIKeyError:
            config = cls.model_construct()

        config_dict = config.model_dump(mode="json", exclude_none=True)

        from revibe.core.tools.manager import ToolManager

        tool_defaults = ToolManager.discover_tool_defaults()
        if tool_defaults:
            config_dict["tools"] = tool_defaults

        return config_dict
