"""Usage counters, failure counters, and the daily quota gate.

One table serves the first and the last of those, and that is the point:
:class:`~app.domains.ai_labs.models.AiLabUsageDaily` is what the admin
console reads about AI Labs, and it is also the only thing the quota gate
consults. There is no second, richer store an admin could be tempted to
query instead — the numbers and the enforcement come from the same
content-free rows.

:class:`~app.domains.ai_labs.models.AiLabFailureDaily` is the one addition,
and it is content-free by the same construction: it adds the *reason*
dimension that ``AiLabUsageDaily.failures`` lacks, holding an error code
and a count. Every failure recorded through :func:`record` books into both,
so "8 requests, 2 failures" can be followed by "both were
AI_LABS_INTERRUPTED on chat" without going near a table that holds a
prompt.

Three properties worth stating:

* **The day is a local day.** Quota windows are computed in the platform's
  reporting timezone (Asia/Shanghai by default), not UTC, because a UTC
  reset lands mid-afternoon in Beijing and an employee's "today" has to be
  a real day.
* **Recording never fails a request.** A counter write that hits an error
  is logged and swallowed. Losing a statistic is a rounding error; failing
  an employee's finished generation because the bookkeeping row would not
  save is not a trade anyone would choose.
* **An error code is not free text.** Vendor error codes arrive as
  arbitrary strings and the async media family quotes rejected prompts back
  inside them, so anything not code-shaped is counted under
  :data:`CODE_UNCLASSIFIED` instead of being stored. See
  :func:`normalize_error_code`.
"""
from __future__ import annotations

import logging
import re
import uuid
from dataclasses import dataclass
from datetime import date, datetime
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from sqlalchemy import case, func, select
from sqlalchemy.dialects.mysql import insert as mysql_insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.exceptions import AppException
from app.domains.ai_labs import catalog
from app.domains.ai_labs.models import AiLabFailureDaily, AiLabUsageDaily
from app.domains.ai_labs.settings import AiLabsConfig, Quotas

logger = logging.getLogger(__name__)

#: Pseudo-capability for chat turns, so the rollup has one row shape.
CAPABILITY_CHAT = "chat"
#: Pseudo-capability for translations.
CAPABILITY_TRANSLATE = "translate"

#: What an error code is allowed to look like in the admin-readable
#: failures table. Our own codes are ``AI_LABS_*``; the vendor's arrive as
#: identifiers too (``DataInspectionFailed``, ``InvalidParameter``).
_CODE_SHAPE = re.compile(r"^[A-Za-z0-9._:-]{1,64}$")

#: Bucket for a code that is not code-shaped. Reached when a vendor puts
#: something else in its ``error_code`` field — and the async media family
#: is documented to quote the offending prompt back inside content-policy
#: rejections, so "something else" can be a fragment of what the employee
#: typed. This surface must not become the place that leaks it, so the
#: failure is still counted, under a name that says only that it was not
#: classifiable. The raw string is logged for whoever is debugging.
CODE_UNCLASSIFIED = "AI_LABS_VENDOR_UNCLASSIFIED"

#: Bucket for a failure recorded before the code was in hand at all, and
#: for the historical rows that predate ``ai_labs_failures_daily``. See
#: ``AiLabsService.failure_overview``, which reconciles the two tables
#: rather than pretending the codes cover every counted failure.
CODE_UNKNOWN = "AI_LABS_UNKNOWN"


def normalize_error_code(raw: str | None) -> str:
    """Coerce an error code into something safe to store and group by.

    Returns :data:`CODE_UNKNOWN` for nothing at all and
    :data:`CODE_UNCLASSIFIED` for anything that is not code-shaped — never
    a truncated or scrubbed version of the input, because a scrubbed
    fragment of a prompt is still a fragment of a prompt.
    """
    candidate = (raw or "").strip()
    if not candidate:
        return CODE_UNKNOWN
    if not _CODE_SHAPE.match(candidate):
        logger.warning(
            "ai_labs: refusing to store a non-code-shaped error code "
            "(length %d) — counting it as %s",
            len(candidate), CODE_UNCLASSIFIED,
        )
        return CODE_UNCLASSIFIED
    return candidate


def local_today(timezone_name: str) -> date:
    """Today's date in the reporting timezone.

    Falls back to UTC on an unknown zone name rather than raising: a typo in
    a settings row must not take the whole surface down, and the rollup being
    off by a few hours is a far smaller problem than a 500.
    """
    try:
        zone = ZoneInfo(timezone_name)
    except (ZoneInfoNotFoundError, ValueError, KeyError):
        logger.warning(
            "ai_labs: unknown reporting timezone %r — falling back to UTC",
            timezone_name,
        )
        return datetime.now(ZoneInfo("UTC")).date()
    return datetime.now(zone).date()


@dataclass(frozen=True)
class DailyTotals:
    """One person's totals for one local day, across every capability."""

    tokens: int = 0
    images: int = 0
    video_seconds: int = 0
    speech: int = 0
    transcribe: int = 0


