from __future__ import annotations

import os
from pathlib import Path
from typing import Literal

from pydantic import BaseModel, Field


def _config_dir() -> Path:
    """Return XDG-compliant config directory.

    Respects $XDG_CONFIG_HOME; defaults to ~/.config/probegpt.
    """
    xdg = os.environ.get("XDG_CONFIG_HOME")
    base = Path(xdg) if xdg else Path.home() / ".config"
    return base / "probegpt"


CONFIG_DIR = _config_dir()
CONFIG_PATH = CONFIG_DIR / "config.json"
CUSTOM_STRATEGIES_DIR = CONFIG_DIR / "strategies"
CUSTOM_PROMPTS_DIR = CONFIG_DIR / "prompts"

ProviderType = Literal["azure_openai", "openrouter", "vertex_ai", "bedrock", "cerebras", "mistral"]

_ALL_PROVIDERS: list[ProviderType] = [
    "azure_openai",
    "openrouter",
    "vertex_ai",
    "bedrock",
    "cerebras",
    "mistral",
]

ALL_STRATEGY_NAMES: list[str] = [
    "objective",
    "mutation",
    "encoding",
    "roleplay",
    "persuasion",
    "linguistic",
    "structural",
    "distraction",
]


class StrategyConfig(BaseModel):
    """Per-run strategy enable flags and relative weights."""

    enabled: list[str] = Field(default_factory=lambda: list(ALL_STRATEGY_NAMES))
    weights: dict[str, float] = Field(
        default_factory=lambda: {name: 1.0 for name in ALL_STRATEGY_NAMES}
    )

    def effective_weight(self, name: str) -> float:
        """Return the weight for a strategy, or 0 if disabled."""
        if name not in self.enabled:
            return 0.0
        return max(0.0, self.weights.get(name, 1.0))


class ProviderConfig(BaseModel):
    provider_type: ProviderType = "azure_openai"

    # Azure OpenAI — credentials via `az login`
    azure_endpoint: str = ""
    azure_deployment: str = ""
    azure_api_version: str = "2024-02-15-preview"

    # OpenRouter — API key from OPENROUTER_API_KEY env var
    openrouter_model: str = ""

    # GCP Vertex AI — credentials via `gcloud auth application-default login`
    vertex_model: str = "gemini-2.5-flash"
    vertex_project: str = ""
    vertex_region: str = "us-central1"

    # AWS Bedrock — credentials via `aws configure`
    bedrock_model: str = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
    bedrock_region: str = "us-east-1"
    bedrock_profile: str = ""

    # Cerebras — API key from CEREBRAS_API_KEY env var
    cerebras_model: str = "llama-3.3-70b"

    # Mistral — API key from MISTRAL_API_KEY env var
    mistral_model: str = "mistral-large-latest"


class Config(BaseModel):
    target: ProviderConfig = ProviderConfig()
    generator: ProviderConfig = ProviderConfig()
    judge: ProviderConfig = ProviderConfig()
    strategies: StrategyConfig = Field(default_factory=StrategyConfig)

    @classmethod
    def from_file(cls, path: Path = CONFIG_PATH) -> Config:
        if not path.exists():
            raise FileNotFoundError(f"Config not found at {path}. Run 'Setup Config' first.")
        return cls.model_validate_json(path.read_text())

    def to_file(self, path: Path = CONFIG_PATH) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(self.model_dump_json(indent=2))
