"""Cost tracker — estimate and accumulate API costs per request."""

from __future__ import annotations

import contextlib
import logging
import os
import threading
import time

from .utils.atomic_io import atomic_write_json, safe_load_json

log = logging.getLogger(__name__)

# Pricing per 1M tokens (input/output) as of early 2026
_MODEL_PRICING = {
    # DeepSeek
    "deepseek-chat": (0.27, 1.10),
    "deepseek-coder": (0.27, 1.10),
    "deepseek-reasoner": (0.55, 2.19),
    "deepseek-v3": (0.27, 1.10),
    # Aliyun DashScope (Qwen)
    "qwen-turbo": (0.28, 0.84),
    "qwen-plus": (0.80, 2.00),
    "qwen-max": (4.00, 12.00),
    "qwen-long": (0.50, 2.00),
    "qwen-vl-plus": (1.12, 3.36),
    "qwen-vl-max": (3.00, 9.00),
    "qwen2.5-72b-instruct": (0.56, 1.68),
    # 88API / OpenAI-compatible
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "gpt-4-turbo": (10.00, 30.00),
    "gpt-3.5-turbo": (0.50, 1.50),
    "o1": (15.00, 60.00),
    "o1-mini": (3.00, 12.00),
    # Anthropic Claude (via 88API)
    "claude-opus-4-6": (15.00, 75.00),
    "claude-sonnet-4-6": (3.00, 15.00),
    "claude-3-5-sonnet": (3.00, 15.00),
    "claude-3-5-haiku": (0.80, 4.00),
    "claude-3-haiku": (0.25, 1.25),
    # Ollama (free / local)
    "llama3.2": (0.0, 0.0),
    "llama3.3": (0.0, 0.0),
    "mistral": (0.0, 0.0),
    "qwen2.5": (0.0, 0.0),
    "qwen2.5:7b": (0.0, 0.0),
    "phi4": (0.0, 0.0),
}

_CHARS_PER_TOKEN_DEFAULT = 4


class CostTracker:
    """Track API costs per request and accumulate daily totals."""

    def __init__(self, data_dir: str | None = None) -> None:
        self._data_dir: str = data_dir or os.path.expanduser("~/.muravey-china")
        self._lock: threading.Lock = threading.Lock()
        self._today_cost: float = 0.0
        self._today_requests: int = 0
        self._today_date: str = time.strftime("%Y-%m-%d")
        self._load_today()

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        if not text:
            return 1
        non_ascii = sum(1 for c in text if ord(c) > 127)
        ascii_count = len(text) - non_ascii
        tokens = ascii_count / _CHARS_PER_TOKEN_DEFAULT + non_ascii / 1.5
        return max(1, int(tokens))

    def estimate_cost(self, model: str, input_text: str, output_text: str) -> float:
        input_tokens = self._estimate_tokens(input_text)
        output_tokens = self._estimate_tokens(output_text)

        pricing = _MODEL_PRICING.get(model)
        if not pricing:
            for key, val in _MODEL_PRICING.items():
                if key in model or model in key:
                    pricing = val
                    break
        if not pricing:
            pricing = (1.0, 3.0)

        input_cost = (input_tokens / 1_000_000) * pricing[0]
        output_cost = (output_tokens / 1_000_000) * pricing[1]
        return input_cost + output_cost

    def record(self, model: str, input_text: str, output_text: str) -> float:
        """Record a request cost. Returns the cost of this request."""
        cost = self.estimate_cost(model, input_text, output_text)
        with self._lock:
            today = time.strftime("%Y-%m-%d")
            if today != self._today_date:
                self._today_cost = 0.0
                self._today_requests = 0
                self._today_date = today

            self._today_cost += cost
            self._today_requests += 1
            self._save_today()

            log.debug(
                "Cost: $%.6f (model=%s, total=$%.4f, requests=%d)",
                cost,
                model,
                self._today_cost,
                self._today_requests,
            )
        return cost

    @property
    def today_cost(self) -> float:
        with self._lock:
            return self._today_cost

    @property
    def today_requests(self) -> int:
        with self._lock:
            return self._today_requests

    def format_cost(self, cost: float) -> str:
        if cost < 0.001:
            return f"${cost:.6f}"
        if cost < 0.01:
            return f"${cost:.4f}"
        return f"${cost:.2f}"

    def today_summary(self) -> str:
        with self._lock:
            return f"Today: {self.format_cost(self._today_cost)} ({self._today_requests} requests)"

    def _save_today(self) -> None:
        path = os.path.join(self._data_dir, "cost_today.json")
        with contextlib.suppress(OSError):
            atomic_write_json(
                path,
                {
                    "date": self._today_date,
                    "cost": self._today_cost,
                    "requests": self._today_requests,
                },
            )

    def _load_today(self) -> None:
        path = os.path.join(self._data_dir, "cost_today.json")
        data = safe_load_json(path, default={})
        with self._lock:
            if data.get("date") == self._today_date:
                self._today_cost = data.get("cost", 0.0)
                self._today_requests = data.get("requests", 0)
