"""LLM Gateway: unified interface for multiple model providers.

Supports OpenAI-compatible APIs (OpenAI, DeepSeek, vLLM, etc.)
"""
import json
import time
import httpx
from typing import Optional, Tuple
from app.config import settings
from app.models.worker import AIWorker
from app.models.cost import CostLog
from sqlalchemy.orm import Session


class LLMGateway:
    """Unified LLM gateway with token counting and cost tracking."""

    @staticmethod
    def _get_api_key(worker: AIWorker) -> str:
        return worker.api_key or settings.DEFAULT_LLM_API_KEY

    @staticmethod
    def _get_base_url(worker: AIWorker) -> str:
        return worker.base_url or settings.DEFAULT_LLM_BASE_URL

    @staticmethod
    def call(
        worker: AIWorker,
        messages: list[dict],
        db: Session,
        tenant_id: int,
        project_id: Optional[int] = None,
        task_id: Optional[int] = None,
    ) -> Tuple[str, CostLog]:
        """Make an LLM API call and record cost.

        Returns (response_content, cost_log).
        """
        api_key = LLMGateway._get_api_key(worker)
        base_url = LLMGateway._get_base_url(worker)
        url = f"{base_url.rstrip('/')}/chat/completions"

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": worker.model_name,
            "messages": messages,
            "temperature": worker.temperature,
        }

        start = time.time()
        cost_log = CostLog(
            tenant_id=tenant_id,
            project_id=project_id,
            task_id=task_id,
            worker_id=worker.id,
            provider=worker.provider,
            model_name=worker.model_name,
        )

        try:
            with httpx.Client(timeout=120) as client:
                resp = client.post(url, json=payload, headers=headers)
                resp.raise_for_status()
                data = resp.json()

            content = data["choices"][0]["message"]["content"]
            usage = data.get("usage", {})

            cost_log.prompt_tokens = usage.get("prompt_tokens", 0)
            cost_log.completion_tokens = usage.get("completion_tokens", 0)
            cost_log.total_tokens = usage.get("total_tokens", 0)
            cost_log.cost_cents = LLMGateway._estimate_cost(
                worker.model_name,
                cost_log.prompt_tokens,
                cost_log.completion_tokens,
            )
            cost_log.duration_ms = int((time.time() - start) * 1000)
            cost_log.status = "success"

            db.add(cost_log)
            db.commit()
            db.refresh(cost_log)

            return content, cost_log

        except Exception as e:
            cost_log.status = "error"
            cost_log.error_message = str(e)
            cost_log.duration_ms = int((time.time() - start) * 1000)
            db.add(cost_log)
            db.commit()
            raise

    @staticmethod
    def _estimate_cost(model: str, prompt_tokens: int, completion_tokens: int) -> int:
        """Estimate cost in cents. Rough pricing, adjustable."""
        # Per 1M tokens, in cents (1 USD = 100 cents)
        pricing = {
            "gpt-4o": {"input": 250, "output": 1000},
            "gpt-4o-mini": {"input": 15, "output": 60},
            "gpt-4-turbo": {"input": 1000, "output": 3000},
            "deepseek-chat": {"input": 14, "output": 28},
            "deepseek-reasoner": {"input": 55, "output": 219},
            "deepseek-v4-flash": {"input": 14, "output": 28},
            "qwen3.6-plus": {"input": 40, "output": 120},
        }
        rates = pricing.get(model, {"input": 15, "output": 60})  # default: cheap
        cost = (
            prompt_tokens * rates["input"] / 1_000_000
            + completion_tokens * rates["output"] / 1_000_000
        )
        return int(cost * 100)