async def daily_totals(
    db: AsyncSession,
    *,
    person_id: uuid.UUID,
    on_date: date,
) -> DailyTotals:
    """Aggregate one person's counters for one day."""
    def _sum_when(condition, column):  # noqa: ANN001, ANN202 — local shorthand
        """``SUM(CASE WHEN cond THEN col ELSE 0 END)``, coalesced to 0.

        ``case`` rather than ``func.if_``: the latter renders as a literal
        ``if_(...)``, which is not a function in any dialect.
        """
        return func.coalesce(func.sum(case((condition, column), else_=0)), 0)

    stmt = select(
        func.coalesce(
            func.sum(
                AiLabUsageDaily.input_tokens + AiLabUsageDaily.output_tokens,
            ), 0,
        ),
        # Only image-producing capabilities count toward the image ceiling;
        # a video's artifact is a clip, not a picture.
        _sum_when(
            AiLabUsageDaily.capability.in_(
                (catalog.CAP_IMAGE, catalog.CAP_IMAGE_EDIT),
            ),
            AiLabUsageDaily.artifacts,
        ),
        _sum_when(
            AiLabUsageDaily.capability == catalog.CAP_VIDEO,
            AiLabUsageDaily.media_seconds,
        ),
        _sum_when(
            AiLabUsageDaily.capability == catalog.CAP_SPEECH,
            AiLabUsageDaily.requests,
        ),
        _sum_when(
            AiLabUsageDaily.capability == catalog.CAP_TRANSCRIBE,
            AiLabUsageDaily.requests,
        ),
    ).where(
        AiLabUsageDaily.person_id == person_id,
        AiLabUsageDaily.usage_date == on_date,
    )
    row = (await db.execute(stmt)).one()
    return DailyTotals(
        tokens=int(row[0] or 0),
        images=int(row[1] or 0),
        video_seconds=int(row[2] or 0),
        speech=int(row[3] or 0),
        transcribe=int(row[4] or 0),
    )


def _quota_error(kind: str, used: int, limit: int) -> AppException:
    return AppException(
        code="AI_LABS_QUOTA_EXCEEDED",
        message=(
            f"Your daily AI Labs {kind} allowance is used up "
            f"({used}/{limit}). It resets at midnight."
        ),
        status_code=429,
    )


async def assert_within_quota(
    db: AsyncSession,
    *,
    person_id: uuid.UUID,
    config: AiLabsConfig,
    capability_key: str,
    #: For video, the length about to be requested — checked before the job
    #: starts, because refunding a clip is not a thing.
    video_seconds: int = 0,
) -> None:
    """Raise 429 if this request would cross a configured ceiling.

    Returns immediately when nothing is configured, which is the default
    state — no query, no cost.
    """
    quotas: Quotas = config.quotas
    if not quotas.any_enforced:
        return

    totals = await daily_totals(
        db,
        person_id=person_id,
        on_date=local_today(config.reporting_timezone),
    )

    if quotas.daily_tokens and totals.tokens >= quotas.daily_tokens:
        raise _quota_error("token", totals.tokens, quotas.daily_tokens)

    if capability_key in (catalog.CAP_IMAGE, catalog.CAP_IMAGE_EDIT):
        if quotas.daily_images and totals.images >= quotas.daily_images:
            raise _quota_error("image", totals.images, quotas.daily_images)

    if capability_key == catalog.CAP_VIDEO and quotas.daily_video_seconds:
        # The ceiling is on produced seconds, so the check has to include
        # what this job would add — otherwise a 10s clip always squeezes
        # past a limit with 1 second left.
        projected = totals.video_seconds + max(0, video_seconds)
        if projected > quotas.daily_video_seconds:
            raise _quota_error(
                "video-second", totals.video_seconds, quotas.daily_video_seconds,
            )

    if capability_key == catalog.CAP_SPEECH:
        if quotas.daily_speech and totals.speech >= quotas.daily_speech:
            raise _quota_error(
                "speech", totals.speech, quotas.daily_speech,
            )

    if capability_key == catalog.CAP_TRANSCRIBE:
        if quotas.daily_transcribe and totals.transcribe >= quotas.daily_transcribe:
            raise _quota_error(
                "transcription", totals.transcribe, quotas.daily_transcribe,
            )


