from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping
from datetime import UTC, datetime, timedelta
from typing import Any, Protocol

from pymongo.errors import OperationFailure

from app.config import config
from app.services.context_budgeting import normalize_budget_text


def _now_utc() -> datetime:
    return datetime.now(UTC)


def _normalize_cache_value(value: Any) -> Any:
    if isinstance(value, str):
        return normalize_budget_text(value)
    if isinstance(value, Mapping):
        return {str(key): _normalize_cache_value(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))}
    if isinstance(value, (list, tuple)):
        return [_normalize_cache_value(item) for item in value]
    if isinstance(value, datetime):
        return value.astimezone(UTC).isoformat()
    return value


def build_llm_task_fingerprint(
    *,
    task: str,
    provider: str,
    model: str,
    cache_version: str,
    prompt_version: str,
    schema_version: str | None = None,
    payload: Any = None,
    metadata: Mapping[str, Any] | None = None,
) -> str:
    canonical: dict[str, Any] = {
        "task": str(task or "").strip(),
        "cache_version": str(cache_version or "").strip(),
        "provider": str(provider or "").strip(),
        "model": str(model or "").strip(),
        "prompt_version": str(prompt_version or "").strip(),
        "payload": _normalize_cache_value(payload),
        "metadata": _normalize_cache_value(dict(metadata or {})),
    }
    if schema_version:
        canonical["schema_version"] = str(schema_version).strip()
    serialized = json.dumps(canonical, ensure_ascii=False, sort_keys=True, default=str, separators=(",", ":"))
    return hashlib.sha256(serialized.encode("utf-8")).hexdigest()


class LLMTaskCacheRepository(Protocol):
    def ensure_indexes(self) -> None: ...

    def get(
        self,
        *,
        username: str,
        task: str,
        fingerprint: str,
    ) -> dict[str, Any] | None: ...

    def upsert(
        self,
        *,
        username: str,
        task: str,
        fingerprint: str,
        provider: str,
        model: str,
        payload: Any,
        metadata: Mapping[str, Any] | None = None,
    ) -> None: ...


class MongoLLMTaskCacheRepository:
    def __init__(
        self,
        mongo_cache: Any,
        *,
        retention_days: int | None = None,
        cache_version: str | None = None,
        enabled: bool | None = None,
    ) -> None:
        self.collection = getattr(mongo_cache, "collection", mongo_cache)
        self.retention_days = int(retention_days or config.LLM_CACHE_RETENTION_DAYS)
        self.cache_version = str(cache_version or config.LLM_CACHE_VERSION)
        self.enabled = config.LLM_CACHE_ENABLED if enabled is None else bool(enabled)

    def ensure_indexes(self) -> None:
        collection = self.collection
        if not self.enabled or collection is None or not hasattr(collection, "create_index"):
            return
        try:
            collection.create_index(
                [("usuario", 1), ("task", 1), ("fingerprint", 1)],
                unique=True,
                partialFilterExpression={
                    "tipo_documento": "llm_task_cache",
                    "usuario": {"$exists": True, "$type": "string"},
                    "task": {"$exists": True, "$type": "string"},
                    "fingerprint": {"$exists": True, "$type": "string"},
                },
            )
            collection.create_index([("expires_at", 1)], expireAfterSeconds=0)
            collection.create_index([("created_at", -1)])
        except OperationFailure as exc:
            if getattr(exc, "code", None) not in {85, 86}:
                raise

    def get(
        self,
        *,
        username: str,
        task: str,
        fingerprint: str,
    ) -> dict[str, Any] | None:
        collection = self.collection
        normalized_username = str(username or "").strip()
        if (
            not self.enabled
            or not normalized_username
            or collection is None
            or not hasattr(collection, "find_one")
        ):
            return None
        now = _now_utc()
        cache = collection.find_one(
            {
                "tipo_documento": "llm_task_cache",
                "usuario": normalized_username,
                "task": str(task or "").strip(),
                "fingerprint": str(fingerprint or "").strip(),
                "$or": [{"expires_at": {"$gt": now}}, {"expires_at": {"$exists": False}}],
            }
        )
        if not cache:
            return None
        if hasattr(collection, "update_one") and cache.get("_id") is not None:
            collection.update_one({"_id": cache["_id"]}, {"$inc": {"hit_count": 1}})
        return dict(cache)

    def upsert(
        self,
        *,
        username: str,
        task: str,
        fingerprint: str,
        provider: str,
        model: str,
        payload: Any,
        metadata: Mapping[str, Any] | None = None,
    ) -> None:
        collection = self.collection
        normalized_username = str(username or "").strip()
        if (
            not self.enabled
            or not normalized_username
            or collection is None
            or not hasattr(collection, "update_one")
        ):
            return
        now = _now_utc()
        document = {
            "tipo_documento": "llm_task_cache",
            "usuario": normalized_username,
            "task": str(task or "").strip(),
            "fingerprint": str(fingerprint or "").strip(),
            "cache_version": self.cache_version,
            "provider": str(provider or "").strip(),
            "model": str(model or "").strip(),
            "payload": _normalize_cache_value(payload),
            "metadata": _normalize_cache_value(dict(metadata or {})),
            "created_at": now,
            "expires_at": now + timedelta(days=max(1, self.retention_days)),
        }
        collection.update_one(
            {
                "tipo_documento": "llm_task_cache",
                "usuario": normalized_username,
                "task": str(task or "").strip(),
                "fingerprint": str(fingerprint or "").strip(),
            },
            {
                "$set": document,
                "$setOnInsert": {"hit_count": 0},
            },
            upsert=True,
        )