from __future__ import annotations

import hashlib
import json
import os
import secrets
import threading
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError


ANON_DAILY_LIMIT = int(os.environ.get("KIMODO_ANON_DAILY_LIMIT", "3"))
ANON_COOLDOWN_SECONDS = int(os.environ.get("KIMODO_ANON_COOLDOWN_SECONDS", "0"))
USER_DAILY_LIMIT = int(os.environ.get("KIMODO_USER_DAILY_LIMIT", "20"))
# Signed-in users are limited only by the (higher) daily cap -- no per-generation
# cooldown, so logging in never introduces a wait that anonymous browsing lacked.
USER_COOLDOWN_SECONDS = int(os.environ.get("KIMODO_USER_COOLDOWN_SECONDS", "0"))
QUOTA_PATH = "quotas/generation_quotas.json"


@dataclass(frozen=True)
class QuotaBucket:
    key: str
    kind: str


@dataclass(frozen=True)
class QuotaResult:
    allowed: bool
    message: str = ""


@dataclass(frozen=True)
class QuotaStatus:
    kind: str               # "hf_user" or "anonymous"
    used: int               # generations consumed today
    limit: int              # daily cap
    remaining: int          # limit - used (>= 0)
    cooldown_remaining: int  # seconds until the next generation is allowed (0 = ready)


def _utc_day(now: float | None = None) -> str:
    return datetime.fromtimestamp(now or time.time(), timezone.utc).strftime("%Y-%m-%d")


def _seconds_until_utc_midnight(now: float) -> float:
    dt = datetime.fromtimestamp(now, timezone.utc)
    nxt = (dt + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)
    return max(0.0, (nxt - dt).total_seconds())


def _format_remaining(seconds: float) -> str:
    """A relative 'time remaining' string (e.g. '1m 30s', '2h 5m'). Relative time
    is timezone-agnostic, so the user reads the same correct value everywhere --
    unlike an absolute clock time, which would be in the server's (UTC) zone."""
    total = int(max(0, round(seconds)))
    hours, rem = divmod(total, 3600)
    minutes, secs = divmod(rem, 60)
    parts = []
    if hours:
        parts.append(f"{hours}h")
    if minutes:
        parts.append(f"{minutes}m")
    if secs and not hours:
        parts.append(f"{secs}s")
    return " ".join(parts) or "0s"


def _stable_salt() -> str:
    salt = os.environ.get("KIMODO_QUOTA_SALT", "").strip()
    if salt:
        return salt
    local = Path(os.environ.get("KIMODO_STORE_PATH", ".kimodo-animations")) / ".quota_salt"
    try:
        local.parent.mkdir(parents=True, exist_ok=True)
        if local.exists():
            existing = local.read_text(encoding="utf-8").strip()
            if existing:
                return existing
        generated = secrets.token_hex(32)
        local.write_text(generated, encoding="utf-8")
        return generated
    except Exception:
        return "kimodo-local-dev-quota-salt"


def _hash_identifier(raw: str) -> str:
    material = f"{_stable_salt()}:{raw}".encode("utf-8", errors="ignore")
    return hashlib.sha256(material).hexdigest()


def _request_headers(request: Any) -> dict[str, str]:
    if request is None:
        return {}
    headers = getattr(request, "headers", None) or {}
    try:
        return {str(k).lower(): str(v) for k, v in dict(headers).items()}
    except Exception:
        return {}


def bucket_for_request(
    creator: dict[str, Any] | None = None,
    request: Any = None,
    anonymous_client_id: str | None = None,
) -> QuotaBucket:
    username = (creator or {}).get("created_by")
    if username:
        return QuotaBucket(f"hf:{username}", "hf_user")

    headers = _request_headers(request)
    raw_identifier = str(anonymous_client_id or "").strip()
    if raw_identifier:
        raw_identifier = f"browser:{raw_identifier}"
    if not raw_identifier:
        raw_identifier = headers.get("x-ip-token") or ""
    if not raw_identifier:
        client = getattr(request, "client", None) if request is not None else None
        raw_identifier = getattr(client, "host", "") or ""
    if not raw_identifier and request is not None:
        raw_identifier = getattr(request, "session_hash", "") or ""
    if not raw_identifier:
        raw_identifier = "unknown-anonymous"
    return QuotaBucket(f"anon:{_hash_identifier(str(raw_identifier))}", "anonymous")


