"""Authoritative request-unit limits for costly DocDoe operations.

Policy is intentionally simple: a successful operation consumes one or more
student-readable units; rejection and provider failure do not. Reservations
count toward the limit and concurrency cap while work is running. A logical
retry reuses the same user/operation/idempotency row, so it cannot double-count.
Prompts and response bodies are never stored here.
"""
from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Literal

from fastapi import HTTPException, status
from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.models.usage_operation import UsageOperation
from app.models.user import User
from app.services.entitlements import EntitlementFeature, require_entitlement_feature, resolve_entitlements
from app.services.plan_catalog import get_monthly_usage_limits, get_plan_config
from app.utils.request_id import request_id_ctx_var

logger = logging.getLogger(__name__)

UsageOperationName = Literal[
    "study_chat_ai",
    "tuition_tutor",
    "correction_ai",
    "source_processing",
    "learn_anything_generation",
    "pyq_style_generation",
    "video_generation",
    "upload_storage",
    "exam_intelligence_generation",
]

_DEFAULT_POLICY: dict[str, dict[str, int | str]] = {
    "study_chat_ai": {"limit": 50, "period": "day", "max_active": 1},
    "tuition_tutor": {"limit": 20, "period": "day", "max_active": 1},
    "correction_ai": {"limit": 20, "period": "day", "max_active": 1},
    "source_processing": {"limit": 10, "period": "month", "max_active": 1},
    "learn_anything_generation": {"limit": 10, "period": "day", "max_active": 1},
    "pyq_style_generation": {"limit": 20, "period": "day", "max_active": 1},
    "video_generation": {"limit": 2, "period": "month", "max_active": 1},
    "upload_storage": {"limit": 100 * 1024 * 1024, "period": "month", "max_active": 2},
    # Audit H3: the two generative GET endpoints in /intelligence/ previously
    # had no per-user quota at all. They share the exam-intelligence feature
    # gate with pyq_style_generation but get their own daily budget.
    "exam_intelligence_generation": {"limit": 20, "period": "day", "max_active": 1},
}

_OPERATION_FEATURES: dict[str, EntitlementFeature] = {
    "study_chat_ai": "study_chat_ai",
    "tuition_tutor": "tuition_ai_tutor",
    "correction_ai": "deterministic_correction",
    "source_processing": "source_upload",
    "learn_anything_generation": "learn_anything_generation",
    "pyq_style_generation": "pyq_style_generation",
    "video_generation": "heavy_generation",
    "upload_storage": "source_upload",
    "exam_intelligence_generation": "pyq_style_generation",
}

# Estimated provider USD cost per AI-token operation, used ONLY for the
# per-user daily budget in reserve_usage(). Values are derived from the
# provider cost table and typical token volumes, not billed amounts. Operations
# that do not consume LLM tokens (source_processing, video_generation,
# upload_storage) are intentionally absent so they are never double-counted.
_OPERATION_ESTIMATED_COST_USD: dict[str, float] = {
    "study_chat_ai": 0.002,
    "tuition_tutor": 0.004,
    "correction_ai": 0.003,
    "learn_anything_generation": 0.008,
    "pyq_style_generation": 0.006,
    "exam_intelligence_generation": 0.006,
}


@dataclass(frozen=True)
class UsagePolicy:
    limit: int
    period: Literal["day", "month"]
    max_active: int


@dataclass(frozen=True)
class UsageReservation:
    record: UsageOperation
    duplicate_completed: bool = False

    @property
    def should_charge(self) -> bool:
        return not self.duplicate_completed


