"""LLM Token 用量记录 — Docker/NAS 内存 7 天 · 客户端 state/system/token/。"""

from __future__ import annotations

import json
import os
import threading
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from niannian.storage.sync.push_after_ingest import is_docker_runtime
from niannian.token.paths import token_usage_file

RETENTION_DAYS = 7


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _parse_ts(ts: str) -> datetime | None:
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt
    except (TypeError, ValueError):
        return None


class TokenUsageStore:
    def __init__(self, home: Path, config: dict[str, Any]) -> None:
        self._home = Path(home)
        self._config = config
        self._lock = threading.Lock()
        self._memory: list[dict[str, Any]] = []
        self._use_memory = is_docker_runtime(config)

    @property
    def backend(self) -> str:
        return "memory" if self._use_memory else "file"

    def record(
        self,
        *,
        user_id: str = "pinn",
        module: str,
        operation: str,
        provider: str = "",
        model: str = "",
        prompt_tokens: int = 0,
        completion_tokens: int = 0,
    ) -> None:
        row = {
            "ts": _now_iso(),
            "user_id": str(user_id or "pinn"),
            "module": str(module or "unknown"),
            "operation": str(operation or "llm"),
            "provider": str(provider or ""),
            "model": str(model or ""),
            "prompt_tokens": max(0, int(prompt_tokens or 0)),
            "completion_tokens": max(0, int(completion_tokens or 0)),
        }
        with self._lock:
            self._purge_locked()
            if self._use_memory:
                self._memory.append(row)
            else:
                path = token_usage_file(self._home)
                with path.open("a", encoding="utf-8") as f:
                    f.write(json.dumps(row, ensure_ascii=False) + "\n")

    def _purge_locked(self) -> None:
        cutoff = datetime.now(timezone.utc) - timedelta(days=RETENTION_DAYS)
        if self._use_memory:
            kept: list[dict[str, Any]] = []
            for row in self._memory:
                dt = _parse_ts(str(row.get("ts") or ""))
                if dt is None or dt >= cutoff:
                    kept.append(row)
            self._memory = kept
            return
        path = token_usage_file(self._home)
        if not path.is_file():
            return
        kept_lines: list[str] = []
        with path.open(encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                dt = _parse_ts(str(row.get("ts") or ""))
                if dt is None or dt >= cutoff:
                    kept_lines.append(line)
        tmp = path.with_suffix(".jsonl.tmp")
        tmp.write_text("\n".join(kept_lines) + ("\n" if kept_lines else ""), encoding="utf-8")
        os.replace(tmp, path)

    def _load_rows_locked(self) -> list[dict[str, Any]]:
        self._purge_locked()
        if self._use_memory:
            return list(self._memory)
        path = token_usage_file(self._home)
        if not path.is_file():
            return []
        rows: list[dict[str, Any]] = []
        with path.open(encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    rows.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
        return rows

    def summary(self, *, user_id: str | None = None, days: int = RETENTION_DAYS) -> dict[str, Any]:
        days = max(1, min(int(days or RETENTION_DAYS), RETENTION_DAYS))
        cutoff = datetime.now(timezone.utc) - timedelta(days=days)
        with self._lock:
            rows = self._load_rows_locked()
        filtered: list[dict[str, Any]] = []
        for row in rows:
            if user_id and str(row.get("user_id") or "") != user_id:
                continue
            dt = _parse_ts(str(row.get("ts") or ""))
            if dt is not None and dt < cutoff:
                continue
            filtered.append(row)

        total_prompt = sum(int(r.get("prompt_tokens") or 0) for r in filtered)
        total_completion = sum(int(r.get("completion_tokens") or 0) for r in filtered)
        by_module: dict[str, dict[str, int]] = defaultdict(lambda: {"prompt_tokens": 0, "completion_tokens": 0, "calls": 0})
        by_day: dict[str, dict[str, int]] = defaultdict(lambda: {"prompt_tokens": 0, "completion_tokens": 0, "calls": 0})
        for row in filtered:
            mod = str(row.get("module") or "unknown")
            pt = int(row.get("prompt_tokens") or 0)
            ct = int(row.get("completion_tokens") or 0)
            by_module[mod]["prompt_tokens"] += pt
            by_module[mod]["completion_tokens"] += ct
            by_module[mod]["calls"] += 1
            dt = _parse_ts(str(row.get("ts") or ""))
            day = dt.strftime("%Y-%m-%d") if dt else "unknown"
            by_day[day]["prompt_tokens"] += pt
            by_day[day]["completion_tokens"] += ct
            by_day[day]["calls"] += 1

        module_rows = [
            {"module": k, **v, "total_tokens": v["prompt_tokens"] + v["completion_tokens"]}
            for k, v in sorted(by_module.items(), key=lambda x: -(x[1]["prompt_tokens"] + x[1]["completion_tokens"]))
        ]
        day_rows = [
            {"date": k, **v, "total_tokens": v["prompt_tokens"] + v["completion_tokens"]}
            for k, v in sorted(by_day.items())
        ]
        recent = sorted(filtered, key=lambda r: str(r.get("ts") or ""), reverse=True)[:80]

        return {
            "backend": self.backend,
            "retention_days": RETENTION_DAYS,
            "window_days": days,
            "totals": {
                "prompt_tokens": total_prompt,
                "completion_tokens": total_completion,
                "total_tokens": total_prompt + total_completion,
                "calls": len(filtered),
            },
            "by_module": module_rows,
            "by_day": day_rows,
            "recent": recent,
        }