class GenerationQuotaStore:
    def __init__(self, path: str | Path | None = None) -> None:
        self._lock = threading.Lock()
        self.path = Path(path or os.environ.get("KIMODO_QUOTA_PATH", ".kimodo-animations/quotas.json"))

    def _load(self) -> dict[str, Any]:
        if not self.path.exists():
            return {"records": {}}
        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
            return data if isinstance(data, dict) else {"records": {}}
        except Exception:
            return {"records": {}}

    def _save(self, data: dict[str, Any]) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(json.dumps(data, separators=(",", ":")), encoding="utf-8")
        tmp.replace(self.path)

    def _limits(self, bucket: QuotaBucket) -> tuple[int, int]:
        if bucket.kind == "hf_user":
            return USER_DAILY_LIMIT, USER_COOLDOWN_SECONDS
        return ANON_DAILY_LIMIT, ANON_COOLDOWN_SECONDS

    def status(self, bucket: QuotaBucket, now: float | None = None) -> QuotaStatus:
        """Read-only snapshot of a bucket's quota for display (never consumes)."""
        now = now or time.time()
        day = _utc_day(now)
        daily_limit, cooldown = self._limits(bucket)
        with self._lock:
            records = self._load().get("records", {})
            record = records.get(bucket.key, {})
            if record.get("day") != day:
                used, last_at = 0, 0.0
            else:
                used = int(record.get("count", 0) or 0)
                last_at = float(record.get("last_generation_at", 0) or 0)
        cooldown_remaining = 0
        if cooldown > 0 and last_at and now - last_at < cooldown:
            cooldown_remaining = int(round(last_at + cooldown - now))
        return QuotaStatus(bucket.kind, used, daily_limit, max(0, daily_limit - used), cooldown_remaining)

    def check(self, bucket: QuotaBucket, now: float | None = None) -> QuotaResult:
        now = now or time.time()
        day = _utc_day(now)
        daily_limit, cooldown = self._limits(bucket)
        with self._lock:
            records = self._load().setdefault("records", {})
            record = records.get(bucket.key, {})
            if record.get("day") != day:
                record = {"day": day, "bucket": bucket.key, "kind": bucket.kind, "count": 0}
            count = int(record.get("count", 0) or 0)
            last_at = float(record.get("last_generation_at", 0) or 0)
            if count >= daily_limit:
                who = "User" if bucket.kind == "hf_user" else "Anonymous"
                sign_in = " Sign in to continue." if bucket.kind != "hf_user" else ""
                remaining = _format_remaining(_seconds_until_utc_midnight(now))
                return QuotaResult(False, f"{who} generation limit reached for today. Resets in {remaining}.{sign_in}")
            if cooldown > 0 and last_at and now - last_at < cooldown:
                reset_at = last_at + cooldown
                who = "User" if bucket.kind == "hf_user" else "Anonymous"
                sign_in = " Sign in to continue with a higher limit." if bucket.kind != "hf_user" else ""
                return QuotaResult(False, f"{who} generation cooldown active. Try again in {_format_remaining(reset_at - now)}.{sign_in}")
            return QuotaResult(True)

    def check_and_consume(self, bucket: QuotaBucket, now: float | None = None) -> QuotaResult:
        now = now or time.time()
        day = _utc_day(now)
        daily_limit, cooldown = self._limits(bucket)
        with self._lock:
            data = self._load()
            records = data.setdefault("records", {})
            record = records.get(bucket.key, {})
            if record.get("day") != day:
                record = {"day": day, "bucket": bucket.key, "kind": bucket.kind, "count": 0}
            count = int(record.get("count", 0) or 0)
            last_at = float(record.get("last_generation_at", 0) or 0)
            if count >= daily_limit:
                who = "User" if bucket.kind == "hf_user" else "Anonymous"
                sign_in = " Sign in to continue." if bucket.kind != "hf_user" else ""
                remaining = _format_remaining(_seconds_until_utc_midnight(now))
                return QuotaResult(False, f"{who} generation limit reached for today. Resets in {remaining}.{sign_in}")
            if cooldown > 0 and last_at and now - last_at < cooldown:
                reset_at = last_at + cooldown
                who = "User" if bucket.kind == "hf_user" else "Anonymous"
                sign_in = " Sign in to continue with a higher limit." if bucket.kind != "hf_user" else ""
                return QuotaResult(False, f"{who} generation cooldown active. Try again in {_format_remaining(reset_at - now)}.{sign_in}")
            record["count"] = count + 1
            record["last_generation_at"] = now
            records[bucket.key] = record
            self._save(data)
            return QuotaResult(True)

    def consume(self, bucket: QuotaBucket, now: float | None = None) -> None:
        now = now or time.time()
        day = _utc_day(now)
        with self._lock:
            data = self._load()
            records = data.setdefault("records", {})
            record = records.get(bucket.key, {})
            if record.get("day") != day:
                record = {"day": day, "bucket": bucket.key, "kind": bucket.kind, "count": 0}
            record["count"] = int(record.get("count", 0) or 0) + 1
            record["last_generation_at"] = now
            records[bucket.key] = record
            self._save(data)


class HFDatasetGenerationQuotaStore(GenerationQuotaStore):
    def __init__(self, repo_id: str, token: str | None = None) -> None:
        super().__init__()
        self.repo_id = repo_id
        self.token = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        self.api = HfApi(token=self.token)
        self.api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, private=False)

    def _load(self) -> dict[str, Any]:
        try:
            path = hf_hub_download(
                repo_id=self.repo_id,
                repo_type="dataset",
                filename=QUOTA_PATH,
                token=self.token,
            )
            data = json.loads(Path(path).read_text(encoding="utf-8"))
            return data if isinstance(data, dict) else {"records": {}}
        except (EntryNotFoundError, RepositoryNotFoundError, FileNotFoundError):
            return {"records": {}}
        except Exception as exc:
            print(f"quota load fell back to local cache: {type(exc).__name__}: {exc}")
            return super()._load()

    def _save(self, data: dict[str, Any]) -> None:
        try:
            blob = json.dumps(data, separators=(",", ":")).encode("utf-8")
            self.api.create_commit(
                repo_id=self.repo_id,
                repo_type="dataset",
                operations=[CommitOperationAdd(path_in_repo=QUOTA_PATH, path_or_fileobj=blob)],
                commit_message="Update generation quotas",
            )
        except Exception as exc:
            print(f"quota dataset save failed; writing local cache: {type(exc).__name__}: {exc}")
            super()._save(data)


def make_quota_store() -> GenerationQuotaStore:
    backend = os.environ.get("KIMODO_STORE_BACKEND", "").strip().lower()
    dataset_repo = os.environ.get("KIMODO_DATASET_REPO", "").strip()
    if backend == "dataset" or dataset_repo:
        return HFDatasetGenerationQuotaStore(dataset_repo or "polats/kimodo-kata-animations")
    return GenerationQuotaStore()


quota_store = make_quota_store()