def _policies(policy_key: str = "free_trial") -> dict[str, UsagePolicy]:
    raw: dict[str, Any] = {}
    configured = get_settings().usage_policy_json.strip()
    if configured:
        try:
            parsed = json.loads(configured)
            if isinstance(parsed, dict):
                raw = parsed
        except (TypeError, ValueError):
            logger.warning("usage_policy_invalid", extra={"event": "usage_policy_invalid"})
    free_config = get_plan_config("free_trial")
    plan_config = get_plan_config(policy_key) or free_config
    generation_multiplier = (
        max(1, plan_config.monthly_generation_limit // max(1, free_config.monthly_generation_limit))
        if plan_config and free_config
        else 1
    )
    legacy_limits = get_monthly_usage_limits(policy_key)
    policy_defaults = {name: dict(values) for name, values in _DEFAULT_POLICY.items()}
    for operation in (
        "study_chat_ai",
        "tuition_tutor",
        "correction_ai",
        "learn_anything_generation",
        "pyq_style_generation",
        "exam_intelligence_generation",
    ):
        policy_defaults[operation]["limit"] = int(policy_defaults[operation]["limit"]) * generation_multiplier
    policy_defaults["source_processing"]["limit"] = int(
        legacy_limits.get("documents_uploaded_count", policy_defaults["source_processing"]["limit"])
    )
    policy_defaults["video_generation"]["limit"] = int(
        plan_config.monthly_video_limit if plan_config else policy_defaults["video_generation"]["limit"]
    )
    policy_defaults["upload_storage"]["limit"] = int(
        legacy_limits.get("storage_bytes_used", policy_defaults["upload_storage"]["limit"])
    )
    policy_overrides = raw.get("policies", {}) if isinstance(raw.get("policies"), dict) else {}
    selected_overrides = (
        policy_overrides.get(policy_key, {})
        if isinstance(policy_overrides.get(policy_key), dict)
        else {}
    )
    result: dict[str, UsagePolicy] = {}
    for operation, defaults in policy_defaults.items():
        global_override = raw.get(operation) if isinstance(raw.get(operation), dict) else {}
        specific_override = (
            selected_overrides.get(operation)
            if isinstance(selected_overrides.get(operation), dict)
            else {}
        )
        override = {**global_override, **specific_override}
        period = override.get("period", defaults["period"])
        result[operation] = UsagePolicy(
            limit=max(0, int(override.get("limit", defaults["limit"]))),
            period="month" if period == "month" else "day",
            max_active=max(1, int(override.get("max_active", defaults["max_active"]))),
        )
    return result


def _period_key(policy: UsagePolicy, now: datetime) -> str:
    return now.strftime("%Y-%m" if policy.period == "month" else "%Y-%m-%d")


def stable_idempotency_key(*parts: object) -> str:
    """Hash metadata into a bounded key without persisting private prompt text."""
    raw = "\x1f".join(str(part) for part in parts)
    return hashlib.sha256(raw.encode("utf-8", errors="ignore")).hexdigest()


def _student_error(code: str, message: str, *, http_status: int) -> HTTPException:
    return HTTPException(
        status_code=http_status,
        detail={
            "code": code,
            "message": message,
            "retryable": code == "OPERATION_IN_PROGRESS",
            "request_id": request_id_ctx_var.get("-"),
        },
    )


def _duplicate_completed_error() -> HTTPException:
    """Reject a replay of an already-completed idempotency key.

    Used by callers whose completed row stores no result: replaying the key
    would otherwise re-run the paid operation for free.
    """
    return _student_error(
        "DUPLICATE_REQUEST",
        "This request has already been completed. Start a new request to continue.",
        http_status=status.HTTP_409_CONFLICT,
    )


def _log(event: str, *, user_id: str, operation: str, units: int, policy: UsagePolicy, remaining: int) -> None:
    logger.info(
        event,
        extra={
            "event": event,
            "operation": operation,
            "user_ref": stable_idempotency_key(user_id)[:16],
            "units": units,
            "remaining": max(0, remaining),
            "limit": policy.limit,
            "period": policy.period,
        },
    )


def _daily_budget_denial(
    db: Session, *, user_id: str, operation: str, now: datetime
) -> tuple[float, float, float] | None:
    """Return (day_cost, this_cost, cap) when the per-user daily AI budget is
    exceeded, else None.

    Uses the same day period_key convention and reserved/completed statuses as
    the unit-limit query above. Fails OPEN: any accounting error returns None so
    a broken cost lookup can never take down AI generation for students.
    """
    try:
        cap = float(get_settings().ai_router_per_user_daily_budget_usd or 0)
        if cap <= 0:
            return None
        cost = float(_OPERATION_ESTIMATED_COST_USD.get(operation, 0) or 0)
        if cost <= 0:
            return None
        day_key = now.strftime("%Y-%m-%d")
        rows = db.execute(
            select(UsageOperation.operation, func.coalesce(func.sum(UsageOperation.units), 0))
            .where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation.in_(list(_OPERATION_ESTIMATED_COST_USD)),
                UsageOperation.period_key == day_key,
                UsageOperation.status.in_(["reserved", "completed"]),
            )
            .group_by(UsageOperation.operation)
        ).all()
        day_cost = sum(
            float(_OPERATION_ESTIMATED_COST_USD.get(op, 0)) * int(units or 0)
            for op, units in rows
        )
        if day_cost + cost > cap:
            return day_cost, cost, cap
        return None
    except Exception:  # noqa: BLE001 - accounting must fail open
        logger.warning(
            "usage_budget_check_failed",
            exc_info=True,
            extra={"event": "usage_budget_check_failed", "operation": operation},
        )
        return None


