"""Usage metering and tier quota enforcement (P11).

Four metered dimensions: tokens per month, messages per day,
concurrent tasks, storage bytes. Tokens are read from the existing
``cost_records`` ledger (``cloud/cost.py``) -- the inference channel's
bookkeeping point -- never re-collected on a second path. Messages,
task starts/ends, storage snapshots and warnings are recorded in the
hash-chained ``usage_ledger`` table in this module.

Tier limits and prices are a PROJECT PROPOSAL, not user-confirmed
pricing and not Meta official pricing: the user delegated the numbers
("set them by market", 2026-11-06) and never confirmed them. Every
payload this module returns carries ``proposal: True``. Real billing
(Stripe / payments / invoicing) is NOT in this module and must be
approved separately; nothing here charges money.

Admission semantics: gates are admission-only. A denial returns a
``quota_exceeded`` dict (HTTP 402 at the server layer); there is no
API here that terminates an in-flight turn. The pre-existing USD
quota path in ``cloud/server.py`` still cuts in-flight work at turn
boundaries via worker polling -- that is a known, documented
discrepancy (see docs/chv-muse-p11-gap.md §5), not something this
module silently resolves.

Token gate counts ``source="actual"`` rows only, matching the
``cost.py`` P1 decision ("estimated rows are never real spend and
must not trip the cap"); estimated tokens are reported separately
and labelled as estimates.
"""

import hashlib
import json
import os
import sqlite3
import time
from datetime import datetime, timezone

from cloud.cost import CostStore
from cloud.db import connect as _db_connect

# ---------------------------------------------------------------------------
# Tier proposal (see module docstring: PROPOSAL, not confirmed pricing).
# ---------------------------------------------------------------------------

GB = 1024 * 1024 * 1024

TIER_NAMES = ("free", "paid", "heavy")

TIER_LIMITS = {
    "free": {
        "label": "Free",
        "price_usd_per_month": 0,
        "messages_per_day": 50,
        "tokens_per_month": 200000,
        "concurrent_tasks": 1,
        "storage_bytes": 1 * GB,
    },
    "paid": {
        "label": "Paid",
        "price_usd_per_month": 19,
        "messages_per_day": 500,
        "tokens_per_month": 2000000,
        "concurrent_tasks": 5,
        "storage_bytes": 10 * GB,
    },
    "heavy": {
        "label": "Heavy",
        "price_usd_per_month": 99,
        # The proposal never gave Heavy a messages/day number;
        # None = no daily message cap (project decision, not proposal).
        "messages_per_day": None,
        "tokens_per_month": 20000000,
        "concurrent_tasks": 20,
        "storage_bytes": 100 * GB,
    },
}

_LEDGER_KINDS = ("message", "task_start", "task_end", "storage",
                 "warning", "usage_query")


def warn_ratio() -> float:
    """Warning threshold as a fraction of a limit (default 0.8).

    Env ``CLOUD_USAGE_WARN_RATIO`` overrides; unparsable or
    out-of-range values fall back to 0.8.
    """
    try:
        value = float(os.environ.get("CLOUD_USAGE_WARN_RATIO", "0.8"))
    except (TypeError, ValueError):
        return 0.8
    if value <= 0 or value > 1:
        return 0.8
    return value


def utc_day(now: float | None = None) -> str:
    """UTC calendar day (YYYY-MM-DD) for ``now`` (default: current)."""
    stamp = now if now is not None else time.time()
    return datetime.fromtimestamp(stamp, tz=timezone.utc).strftime("%Y-%m-%d")


def utc_month(now: float | None = None) -> str:
    """UTC calendar month (YYYY-MM) for ``now`` (default: current)."""
    stamp = now if now is not None else time.time()
    return datetime.fromtimestamp(stamp, tz=timezone.utc).strftime("%Y-%m")


def disk_usage_bytes(path: str) -> int:
    """Sum of file sizes under ``path`` (best effort, bytes).

    Does not follow symlinks; unreadable files/dirs are skipped.
    Missing path counts as 0.
    """
    total = 0
    if not path or not os.path.exists(path):
        return 0
    if os.path.isfile(path) and not os.path.islink(path):
        try:
            return os.path.getsize(path)
        except OSError:
            return 0
    for root, dirs, files in os.walk(path, followlinks=False):
        for name in files:
            full = os.path.join(root, name)
            try:
                if not os.path.islink(full):
                    total += os.path.getsize(full)
            except OSError:
                continue
    return total


