"""Plan enforcement and usage metering.

Enforcement is a no-op unless ``billing_enabled`` is set, so small self-hosted
deployments are never surprised by a limit. Usage is metered into
``UsageRecord`` per billing period (``YYYY-MM``).
"""

from __future__ import annotations

import uuid
from datetime import UTC, datetime

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from legate.billing.plans import UNLIMITED, Plan, get_plan
from legate.config import get_settings
from legate.db.models.agent import Agent
from legate.db.models.membership import Membership
from legate.db.models.run import Run
from legate.db.models.usage import UsageRecord
from legate.db.models.workflow import Workflow
from legate.db.models.workspace import Workspace

_RESOURCE_MODEL = {"agents": Agent, "workflows": Workflow, "members": Membership}


class BillingLimitError(Exception):
    """Raised when a workspace exceeds its plan limit."""


def current_period(now: datetime | None = None) -> str:
    return (now or datetime.now(UTC)).strftime("%Y-%m")


def record_usage(
    db: Session,
    *,
    workspace_id: uuid.UUID,
    metric: str,
    quantity: int = 1,
    period: str | None = None,
) -> UsageRecord:
    period = period or current_period()
    record = db.scalars(
        select(UsageRecord).where(
            UsageRecord.workspace_id == workspace_id,
            UsageRecord.metric == metric,
            UsageRecord.period == period,
        )
    ).first()
    if record is None:
        record = UsageRecord(workspace_id=workspace_id, metric=metric, quantity=0, period=period)
        db.add(record)
    record.quantity += quantity
    db.flush()
    return record


def usage_for(
    db: Session, *, workspace_id: uuid.UUID, metric: str, period: str | None = None
) -> int:
    period = period or current_period()
    total = db.scalar(
        select(func.coalesce(func.sum(UsageRecord.quantity), 0)).where(
            UsageRecord.workspace_id == workspace_id,
            UsageRecord.metric == metric,
            UsageRecord.period == period,
        )
    )
    return int(total or 0)


def _count(db: Session, model, workspace_id: uuid.UUID) -> int:
    total = db.scalar(
        select(func.count()).select_from(model).where(model.workspace_id == workspace_id)
    )
    return int(total or 0)


def plan_for_workspace(workspace: Workspace) -> Plan:
    return get_plan(workspace.plan)


def enforce_resource(db: Session, *, workspace: Workspace, resource: str) -> None:
    """Block creating another ``resource`` (agents/workflows/members) past the plan."""
    if not get_settings().billing_enabled:
        return
    plan = plan_for_workspace(workspace)
    limit = getattr(plan, f"max_{resource}")
    if limit == UNLIMITED:
        return
    if _count(db, _RESOURCE_MODEL[resource], workspace.id) >= limit:
        raise BillingLimitError(
            f"The {plan.title} plan allows up to {limit} {resource}. Upgrade to add more."
        )


def enforce_run_quota(db: Session, *, workspace: Workspace) -> None:
    if not get_settings().billing_enabled:
        return
    plan = plan_for_workspace(workspace)
    if plan.max_runs_per_month == UNLIMITED:
        return
    if usage_for(db, workspace_id=workspace.id, metric="runs") >= plan.max_runs_per_month:
        raise BillingLimitError(
            f"The {plan.title} plan's monthly run quota ({plan.max_runs_per_month}) is exhausted."
        )


def require_feature(workspace: Workspace, feature: str) -> None:
    if not get_settings().billing_enabled:
        return
    if not plan_for_workspace(workspace).allows(feature):
        raise BillingLimitError(f"The {workspace.plan} plan does not include '{feature}'.")


def usage_summary(db: Session, workspace: Workspace) -> dict:
    plan = plan_for_workspace(workspace)
    period = current_period()
    tokens = db.scalar(
        select(func.coalesce(func.sum(Run.total_tokens), 0)).where(
            Run.workspace_id == workspace.id, Run.created_at >= f"{period}-01"
        )
    )
    return {
        "plan": plan.as_dict(),
        "period": period,
        "usage": {
            "agents": _count(db, Agent, workspace.id),
            "workflows": _count(db, Workflow, workspace.id),
            "members": _count(db, Membership, workspace.id),
            "runs": usage_for(db, workspace_id=workspace.id, metric="runs", period=period),
            "tokens": int(tokens or 0),
        },
        "enforced": get_settings().billing_enabled,
    }
