import logging
import os
import yaml
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Any, Optional, Tuple


logger = logging.getLogger(__name__)

LLM_PROVIDER = "openai"
LLM_MODEL = "gpt-5.2-chat-latest"
LLM_TEMPERATURE = 0.7

# Default models per provider (used when BYO key determines the provider)
PROVIDER_MODEL_DEFAULTS = {
    "anthropic": "claude-sonnet-4-5-20250929",
    "openai": "gpt-5.2-chat-latest",
    "google-gla": "gemini-3-flash-preview",
    "google-vertex": "gemini-3-flash-preview",
}

# BYO key env var → (provider, pydantic-ai API key env var)
BYO_KEY_PROVIDER_MAP = {
    "BYO_ANTHROPIC_API_KEY": ("anthropic", "ANTHROPIC_API_KEY"),
    "BYO_OPENAI_API_KEY": ("openai", "OPENAI_API_KEY"),
    "BYO_GOOGLE_API_KEY": ("google-gla", "GEMINI_API_KEY"),
}

# Priority order for BYO keys (highest first)
BYO_KEY_PRIORITY = ["BYO_ANTHROPIC_API_KEY", "BYO_OPENAI_API_KEY", "BYO_GOOGLE_API_KEY"]


def resolve_llm_provider() -> Tuple[Optional[str], Optional[str], Optional[str]]:
    """Resolve LLM provider from BYO API key env vars.

    Priority: Anthropic > OpenAI > Google.

    Returns:
        Tuple of (provider, model, byo_env_var_name) if a BYO key is found,
        or (None, None, None) if no BYO key is set.
    """
    for byo_env_var in BYO_KEY_PRIORITY:
        key_value = os.environ.get(byo_env_var)
        if key_value:
            provider, target_env_var = BYO_KEY_PROVIDER_MAP[byo_env_var]
            model = PROVIDER_MODEL_DEFAULTS[provider]
            # Set the env var that pydantic-ai expects for this provider
            os.environ[target_env_var] = key_value
            logger.info("BYO key detected: %s → provider=%s, model=%s", byo_env_var, provider, model)
            return provider, model, byo_env_var
    return None, None, None

# Resolve relative to project root: config.py → llm_agent → puli_mcp_server → src → project root
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent
PROMPT_FILE_PATH = _PROJECT_ROOT / "prompts" / "analyze_code_prompt.yaml"


def load_prompt(prompt_file_path: Path) -> str:
    """Load prompt from file."""
    try:
        with open(prompt_file_path, 'r', encoding='utf-8') as f:
            prompt_text = f.read()
        data = yaml.safe_load(prompt_text)
        system_prompt = data.get("prompt")
    except yaml.YAMLError as e:
        raise ValueError(f"Error parsing YAML prompt file: {e}")
    return system_prompt


def _setup_vertex_env(config: Dict[str, Any]) -> None:
    """Set GCP environment variables required by pydantic-ai's google-vertex provider.

    pydantic-ai reads GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION from the
    environment when using the google-vertex provider (via ADC).

    Note: This mutates os.environ at startup. Safe because the MCP server runs
    as a single-user process (one per IDE session) with a fixed project/location.

    Args:
        config: Remote configuration dict containing VERTEX_PROJECT and VERTEX_LOCATION.

    Raises:
        ValueError: If VERTEX_PROJECT is not found in config or environment.
    """
    gcp_project = config.get("VERTEX_PROJECT") or os.environ.get("GOOGLE_CLOUD_PROJECT")
    gcp_location = config.get("VERTEX_LOCATION") or os.environ.get("GOOGLE_CLOUD_LOCATION")

    if not gcp_project:
        raise ValueError(
            "google-vertex provider requires a GCP project. "
            "Set VERTEX_PROJECT in remote config or GOOGLE_CLOUD_PROJECT env var."
        )

    os.environ["GOOGLE_CLOUD_PROJECT"] = gcp_project
    if gcp_location:
        os.environ["GOOGLE_CLOUD_LOCATION"] = gcp_location

    logger.info(
        "Vertex AI env configured: project=%s, location=%s",
        gcp_project,
        gcp_location or "(default)",
    )


@dataclass
class LLMAgentConfig:
    """Configuration for LLM Agent, loaded from environment variables."""

    provider: str
    model: str
    temperature: float
    system_prompt: str
    vertex_project: Optional[str] = None
    vertex_location: Optional[str] = None

    @classmethod
    def from_env(cls) -> "LLMAgentConfig":
        """Load configuration from environment variables and local prompt file."""
        provider = os.environ.get("LLM_PROVIDER", LLM_PROVIDER)
        model = os.environ.get("LLM_MODEL", LLM_MODEL)
        temperature = float(os.environ.get("LLM_TEMPERATURE", LLM_TEMPERATURE))

        # Load prompt from absolute path
        try:
            system_prompt = load_prompt(PROMPT_FILE_PATH)
            if not system_prompt:
                raise ValueError("System prompt 'prompt' key not found in YAML file")
            
        except FileNotFoundError:
            raise ValueError(f"Prompt file not found at {PROMPT_FILE_PATH}")
        except Exception as e:
            raise ValueError(f"Could not load prompt from {PROMPT_FILE_PATH}: {e}")

        return cls(
            provider=provider,
            model=model,
            temperature=temperature,
            system_prompt=system_prompt,
            vertex_project=os.environ.get("VERTEX_PROJECT"),
            vertex_location=os.environ.get("VERTEX_LOCATION"),
        )

    @classmethod
    def from_remote(
        cls,
        config: Dict[str, Any],
        prompts: Dict[str, Any],
        provider_override: Optional[str] = None,
        model_override: Optional[str] = None,
    ) -> "LLMAgentConfig":
        """Load configuration from remote config dict (fetched from proxy).

        Args:
            config: Configuration dictionary from proxy /config/mcp endpoint
            prompts: Prompts dictionary from proxy /config/mcp endpoint
            provider_override: If set, overrides the provider from remote config (determined by BYO key)
            model_override: If set, overrides the model from remote config (determined by BYO key)

        Returns:
            LLMAgentConfig instance with settings from remote config
        """
        provider = provider_override or config.get("LLM_PROVIDER", LLM_PROVIDER)
        model = model_override or config.get("LLM_MODEL", LLM_MODEL)
        temperature = config.get("LLM_TEMPERATURE", LLM_TEMPERATURE)

        # Set GCP env vars for Vertex-based providers (pydantic-ai reads these via ADC)
        if provider in ("google-vertex", "google-vertex-anthropic"):
            _setup_vertex_env(config)

        system_prompt = prompts.get("analyze_prompt") or load_prompt(PROMPT_FILE_PATH)
        if not system_prompt:
            print(f"analyze_prompt not found in prompts configuration: {prompts}")
            raise ValueError("analyze_prompt not found in prompts configuration")

        vertex_project = config.get("VERTEX_PROJECT") or os.environ.get("VERTEX_PROJECT")
        vertex_location = config.get("VERTEX_LOCATION") or os.environ.get("VERTEX_LOCATION")

        return cls(
            provider=provider,
            model=model,
            temperature=temperature,
            system_prompt=system_prompt,
            vertex_project=vertex_project,
            vertex_location=vertex_location,
        )