def reserve_usage(
    db: Session,
    *,
    user_id: str,
    operation: UsageOperationName,
    idempotency_key: str,
    units: int = 1,
    reject_completed_duplicate: bool = False,
) -> UsageReservation:
    """Reserve quota before work.

    A completed row whose idempotency key is reused is normally returned as a
    no-charge duplicate. That is correct for resource-backed operations
    (uploads, video renders) which look up the stored ``resource_id`` and return
    the existing artefact. Callers whose completed row stores no result — chat,
    where replaying the key would otherwise re-run a paid AI call for free —
    pass ``reject_completed_duplicate=True`` to reject the replay instead.
    """
    if units <= 0:
        raise ValueError("Usage units must be positive")
    entitlement = require_entitlement_feature(db, user_id, _OPERATION_FEATURES[operation])
    policy = _policies(entitlement.usage_policy_key)[operation]
    if not get_settings().usage_limits_enabled:
        policy = UsagePolicy(limit=2_147_483_647, period=policy.period, max_active=policy.max_active)
    now = datetime.now(timezone.utc)
    period_key = _period_key(policy, now)
    key = idempotency_key.strip()[:180] or stable_idempotency_key(request_id_ctx_var.get("-"))

    existing = db.scalar(
        select(UsageOperation).where(
            UsageOperation.user_id == user_id,
            UsageOperation.operation == operation,
            UsageOperation.idempotency_key == key,
        )
    )
    if existing is not None and existing.status == "completed":
        if reject_completed_duplicate:
            _log("usage_duplicate_rejected", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
            raise _duplicate_completed_error()
        _log("usage_allowed", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
        return UsageReservation(existing, duplicate_completed=True)
    if existing is not None and existing.status == "reserved":
        cutoff = now - timedelta(seconds=get_settings().usage_reservation_ttl_seconds)
        updated_at = existing.updated_at
        if updated_at.tzinfo is None:
            updated_at = updated_at.replace(tzinfo=timezone.utc)
        if updated_at >= cutoff:
            raise _student_error(
                "OPERATION_IN_PROGRESS",
                "That request is already running.",
                http_status=status.HTTP_409_CONFLICT,
            )
        existing.status = "released"
        db.add(existing)
        db.commit()

    # PostgreSQL serializes reservations per user. SQLite remains deterministic
    # in local tests through its write lock plus the unique logical key.
    db.scalar(select(User.id).where(User.id == user_id).with_for_update())
    used = int(
        db.scalar(
            select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.period_key == period_key,
                UsageOperation.status.in_(["reserved", "completed"]),
            )
        )
        or 0
    )
    if used + units > policy.limit:
        _log("usage_denied", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used)
        raise _student_error(
            "USAGE_LIMIT_REACHED",
            "You've reached today's AI study limit. Try again when your allowance resets.",
            http_status=status.HTTP_402_PAYMENT_REQUIRED,
        )

    # Per-user daily USD budget: a cost backstop on top of the count limits
    # above. Denies through the same path/shape as the unit limit.
    budget_denial = _daily_budget_denial(db, user_id=user_id, operation=operation, now=now)
    if budget_denial is not None:
        day_cost, operation_cost, cap = budget_denial
        logger.warning(
            "usage_budget_denied",
            extra={
                "event": "usage_budget_denied",
                "operation": operation,
                "user_ref": stable_idempotency_key(user_id)[:16],
                "estimated_day_cost_usd": round(day_cost, 6),
                "estimated_operation_cost_usd": round(operation_cost, 6),
                "cap_usd": cap,
            },
        )
        raise _student_error(
            "USAGE_LIMIT_REACHED",
            "You've reached today's AI study budget. Try again when your allowance resets.",
            http_status=status.HTTP_402_PAYMENT_REQUIRED,
        )

    active = int(
        db.scalar(
            select(func.count()).select_from(UsageOperation).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.status == "reserved",
            )
        )
        or 0
    )
    if active >= policy.max_active:
        raise _student_error(
            "OPERATION_IN_PROGRESS",
            "That request is already running.",
            http_status=status.HTTP_409_CONFLICT,
        )

    if existing is not None:
        existing.status = "reserved"
        existing.period_key = period_key
        existing.units = units
        existing.request_id = request_id_ctx_var.get("-")
        existing.completed_at = None
        record = existing
    else:
        record = UsageOperation(
            user_id=user_id,
            operation=operation,
            period_key=period_key,
            idempotency_key=key,
            units=units,
            status="reserved",
            request_id=request_id_ctx_var.get("-"),
        )
    db.add(record)
    try:
        db.commit()
        db.refresh(record)
    except IntegrityError:
        db.rollback()
        raced = db.scalar(
            select(UsageOperation).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.idempotency_key == key,
            )
        )
        if raced is not None and raced.status == "completed":
            if reject_completed_duplicate:
                raise _duplicate_completed_error()
            return UsageReservation(raced, duplicate_completed=True)
        raise _student_error(
            "OPERATION_IN_PROGRESS",
            "That request is already running.",
            http_status=status.HTTP_409_CONFLICT,
        )
    _log("usage_allowed", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used - units)
    return UsageReservation(record)


