"""Durable module action budget counters."""

from __future__ import annotations

import sqlite3
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator


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


@dataclass(frozen=True)
class ModuleActionBudgetDecision:
    allowed: bool
    reason: str | None
    day_key: str
    daily_used: int
    entity_used: int


class ModuleActionBudgetStore:
    """SQLite-backed per-module global and per-entity action budgets."""

    _GLOBAL_SCOPE = "global"
    _ENTITY_SCOPE = "entity"
    _GLOBAL_KEY = "__all__"

    def __init__(self, db_path: Path | str) -> None:
        self.db_path = Path(db_path).resolve()
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._ensure_schema()

    @contextmanager
    def _connect(self) -> Iterator[sqlite3.Connection]:
        conn = sqlite3.connect(self.db_path, timeout=10, isolation_level=None)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA busy_timeout=10000")
        try:
            yield conn
        finally:
            conn.close()

    def _ensure_schema(self) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS module_action_budgets (
                    module_id TEXT NOT NULL,
                    day_key TEXT NOT NULL,
                    scope TEXT NOT NULL,
                    scope_key TEXT NOT NULL,
                    count INTEGER NOT NULL,
                    updated_at TEXT NOT NULL,
                    PRIMARY KEY(module_id, day_key, scope, scope_key)
                )
                """
            )
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_module_action_budgets_scope
                ON module_action_budgets(module_id, day_key, scope)
                """
            )

    def consume(
        self,
        *,
        module_id: str,
        day_key: str,
        entity_key: str,
        daily_limit: int,
        entity_limit: int,
    ) -> ModuleActionBudgetDecision:
        normalized_module_id = str(module_id or "").strip().lower()
        if not normalized_module_id:
            raise ValueError("module_id is required")
        normalized_entity = str(entity_key or "").strip().lower() or "entity:unknown"

        with self._connect() as conn:
            conn.execute("BEGIN IMMEDIATE")
            try:
                self._cleanup_stale_days(conn, normalized_module_id, day_key)
                daily_used = self._get_count(
                    conn,
                    normalized_module_id,
                    day_key,
                    self._GLOBAL_SCOPE,
                    self._GLOBAL_KEY,
                )
                entity_used = self._get_count(
                    conn,
                    normalized_module_id,
                    day_key,
                    self._ENTITY_SCOPE,
                    normalized_entity,
                )

                if daily_used >= daily_limit:
                    conn.execute("COMMIT")
                    return ModuleActionBudgetDecision(
                        allowed=False,
                        reason="daily_budget_exhausted",
                        day_key=day_key,
                        daily_used=daily_used,
                        entity_used=entity_used,
                    )
                if entity_used >= entity_limit:
                    conn.execute("COMMIT")
                    return ModuleActionBudgetDecision(
                        allowed=False,
                        reason="entity_budget_exhausted",
                        day_key=day_key,
                        daily_used=daily_used,
                        entity_used=entity_used,
                    )

                daily_used += 1
                entity_used += 1
                now_iso = _utc_now_iso()
                self._set_count(
                    conn,
                    normalized_module_id,
                    day_key,
                    self._GLOBAL_SCOPE,
                    self._GLOBAL_KEY,
                    daily_used,
                    now_iso,
                )
                self._set_count(
                    conn,
                    normalized_module_id,
                    day_key,
                    self._ENTITY_SCOPE,
                    normalized_entity,
                    entity_used,
                    now_iso,
                )
                conn.execute("COMMIT")
                return ModuleActionBudgetDecision(
                    allowed=True,
                    reason=None,
                    day_key=day_key,
                    daily_used=daily_used,
                    entity_used=entity_used,
                )
            except Exception:
                conn.execute("ROLLBACK")
                raise

    def daily_usage(self, *, module_id: str, day_key: str) -> int:
        normalized_module_id = str(module_id or "").strip().lower()
        if not normalized_module_id:
            return 0
        with self._connect() as conn:
            return self._get_count(
                conn,
                normalized_module_id,
                day_key,
                self._GLOBAL_SCOPE,
                self._GLOBAL_KEY,
            )

    def entity_usage(self, *, module_id: str, day_key: str, entity_key: str) -> int:
        normalized_module_id = str(module_id or "").strip().lower()
        normalized_entity = str(entity_key or "").strip().lower() or "entity:unknown"
        if not normalized_module_id:
            return 0
        with self._connect() as conn:
            return self._get_count(
                conn,
                normalized_module_id,
                day_key,
                self._ENTITY_SCOPE,
                normalized_entity,
            )

    def _cleanup_stale_days(self, conn: sqlite3.Connection, module_id: str, day_key: str) -> None:
        conn.execute(
            "DELETE FROM module_action_budgets WHERE module_id = ? AND day_key < ?",
            (module_id, day_key),
        )

    def _get_count(
        self,
        conn: sqlite3.Connection,
        module_id: str,
        day_key: str,
        scope: str,
        scope_key: str,
    ) -> int:
        row = conn.execute(
            """
            SELECT count
            FROM module_action_budgets
            WHERE module_id = ? AND day_key = ? AND scope = ? AND scope_key = ?
            """,
            (module_id, day_key, scope, scope_key),
        ).fetchone()
        if not row:
            return 0
        try:
            return int(row["count"])
        except Exception:
            return 0

    def _set_count(
        self,
        conn: sqlite3.Connection,
        module_id: str,
        day_key: str,
        scope: str,
        scope_key: str,
        count: int,
        updated_at: str,
    ) -> None:
        conn.execute(
            """
            INSERT INTO module_action_budgets(module_id, day_key, scope, scope_key, count, updated_at)
            VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(module_id, day_key, scope, scope_key)
            DO UPDATE SET
                count = excluded.count,
                updated_at = excluded.updated_at
            """,
            (module_id, day_key, scope, scope_key, int(count), updated_at),
        )


__all__ = ["ModuleActionBudgetStore", "ModuleActionBudgetDecision"]
