"""BilgeAI configuration management."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Literal

import yaml
from pydantic import BaseModel, Field, field_validator

from bilgeai.core.constants import (
    CONFIG_FILE,
    PROVIDER_DEFAULTS,
    SCHEMA_VERSION,
    SUPPORTED_MAX_VERSION,
)
from bilgeai.core.exceptions import ConfigError

logger = logging.getLogger(__name__)


class ModelConfig(BaseModel):
    """Model configuration."""

    name: str = "meta-llama/Llama-3-8B"


class DataConfig(BaseModel):
    """Data configuration."""

    path: str = "./train.jsonl"
    format: str = "auto"

    @field_validator("path")
    @classmethod
    def strip_quotes(cls, v: str) -> str:
        return v.strip("\"'")


class TrainingOverrides(BaseModel):
    """User-specified training parameter overrides."""

    batch_size: int | None = None
    learning_rate: float | None = None
    num_epochs: int | None = None
    max_seq_length: int | None = None
    lora_rank: int | None = None
    quantization: str | None = None
    num_nodes: int | None = None
    gpus_per_node: int | None = None
    report_to: str | None = None
    wandb_project: str | None = None


class LLMConfig(BaseModel):
    """LLM provider configuration."""

    provider: Literal["anthropic", "openai", "google", "local"] = "anthropic"
    model: str = "claude-sonnet-4-20250514"
    api_key_env: str = "ANTHROPIC_API_KEY"
    base_url: str | None = None


class PrivacyConfig(BaseModel):
    """Privacy settings."""

    scan_pii: bool = True


class BilgeConfig(BaseModel):
    """Root configuration model for bilge.yaml."""

    schema_version: str = Field(default=SCHEMA_VERSION)
    model: ModelConfig = Field(default_factory=ModelConfig)
    data: DataConfig = Field(default_factory=DataConfig)
    training: TrainingOverrides = Field(default_factory=TrainingOverrides)
    llm: LLMConfig = Field(default_factory=LLMConfig)
    privacy: PrivacyConfig = Field(default_factory=PrivacyConfig)

    @classmethod
    def load(cls, path: Path | None = None) -> BilgeConfig:
        """Load config from bilge.yaml with auto-migration support."""
        config_path = path or Path.cwd() / CONFIG_FILE
        if not config_path.exists():
            raise ConfigError(
                f"Config file not found: {config_path}",
                suggestion="Run 'bilge init' to create a configuration file.",
            )
        with open(config_path) as f:
            data = yaml.safe_load(f) or {}

        # Downgrade protection: check schema version
        file_version = data.get("schema_version", "0.0.1")
        if _version_gt(file_version, SUPPORTED_MAX_VERSION):
            raise ConfigError(
                f"This bilge.yaml uses v{file_version} format. "
                f"Your BilgeAI version supports up to v{SCHEMA_VERSION}. "
                "Please upgrade: pip install --upgrade bilgeai",
            )

        # Auto-migrate from v0.1.0 format
        if _is_legacy_format(data):
            data = _migrate_v010_to_v020(data)
            logger.info("bilge.yaml migrated from legacy format to v0.1.0.")
            # Save migrated config
            config = cls.model_validate(data)
            config.save(config_path)
            return config

        return cls.model_validate(data)

    def save(self, path: Path | None = None) -> None:
        """Save config to bilge.yaml."""
        config_path = path or Path.cwd() / CONFIG_FILE
        data = self.model_dump(exclude_none=True)
        with open(config_path, "w") as f:
            yaml.dump(data, f, default_flow_style=False, allow_unicode=True, sort_keys=False)

    def get_overrides(self) -> dict[str, Any]:
        """Return non-None training overrides as a dict."""
        return {k: v for k, v in self.training.model_dump().items() if v is not None}


def _version_gt(v1: str, v2: str) -> bool:
    """Check if version v1 is greater than v2."""
    parts1 = [int(x) for x in v1.split(".")]
    parts2 = [int(x) for x in v2.split(".")]
    return parts1 > parts2


def _is_legacy_format(data: dict) -> bool:
    """Check if config uses legacy format (ollama/remote booleans)."""
    llm = data.get("llm", {})
    return "ollama" in llm or "remote" in llm


def _migrate_v010_to_v020(data: dict) -> dict:
    """Migrate legacy config format to v0.1.0."""
    llm = data.get("llm", {})

    if llm.get("ollama", False):
        # Ollama was enabled -> map to local provider
        new_llm = {
            "provider": "local",
            "model": "local-model",
            "api_key_env": "",
            "base_url": PROVIDER_DEFAULTS["local"]["base_url"],
        }
    elif llm.get("remote", False):
        # Remote was enabled -> map to anthropic (default)
        old_key_env = llm.get("api_key_env", "BILGE_API_KEY")
        new_llm = {
            "provider": "anthropic",
            "model": PROVIDER_DEFAULTS["anthropic"]["model"],
            "api_key_env": old_key_env if old_key_env != "BILGE_API_KEY" else "ANTHROPIC_API_KEY",
        }
    else:
        # Neither was enabled -> default to anthropic
        new_llm = {
            "provider": "anthropic",
            "model": PROVIDER_DEFAULTS["anthropic"]["model"],
            "api_key_env": "ANTHROPIC_API_KEY",
        }

    data["llm"] = new_llm
    data["schema_version"] = SCHEMA_VERSION
    return data