def finalize_usage(db: Session, reservation: UsageReservation, *, resource_id: str | None = None) -> None:
    if not reservation.should_charge:
        return
    db.execute(
        update(UsageOperation)
        .where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
        .values(status="completed", completed_at=datetime.now(timezone.utc), resource_id=resource_id)
    )
    db.commit()
    entitlement = resolve_entitlements(db, reservation.record.user_id)
    policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
    _log("usage_recorded", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


def release_usage(db: Session, reservation: UsageReservation) -> None:
    if not reservation.should_charge:
        return
    db.execute(
        update(UsageOperation)
        .where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
        .values(status="released", completed_at=None)
    )
    db.commit()
    entitlement = resolve_entitlements(db, reservation.record.user_id)
    policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
    _log("usage_reservation_released", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


def get_usage_status(db: Session, user_id: str) -> dict[str, Any]:
    entitlement = resolve_entitlements(db, user_id)
    items: list[dict[str, Any]] = []
    for operation, policy in _policies(entitlement.usage_policy_key).items():
        period_key = _period_key(policy, datetime.now(timezone.utc))
        used = int(
            db.scalar(
                select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
                    UsageOperation.user_id == user_id,
                    UsageOperation.operation == operation,
                    UsageOperation.period_key == period_key,
                    UsageOperation.status.in_(["reserved", "completed"]),
                )
            )
            or 0
        )
        items.append(
            {
                "operation": operation,
                "period": policy.period,
                "used": used,
                "limit": policy.limit,
                "remaining": max(0, policy.limit - used),
                "allowed": used < policy.limit,
            }
        )
    return {
        "tier": entitlement.tier,
        "entitlement_status": entitlement.status,
        "usage_policy_key": entitlement.usage_policy_key,
        "items": items,
    }