from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import openai

from .pricing import get_numeric_token_rates, get_provider_tier, normalize_pricing_provider

PROVIDER_OPENAI = "openai"
PROVIDER_AZURE = "azure"
PROVIDER_AZURE_ALIASES = {"azure", "azure_openai", "azure-openai"}
PROVIDER_OPENAI_ALIASES = {"openai"}


def _normalize_provider(value: Optional[str], *, strict: bool = False) -> str:
    raw = (value or PROVIDER_OPENAI).strip().lower()
    if raw in PROVIDER_OPENAI_ALIASES:
        return PROVIDER_OPENAI
    if raw in PROVIDER_AZURE_ALIASES:
        return PROVIDER_AZURE
    if strict:
        raise ValueError("model_config['provider'] must be 'openai' or 'azure'")
    return PROVIDER_OPENAI


def _clean_string(value: Any) -> Optional[str]:
    if value is None:
        return None
    if not isinstance(value, str):
        raise TypeError("model_config values must be strings when provided.")
    cleaned = value.strip()
    return cleaned or None


def _normalize_base_url(base_url: Optional[str]) -> Optional[str]:
    if not base_url:
        return None
    if "openai.azure.com" in base_url and "/openai/" not in base_url:
        return base_url.rstrip("/") + "/openai/v1"
    return base_url


def _coerce_model_config(model_config: Optional[Dict[str, Any]]) -> Dict[str, Optional[str]]:
    if model_config is None:
        return {}
    if not isinstance(model_config, dict):
        raise TypeError("model_config must be a dict when provided.")

    provider_raw = model_config.get("provider", model_config.get("name"))
    provider = _normalize_provider(_clean_string(provider_raw), strict=True) if provider_raw is not None else None
    model = _clean_string(model_config.get("model", model_config.get("deployment_name")))
    pricing_model = _clean_string(
        model_config.get(
            "pricing_model",
            model_config.get("deployment_raw_name", model_config.get("raw_model")),
        )
    )
    api_key = _clean_string(model_config.get("api_key"))
    base_url = _normalize_base_url(
        _clean_string(model_config.get("base_url", model_config.get("endpoint", model_config.get("azure_endpoint"))))
    )

    return {
        "provider": provider,
        "model": model,
        "pricing_model": pricing_model,
        "api_key": api_key,
        "base_url": base_url,
    }