def _entry_hash(prev_hash: str | None, fields: dict) -> str:
    """Hash-chain digest over the previous hash + canonical fields."""
    payload = json.dumps({"prev": prev_hash or "", "fields": fields},
                         sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


class TierStore:
    """Per-user tier assignment (``user_tiers`` table).

    No row means the user is NOT tier-gated: only the pre-existing
    USD quota applies. Assigning a tier is an explicit admin action;
    nothing here defaults a user to Free.
    """

    _DDL = """CREATE TABLE IF NOT EXISTS user_tiers (
                user_id TEXT PRIMARY KEY,
                tier TEXT,
                updated_at DOUBLE PRECISION
            )"""

    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        self._conn = _db_connect(db_path)
        self._conn.execute(self._DDL)
        self._conn.commit()

    def get_tier(self, user_id: str) -> str | None:
        """Tier name for ``user_id``, or None when unassigned."""
        row = self._conn.execute(
            "SELECT tier FROM user_tiers WHERE user_id = ?",
            (user_id,)).fetchone()
        if row is None:
            return None
        return row[0]

    def set_tier(self, user_id: str, tier: str) -> None:
        """Assign ``tier`` to ``user_id``; unknown tier names rejected."""
        if tier not in TIER_NAMES:
            raise ValueError("unknown tier %r (want one of %r)"
                             % (tier, TIER_NAMES))
        self._conn.execute(
            "INSERT INTO user_tiers (user_id, tier, updated_at)"
            " VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE"
            " SET tier = excluded.tier,"
            " updated_at = excluded.updated_at",
            (user_id, tier, time.time()))
        self._conn.commit()

    def clear_tier(self, user_id: str) -> None:
        """Remove the tier row: the user is tier-gated no more."""
        self._conn.execute(
            "DELETE FROM user_tiers WHERE user_id = ?", (user_id,))
        self._conn.commit()

    def limits_for(self, user_id: str) -> dict | None:
        """Tier limits dict for ``user_id``, or None when unassigned."""
        tier = self.get_tier(user_id)
        if tier is None:
            return None
        return dict(TIER_LIMITS[tier])


class UsageLedger:
    """Append-only, hash-chained usage ledger (``usage_ledger``).

    One chain per user: each row stores the previous row's hash and
    its own digest over the previous hash + canonical fields. Any
    UPDATE/DELETE breaks ``verify_chain``. The ledger never stores
    conversation text, file contents, or credential values -- only
    counts, ids, and hashes.
    """

    _DDL = """CREATE TABLE IF NOT EXISTS usage_ledger (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT,
                kind TEXT,
                amount INTEGER,
                unit TEXT,
                task_id TEXT,
                model TEXT,
                source TEXT,
                day TEXT,
                created_at REAL,
                prev_hash TEXT,
                entry_hash TEXT
            )"""

    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        if db_path != ":memory:":
            parent = os.path.dirname(os.path.abspath(db_path))
            os.makedirs(parent, exist_ok=True)
        self._conn = _db_connect(db_path)
        self._conn.execute(self._DDL)
        self._conn.commit()

    def _new_conn(self) -> sqlite3.Connection:
        """Fresh connection per write/read (mirrors CostStore).

        ``:memory:`` reuses the shared connection; anything else
        opens a new one so concurrent threads never share a cursor.
        """
        if self.db_path == ":memory:":
            return self._conn
        conn = _db_connect(self.db_path)
        conn.execute(self._DDL)
        conn.commit()
        return conn

    def record(self, user_id: str, kind: str, amount: int,
               unit: str = "", task_id: str | None = None,
               model: str | None = None, source: str | None = None,
               day: str | None = None) -> int:
        """Append one ledger event; returns its row id.

        ``kind`` must be one of the known event kinds; ``amount`` is
        a count (1 per message, bytes for storage snapshots).
        """
        if kind not in _LEDGER_KINDS:
            raise ValueError("unknown ledger kind %r" % (kind,))
        day = day or utc_day()
        created_at = time.time()
        conn = self._new_conn()
        try:
            prev = conn.execute(
                "SELECT entry_hash FROM usage_ledger"
                " WHERE user_id = ? ORDER BY id DESC LIMIT 1",
                (user_id,)).fetchone()
            prev_hash = prev[0] if prev is not None else None
            fields = {"user_id": user_id, "kind": kind,
                      "amount": int(amount), "unit": unit,
                      "task_id": task_id or "", "model": model or "",
                      "source": source or "", "day": day,
                      "created_at": created_at}
            digest = _entry_hash(prev_hash, fields)
            cursor = conn.execute(
                """INSERT INTO usage_ledger
                   (user_id, kind, amount, unit, task_id, model,
                    source, day, created_at, prev_hash, entry_hash)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (user_id, kind, int(amount), unit, task_id, model,
                 source, day, created_at, prev_hash, digest))
            conn.commit()
            return int(cursor.lastrowid or 0)
        finally:
            if conn is not self._conn:
                conn.close()

    def sum_kind(self, user_id: str, kind: str,
                 day: str | None = None,
                 month: str | None = None) -> int:
        """Sum ``amount`` for ``user_id``/``kind`` in a day or month."""
        conn = self._new_conn()
        try:
            if day is not None:
                row = conn.execute(
                    "SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
                    " WHERE user_id = ? AND kind = ? AND day = ?",
                    (user_id, kind, day)).fetchone()
            elif month is not None:
                row = conn.execute(
                    "SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
                    " WHERE user_id = ? AND kind = ? AND day LIKE ?",
                    (user_id, kind, month + "%")).fetchone()
            else:
                row = conn.execute(
                    "SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
                    " WHERE user_id = ? AND kind = ?",
                    (user_id, kind)).fetchone()
            return int(row[0] or 0)
        finally:
            if conn is not self._conn:
                conn.close()

    def verify_chain(self, user_id: str) -> tuple[bool, int | None]:
        """Recompute the hash chain; (True, None) or (False, bad id)."""
        conn = self._new_conn()
        try:
            rows = conn.execute(
                "SELECT id, user_id, kind, amount, unit, task_id,"
                " model, source, day, created_at, prev_hash,"
                " entry_hash FROM usage_ledger"
                " WHERE user_id = ? ORDER BY id",
                (user_id,)).fetchall()
        finally:
            if conn is not self._conn:
                conn.close()
        prev_hash: str | None = None
        for row in rows:
            fields = {"user_id": row[1], "kind": row[2],
                      "amount": int(row[3]), "unit": row[4],
                      "task_id": row[5] or "", "model": row[6] or "",
                      "source": row[7] or "", "day": row[8],
                      "created_at": row[9]}
            if row[10] != prev_hash:
                return False, int(row[0])
            if _entry_hash(prev_hash, fields) != row[11]:
                return False, int(row[0])
            prev_hash = row[11]
        return True, None

    def warning_recorded(self, user_id: str, dimension: str,
                         window: str) -> bool:
        """True when a warning for (dimension, window) is on record."""
        conn = self._new_conn()
        try:
            row = conn.execute(
                "SELECT COUNT(*) FROM usage_ledger"
                " WHERE user_id = ? AND kind = 'warning'"
                " AND unit = ? AND day = ?",
                (user_id, dimension, window)).fetchone()
            return int(row[0] or 0) > 0
        finally:
            if conn is not self._conn:
                conn.close()

    def record_warning(self, user_id: str, dimension: str,
                       window: str) -> int:
        """Record a once-per-window warning event; returns row id."""
        return self.record(user_id, "warning", 1, unit=dimension,
                           day=window)


class UsageMeter:
    """Four-dimension usage metering for one database.

    Tokens come from ``cost_records`` (the inference channel's
    bookkeeping point, joined to ``tasks`` exactly like
    ``_quota_status`` does); messages come from the ledger;
    concurrency and storage come from injected measurers so the
    meter stays unit-testable without a server or a disk layout.
    An unwired measurer reports ``None`` (unmeasured) in summaries
    and raises in gates -- never silently passes as zero.
    """

    def __init__(self, db_path: str,
                 running_task_count=None,
                 storage_bytes=None) -> None:
        self.db_path = db_path
        self.tiers = TierStore(db_path)
        self.ledger = UsageLedger(db_path)
        self._running_task_count = running_task_count
        self._storage_bytes = storage_bytes

    def monthly_tokens(self, user_id: str,
                       month: str | None = None) -> dict:
        """Actual / estimated / total tokens for the UTC month.

        ``actual`` is provider-reported usage; ``estimated`` is the
        character-count fallback, labelled as an estimate. Only
        ``actual`` ever trips a gate.
        """
        month = month or utc_month()
        _cost_store = CostStore(self.db_path)
        try:
            _cost_store._conn.close()
        except Exception:
            pass
        conn = _db_connect(self.db_path)
        try:
            rows = conn.execute(
                "SELECT c.source,"
                " COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
                " FROM cost_records c JOIN tasks t ON c.task_id = t.id"
                " WHERE t.created_by = ? AND c.day LIKE ?"
                " GROUP BY c.source",
                (user_id, month + "%")).fetchall()
        finally:
            conn.close()
        totals = {"actual": 0, "estimated": 0}
        for source, total in rows:
            if source in totals:
                totals[source] = int(total or 0)
        totals["total"] = totals["actual"] + totals["estimated"]
        return totals

    def daily_messages(self, user_id: str,
                       day: str | None = None) -> int:
        """Messages recorded for the UTC day (fixed window)."""
        return self.ledger.sum_kind(user_id, "message",
                                    day=day or utc_day())

    def record_message(self, user_id: str,
                       task_id: str | None = None) -> int:
        """Record one message event; returns the ledger row id."""
        return self.ledger.record(user_id, "message", 1,
                                  unit="message", task_id=task_id)

    def record_task_start(self, user_id: str,
                          task_id: str | None = None) -> int:
        """Record a task-start audit event; returns row id."""
        return self.ledger.record(user_id, "task_start", 1,
                                  unit="task", task_id=task_id)

    def record_task_end(self, user_id: str,
                        task_id: str | None = None) -> int:
        """Record a task-end audit event; returns row id."""
        return self.ledger.record(user_id, "task_end", 1,
                                  unit="task", task_id=task_id)

    def record_storage_snapshot(self, user_id: str,
                                size_bytes: int) -> int:
        """Record a measured storage snapshot; returns row id."""
        return self.ledger.record(user_id, "storage", int(size_bytes),
                                  unit="bytes")

    def running_tasks(self, user_id: str) -> int | None:
        """Concurrent running tasks, or None when unwired."""
        if self._running_task_count is None:
            return None
        return int(self._running_task_count(user_id))

    def storage_used(self, user_id: str) -> int | None:
        """Storage bytes, or None when unwired."""
        if self._storage_bytes is None:
            return None
        return int(self._storage_bytes(user_id))

    def summary(self, user_id: str) -> dict:
        """Current-period usage, limits, remaining, warnings.

        Warnings are recorded once per dimension per window (month
        for tokens, day for the rest) and included in the payload.
        """
        month = utc_month()
        day = utc_day()
        limits = self.tiers.limits_for(user_id)
        tier = self.tiers.get_tier(user_id)
        tokens = self.monthly_tokens(user_id, month)
        messages = self.daily_messages(user_id, day)
        running = self.running_tasks(user_id)
        storage = self.storage_used(user_id)
        ledger_ok, ledger_bad = self.ledger.verify_chain(user_id)
        used = {"tokens_actual": tokens["actual"],
                "tokens_estimated": tokens["estimated"],
                "messages": messages,
                "concurrent_tasks": running,
                "storage_bytes": storage}
        payload: dict = {
            "proposal": True,
            "user_id": user_id,
            "tier": tier,
            "month": month,
            "day": day,
            "used": used,
            "limits": limits,
            "remaining": _remaining(limits, used),
            "ledger_verified": ledger_ok,
            "ledger_first_bad_id": ledger_bad,
            "warnings": [],
        }
        if limits is not None:
            payload["warnings"] = self._check_warnings(
                user_id, limits, used, month, day)
        return payload

    def _check_warnings(self, user_id: str, limits: dict,
                        used: dict, month: str, day: str) -> list:
        """80%-of-limit warnings, once per dimension per window."""
        ratio = warn_ratio()
        warnings = []
        checks = (("tokens", limits.get("tokens_per_month"),
                   used["tokens_actual"], month),
                  ("messages", limits.get("messages_per_day"),
                   used["messages"], day),
                  ("concurrent_tasks", limits.get("concurrent_tasks"),
                   used["concurrent_tasks"], day),
                  ("storage", limits.get("storage_bytes"),
                   used["storage_bytes"], day))
        for dimension, limit, value, window in checks:
            if limit is None or value is None:
                continue
            if limit > 0 and value >= limit * ratio:
                if not self.ledger.warning_recorded(
                        user_id, dimension, window):
                    self.ledger.record_warning(user_id, dimension,
                                               window)
                warnings.append({"dimension": dimension,
                                 "limit": limit, "used": value,
                                 "ratio": ratio})
        return warnings


def _remaining(limits: dict | None, used: dict) -> dict:
    """Remaining quota per dimension (None = unlimited/unmeasured)."""
    if limits is None:
        return {"tokens": None, "messages": None,
                "concurrent_tasks": None, "storage_bytes": None}
    pairs = (("tokens", limits.get("tokens_per_month"),
              used["tokens_actual"]),
             ("messages", limits.get("messages_per_day"),
              used["messages"]),
             ("concurrent_tasks", limits.get("concurrent_tasks"),
              used["concurrent_tasks"]),
             ("storage_bytes", limits.get("storage_bytes"),
              used["storage_bytes"]))
    remaining: dict[str, int | None] = {}
    for name, limit, value in pairs:
        if limit is None or value is None:
            remaining[name] = None
        else:
            remaining[name] = max(0, limit - value)
    return remaining


def _denial(dimension: str, limit: int, used: int,
            detail: str) -> dict:
    """One quota denial payload (HTTP 402 at the server layer)."""
    return {"error": "quota_exceeded", "dimension": dimension,
            "limit": limit, "used": used, "detail": detail}


class QuotaGate:
    """Admission-only tier gates (pure decisions, no side effects).

    Every method returns None (allow) or a denial dict. Nothing here
    touches a running task: an in-flight turn is never cut off by
    these gates; a denial only blocks the NEXT admission.
    """

    def __init__(self, meter: UsageMeter) -> None:
        self.meter = meter

    def check_message_admission(self, user_id: str) -> dict | None:
        """Gate a new message: messages/day, then tokens/month."""
        limits = self.meter.tiers.limits_for(user_id)
        if limits is None:
            return None
        day_limit = limits.get("messages_per_day")
        if day_limit is not None:
            used = self.meter.daily_messages(user_id)
            if used >= day_limit:
                return _denial(
                    "messages", day_limit, used,
                    "daily message quota exceeded for tier %r"
                    " (used %d of %d); wait until tomorrow (UTC)"
                    " or upgrade the tier"
                    % (self.meter.tiers.get_tier(user_id), used,
                       day_limit))
        return self._check_tokens(user_id, limits)

    def check_task_admission(
            self, user_id: str, running_tasks: int | None = None,
            storage_bytes: int | None = None) -> dict | None:
        """Gate a new task: tokens, then concurrency, then storage."""
        limits = self.meter.tiers.limits_for(user_id)
        if limits is None:
            return None
        denied = self._check_tokens(user_id, limits)
        if denied is not None:
            return denied
        running = running_tasks
        if running is None:
            running = self.meter.running_tasks(user_id)
        if running is None:
            raise ValueError("no running-task measurer wired")
        task_limit = limits.get("concurrent_tasks")
        if task_limit is not None and running >= task_limit:
            return _denial(
                "concurrent_tasks", task_limit, running,
                "concurrent task quota exceeded for tier %r"
                " (running %d of %d); wait for a running task"
                " to finish" % (self.meter.tiers.get_tier(user_id),
                                running, task_limit))
        storage = storage_bytes
        if storage is None:
            storage = self.meter.storage_used(user_id)
        if storage is None:
            raise ValueError("no storage measurer wired")
        storage_limit = limits.get("storage_bytes")
        if storage_limit is not None and storage >= storage_limit:
            return _denial(
                "storage", storage_limit, storage,
                "storage quota exceeded for tier %r"
                " (used %d of %d bytes); free space by removing"
                " old task worktrees or snapshots, then retry"
                % (self.meter.tiers.get_tier(user_id), storage,
                   storage_limit))
        return None

    def _check_tokens(self, user_id: str,
                      limits: dict) -> dict | None:
        """Monthly actual-token gate (estimated tokens never gate)."""
        token_limit = limits.get("tokens_per_month")
        if token_limit is None:
            return None
        used = self.meter.monthly_tokens(user_id)["actual"]
        if used >= token_limit:
            return _denial(
                "tokens", token_limit, used,
                "monthly token quota exceeded for tier %r"
                " (used %d of %d actual tokens; estimated tokens"
                " are not counted); wait until next month (UTC)"
                " or upgrade the tier"
                % (self.meter.tiers.get_tier(user_id), used,
                   token_limit))
        return None