async def record(
    *,
    person_id: uuid.UUID,
    capability_key: str,
    model: str,
    timezone_name: str,
    requests: int = 1,
    failures: int = 0,
    input_tokens: int = 0,
    output_tokens: int = 0,
    artifacts: int = 0,
    media_seconds: int = 0,
    #: Required in practice whenever ``failures`` is non-zero: without it
    #: the failure is counted but lands in the ``AI_LABS_UNKNOWN`` bucket,
    #: which is exactly the blind spot ``ai_labs_failures_daily`` exists to
    #: close. Optional in the signature so a caller recording a *success*
    #: does not have to pass ``None``.
    error_code: str | None = None,
) -> None:
    """Add to today's counters, in **its own transaction**. Never raises.

    Two reasons it does not join the caller's session, and they point the
    same way:

    * *Semantics.* The tokens were spent at the vendor the moment the
      request returned. If the caller's transaction then rolls back — a
      failed message insert, a lost race on a lease — the spend still
      happened, and a counter that rolled back with it would under-report
      real consumption and hand out quota that was already used.
    * *Safety.* A failed ``INSERT`` inside the caller's transaction poisons
      that transaction; swallowing the exception there would leave the
      session unusable and make the *next* statement fail with something
      unrelated. Catching a failure only works if the failure is isolated,
      and isolating it means a separate session. Same reasoning as
      :func:`app.audit.service.log_denied_action`.

    ``INSERT ... ON DUPLICATE KEY UPDATE`` on the unique
    ``(person, day, capability, model)`` key, so two concurrent turns from
    one person accumulate rather than racing over a read-modify-write.
    """
    try:
        # See the note in ``chat._persist_answer``: the factory must be
        # looked up at call time so the test harness's patched, schema-
        # isolated factory is the one used.
        from app.core.database import async_session_factory

        stmt = mysql_insert(AiLabUsageDaily).values(
            id=uuid.uuid4(),
            person_id=person_id,
            usage_date=local_today(timezone_name),
            capability=capability_key,
            model=(model or "")[:64],
            requests=requests,
            failures=failures,
            input_tokens=max(0, input_tokens),
            output_tokens=max(0, output_tokens),
            artifacts=max(0, artifacts),
            media_seconds=max(0, media_seconds),
        )
        stmt = stmt.on_duplicate_key_update(
            requests=AiLabUsageDaily.requests + stmt.inserted.requests,
            failures=AiLabUsageDaily.failures + stmt.inserted.failures,
            input_tokens=(
                AiLabUsageDaily.input_tokens + stmt.inserted.input_tokens
            ),
            output_tokens=(
                AiLabUsageDaily.output_tokens + stmt.inserted.output_tokens
            ),
            artifacts=AiLabUsageDaily.artifacts + stmt.inserted.artifacts,
            media_seconds=(
                AiLabUsageDaily.media_seconds + stmt.inserted.media_seconds
            ),
            updated_at=func.now(),
        )
        async with async_session_factory() as own:
            await own.execute(stmt)
            await own.commit()
    except Exception as exc:  # noqa: BLE001 — bookkeeping must not fail a turn
        logger.warning(
            "ai_labs: failed to record usage for person=%s capability=%s: %s",
            person_id, capability_key, exc,
        )

    if failures > 0:
        # Deliberately after the commit above and in its own session, so a
        # problem writing the reason cannot cost us the spend counter — the
        # quota gate reads that one, and an under-reported quota hands out
        # allowance that was already used.
        await _record_failure_reason(
            person_id=person_id,
            capability_key=capability_key,
            model=model,
            timezone_name=timezone_name,
            error_code=error_code,
            failures=failures,
        )


async def _record_failure_reason(
    *,
    person_id: uuid.UUID,
    capability_key: str,
    model: str,
    timezone_name: str,
    error_code: str | None,
    failures: int,
) -> None:
    """Add to today's per-reason failure counters. Never raises.

    Same ``INSERT ... ON DUPLICATE KEY UPDATE`` shape and same
    own-transaction reasoning as :func:`record`. ``last_seen_at`` is
    overwritten rather than accumulated: the question it answers is "is
    this still happening?", so the newest occurrence is the useful one.
    """
    try:
        from app.core.database import async_session_factory

        now = datetime.now(ZoneInfo("UTC"))
        stmt = mysql_insert(AiLabFailureDaily).values(
            id=uuid.uuid4(),
            person_id=person_id,
            usage_date=local_today(timezone_name),
            capability=capability_key,
            model=(model or "")[:64],
            error_code=normalize_error_code(error_code),
            failures=failures,
            last_seen_at=now,
        )
        stmt = stmt.on_duplicate_key_update(
            failures=AiLabFailureDaily.failures + stmt.inserted.failures,
            last_seen_at=stmt.inserted.last_seen_at,
            updated_at=func.now(),
        )
        async with async_session_factory() as own:
            await own.execute(stmt)
            await own.commit()
    except Exception as exc:  # noqa: BLE001 — bookkeeping must not fail a turn
        logger.warning(
            "ai_labs: failed to record the failure reason for person=%s "
            "capability=%s: %s",
            person_id, capability_key, exc,
        )


__all__ = [
    "CAPABILITY_CHAT",
    "CAPABILITY_TRANSLATE",
    "CODE_UNCLASSIFIED",
    "CODE_UNKNOWN",
    "DailyTotals",
    "assert_within_quota",
    "daily_totals",
    "local_today",
    "normalize_error_code",
    "record",
]