def merge_model_into_model_config(
    model: Optional[str],
    model_config: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
    config = _coerce_model_config(model_config)
    explicit_model = _clean_string(model)
    if explicit_model:
        config["model"] = explicit_model
    return config or None


def _resolve_connection_values(
    provider: str,
    model_config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Optional[str]]:
    config = _coerce_model_config(model_config)
    api_key = config.get("api_key")
    base_url = config.get("base_url")

    if provider == PROVIDER_AZURE:
        api_key = api_key or os.getenv("AZURE_OPENAI_API_KEY")
        base_url = base_url or os.getenv("AZURE_OPENAI_BASE_URL") or os.getenv("AZURE_OPENAI_ENDPOINT")
    else:
        api_key = api_key or os.getenv("OPENAI_API_KEY")
        base_url = base_url

    return {
        "api_key": api_key,
        "base_url": _normalize_base_url(base_url),
    }


def get_llm_provider(model_config: Optional[Dict[str, Any]] = None) -> str:
    config = _coerce_model_config(model_config)
    if config.get("provider"):
        return str(config["provider"])
    provider = os.getenv("NAVEXA_LLM_PROVIDER", PROVIDER_OPENAI)
    return _normalize_provider(provider, strict=False)


def has_llm_credentials(model_config: Optional[Dict[str, Any]] = None) -> bool:
    if model_config is not None:
        provider = get_llm_provider(model_config=model_config)
        resolved = _resolve_connection_values(provider, model_config=model_config)
        if provider == PROVIDER_AZURE:
            return bool(resolved["api_key"] and resolved["base_url"])
        return bool(resolved["api_key"])

    global _GLOBAL
    if _GLOBAL is not None:
        provider = _GLOBAL.settings.provider or get_llm_provider()
        resolved = _resolve_connection_values(
            provider,
            {
                "provider": provider,
                "api_key": _GLOBAL.settings.api_key,
                "base_url": _GLOBAL.settings.base_url,
            },
        )
        api_key = resolved["api_key"]
        base_url = resolved["base_url"]
        if provider == PROVIDER_AZURE:
            return bool(api_key and base_url)
        return bool(api_key)

    provider = get_llm_provider()
    resolved = _resolve_connection_values(provider)
    if provider == PROVIDER_AZURE:
        return bool(resolved["api_key"] and resolved["base_url"])
    return bool(resolved["api_key"])


def resolve_model_name(
    user_model: Optional[str],
    *,
    use_llm_requested: bool = True,
    model_config: Optional[Dict[str, Any]] = None,
) -> Optional[str]:
    if not use_llm_requested:
        return None
    if user_model and user_model.strip():
        return user_model.strip()
    config = _coerce_model_config(model_config)
    explicit_model = config.get("model")
    if explicit_model:
        return explicit_model
    provider = get_llm_provider(model_config=model_config)
    if provider == PROVIDER_AZURE:
        value = os.getenv("AZURE_DEPLOYMENT_NAME")
        if value and value.strip():
            return value.strip()
    else:
        value = os.getenv("OPENAI_MODEL_NAME")
        if value and value.strip():
            return value.strip()
    return None


def resolve_pricing_model_name(selected_model: Optional[str], *, model_config: Optional[Dict[str, Any]] = None) -> str:
    config = _coerce_model_config(model_config)
    explicit_pricing_model = config.get("pricing_model")
    if explicit_pricing_model:
        return explicit_pricing_model
    provider = get_llm_provider(model_config=model_config)
    if provider == PROVIDER_AZURE:
        raw = os.getenv("AZURE_DEPLOYMENT_RAW_NAME")
        if raw and raw.strip():
            return raw.strip()
    if selected_model and selected_model.strip():
        return selected_model.strip()
    return "no-model-defined"


def _get(obj: Any, key: str, default: Any = None) -> Any:
    if obj is None:
        return default
    if isinstance(obj, dict):
        return obj.get(key, default)
    return getattr(obj, key, default)


@dataclass
class UsageCostTracker:
    pricing_provider: str = PROVIDER_OPENAI
    pricing_model: str = "no-model-defined"
    tier: str = "standard"
    calls: int = 0
    input_tokens: int = 0
    cached_input_tokens: int = 0
    output_tokens: int = 0

    def add_usage(self, usage: Any) -> None:
        in_toks = _get(usage, "prompt_tokens", _get(usage, "input_tokens", 0))
        out_toks = _get(usage, "completion_tokens", _get(usage, "output_tokens", 0))
        details = _get(usage, "prompt_tokens_details", _get(usage, "input_tokens_details", None))
        cached = _get(details, "cached_tokens", 0) if details is not None else 0

        self.calls += 1
        self.input_tokens += int(in_toks or 0)
        self.output_tokens += int(out_toks or 0)
        self.cached_input_tokens += int(cached or 0)

    def cost_usd(self) -> float:
        prices = get_numeric_token_rates(self.pricing_provider, self.pricing_model)
        if prices is None:
            return 0.0
        return (
            (self.input_tokens / 1_000_000) * prices["input"]
            + (self.cached_input_tokens / 1_000_000) * prices["cached_input"]
            + (self.output_tokens / 1_000_000) * prices["output"]
        )

    def summary(self) -> Dict[str, Any]:
        return {
            "pricing_model": self.pricing_model,
            "tier": self.tier,
            "calls": self.calls,
            "input_tokens": self.input_tokens,
            "cached_input_tokens": self.cached_input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.input_tokens + self.output_tokens,
            "estimated_cost_usd": round(self.cost_usd(), 6),
        }


@dataclass
class LLMSettings:
    model: Optional[str] = None
    provider: Optional[str] = None
    api_key: Optional[str] = None
    base_url: Optional[str] = None
    pricing_model: Optional[str] = None
    system_prompt: str = (
        "You process regulatory and medical PDF text for structure extraction and summarization. "
        "Keep extraction faithful and concise."
    )


def build_llm_settings(
    *,
    model: Optional[str] = None,
    model_config: Optional[Dict[str, Any]] = None,
    use_llm_requested: bool = True,
    system_prompt: Optional[str] = None,
) -> LLMSettings:
    provider = get_llm_provider(model_config=model_config)
    connection = _resolve_connection_values(provider, model_config=model_config)
    selected_model = resolve_model_name(model, use_llm_requested=use_llm_requested, model_config=model_config)
    pricing_model = resolve_pricing_model_name(selected_model, model_config=model_config)
    return LLMSettings(
        model=selected_model,
        provider=provider,
        api_key=connection["api_key"],
        base_url=connection["base_url"],
        pricing_model=pricing_model,
        system_prompt=system_prompt
        or (
            "You process regulatory and medical PDF text for structure extraction and summarization. "
            "Keep extraction faithful and concise."
        ),
    )


def serialize_model_config(
    *,
    model: Optional[str],
    model_config: Optional[Dict[str, Any]] = None,
    use_llm_requested: bool = True,
) -> Dict[str, Any]:
    settings = build_llm_settings(model=model, model_config=model_config, use_llm_requested=use_llm_requested)
    source = "env"
    if model_config is not None and model and str(model).strip():
        source = "model_override"
    elif model_config is not None:
        source = "model_config"
    elif model and str(model).strip():
        source = "model_arg"
    return {
        "provider": settings.provider or PROVIDER_OPENAI,
        "model": settings.model,
        "pricing_model": settings.pricing_model or "no-model-defined",
        "api_key_configured": bool(settings.api_key),
        "base_url_configured": bool(settings.base_url),
        "source": source,
    }


class LLMClient:
    def __init__(self, settings: Optional[LLMSettings] = None) -> None:
        self.settings = settings or LLMSettings()
        provider = self.settings.provider or get_llm_provider()
        pricing_provider = normalize_pricing_provider(provider)
        self.tracker = UsageCostTracker(
            pricing_provider=pricing_provider,
            pricing_model=self.settings.pricing_model
            or resolve_pricing_model_name(self.settings.model, model_config={"provider": provider}),
            tier=get_provider_tier(pricing_provider),
        )
        self.client: Optional[openai.OpenAI] = None
        self.async_client: Optional[openai.AsyncOpenAI] = None

    def _kwargs(self) -> Dict[str, Any]:
        provider = self.settings.provider or get_llm_provider()
        resolved = _resolve_connection_values(
            provider,
            {
                "provider": provider,
                "api_key": self.settings.api_key,
                "base_url": self.settings.base_url,
            },
        )
        api_key = resolved["api_key"]
        base_url = resolved["base_url"]
        kwargs: Dict[str, Any] = {}
        if api_key:
            kwargs["api_key"] = api_key
        if base_url:
            kwargs["base_url"] = base_url
        return kwargs

    def _sync(self) -> openai.OpenAI:
        if self.client is None:
            self.client = openai.OpenAI(**self._kwargs())
        return self.client

    def _async(self) -> openai.AsyncOpenAI:
        if self.async_client is None:
            self.async_client = openai.AsyncOpenAI(**self._kwargs())
        return self.async_client

    def _messages(self, prompt: str, chat_history: Optional[List[Dict[str, str]]] = None) -> List[Dict[str, str]]:
        messages = [{"role": "system", "content": self.settings.system_prompt}]
        if chat_history:
            messages.extend(chat_history)
        messages.append({"role": "user", "content": prompt})
        return messages

    def chat(self, model: str, prompt: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Any:
        self.tracker.pricing_model = resolve_pricing_model_name(model)
        response = self._sync().chat.completions.create(
            model=model,
            messages=self._messages(prompt, chat_history=chat_history),
            temperature=0,
        )
        if getattr(response, "usage", None) is not None:
            self.tracker.add_usage(response.usage)
        return response

    async def chat_async(self, model: str, prompt: str) -> Any:
        self.tracker.pricing_model = resolve_pricing_model_name(model)
        response = await self._async().chat.completions.create(
            model=model,
            messages=self._messages(prompt),
            temperature=0,
        )
        if getattr(response, "usage", None) is not None:
            self.tracker.add_usage(response.usage)
        return response


_GLOBAL: Optional[LLMClient] = None


def get_llm_client() -> LLMClient:
    global _GLOBAL
    if _GLOBAL is None:
        _GLOBAL = LLMClient()
    return _GLOBAL


def configure_llm_client(settings: Optional[LLMSettings] = None) -> LLMClient:
    global _GLOBAL
    desired = settings or LLMSettings()
    if _GLOBAL is None or _GLOBAL.settings != desired:
        _GLOBAL = LLMClient(settings=desired)
    return _GLOBAL


def reset_llm_client(settings: Optional[LLMSettings] = None) -> LLMClient:
    global _GLOBAL
    _GLOBAL = LLMClient(settings=settings)
    return _GLOBAL
