from __future__ import annotations

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

import openai


GPT41_PRICING_USD_PER_1M = {
    "standard": {"input": 2.00, "cached_input": 0.50, "output": 8.00}
}

PROVIDER_OPENAI = "openai"
PROVIDER_AZURE = "azure"


def get_llm_provider() -> str:
    provider = (os.getenv("NAVEXA_LLM_PROVIDER", PROVIDER_OPENAI) or PROVIDER_OPENAI).strip().lower()
    if provider not in {PROVIDER_OPENAI, PROVIDER_AZURE}:
        return PROVIDER_OPENAI
    return provider


def has_llm_credentials() -> bool:
    provider = get_llm_provider()
    if provider == PROVIDER_AZURE:
        return bool(os.getenv("AZURE_OPENAI_API_KEY") and os.getenv("AZURE_OPENAI_BASE_URL"))
    return bool(os.getenv("OPENAI_API_KEY"))


def resolve_model_name(user_model: Optional[str], *, use_llm_requested: bool = True) -> Optional[str]:
    if not use_llm_requested:
        return None
    if user_model and user_model.strip():
        return user_model.strip()
    provider = get_llm_provider()
    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]) -> str:
    provider = get_llm_provider()
    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 "gpt-4.1"


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_model: str = "gpt-4.1"
    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 = GPT41_PRICING_USD_PER_1M[self.tier]
        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
    system_prompt: str = (
        "You process regulatory and medical PDF text for structure extraction and summarization. "
        "Keep extraction faithful and concise."
    )


class LLMClient:
    def __init__(self, settings: Optional[LLMSettings] = None) -> None:
        self.settings = settings or LLMSettings()
        self.tracker = UsageCostTracker(pricing_model=resolve_pricing_model_name(self.settings.model))
        self.client: Optional[openai.OpenAI] = None
        self.async_client: Optional[openai.AsyncOpenAI] = None

    def _kwargs(self) -> Dict[str, Any]:
        provider = get_llm_provider()
        if provider == PROVIDER_AZURE:
            api_key = os.getenv("AZURE_OPENAI_API_KEY")
            base_url = os.getenv("AZURE_OPENAI_BASE_URL")
        else:
            api_key = os.getenv("OPENAI_API_KEY")
            base_url = None
        if base_url and "openai.azure.com" in base_url and "/openai/" not in base_url:
            base_url = base_url.rstrip("/") + "/openai/v1"
        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:
        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:
        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 reset_llm_client(settings: Optional[LLMSettings] = None) -> LLMClient:
    global _GLOBAL
    _GLOBAL = LLMClient(settings=settings)
    return _GLOBAL
