"""
Usage metering recorder.

Provides best-effort helpers to persist LLM and TTS usage records.
Billing settlement stays asynchronous; request threads stop after raw
`bill_usage` persistence and may only enqueue follow-up settlement work.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional

from flask import Flask

from flaskr.dao import db
from flaskr.service.shifu.demo_courses import is_builtin_demo_shifu
from flaskr.util.uuid import generate_id

from .consts import (
    BILL_USAGE_SCENE_PROD,
    BILL_USAGE_TYPE_LLM,
    BILL_USAGE_TYPE_TTS,
    normalize_usage_scene,
)
from .models import BillUsageRecord


@dataclass(frozen=True)
class UsageContext:
    """Context fields shared by usage records."""

    user_bid: str = ""
    shifu_bid: str = ""
    outline_item_bid: str = ""
    progress_record_bid: str = ""
    generated_block_bid: str = ""
    audio_bid: str = ""
    request_id: str = ""
    trace_id: str = ""
    usage_scene: int = BILL_USAGE_SCENE_PROD
    billable: Optional[int] = None


def _resolve_billable(app: Flask, *, context: UsageContext, usage_scene: int) -> int:
    if context.billable is not None:
        return int(context.billable)
    if context.shifu_bid and is_builtin_demo_shifu(app, context.shifu_bid):
        return 0
    normalize_usage_scene(usage_scene)
    return 1


def _persist_usage_record(app: Flask, record: BillUsageRecord) -> bool:
    """Persist raw usage only; async settlement owns later credit mutations."""
    try:
        with app.app_context():
            db.session.add(record)
            db.session.commit()
        return True
    except Exception as exc:
        try:
            app.logger.error("Usage metering persist failed: %s", exc, exc_info=True)
        except Exception:
            # Ignore logging failures to avoid masking the original persistence error.
            pass
        try:
            db.session.rollback()
        except Exception:
            # Ignore rollback failures; session may already be invalidated.
            pass
        return False


def _should_enqueue_usage_settlement(
    *,
    billable: int,
    status: int,
    record_level: int,
) -> bool:
    return (
        int(billable or 0) == 1
        and int(status or 0) == 0
        and int(record_level or 0) == 0
    )


def _enqueue_usage_settlement(app: Flask, *, usage_bid: str) -> None:
    normalized_usage_bid = str(usage_bid or "").strip()
    if not normalized_usage_bid:
        return
    try:
        from flaskr.common.celery_app import get_celery_app

        celery_app = get_celery_app(flask_app=app)
        task = celery_app.tasks.get("billing.settle_usage")
        if task is None:
            app.logger.warning(
                "billing.settle_usage is unavailable for usage_bid=%s",
                normalized_usage_bid,
            )
            return
        task.apply_async(kwargs={"usage_bid": normalized_usage_bid})
    except Exception as exc:
        try:
            app.logger.error(
                "Usage settlement enqueue failed for usage_bid=%s: %s",
                normalized_usage_bid,
                exc,
                exc_info=True,
            )
        except Exception:
            pass


def record_llm_usage(
    app: Flask,
    context: UsageContext,
    *,
    provider: str,
    model: str,
    is_stream: bool,
    input: int,
    input_cache: int = 0,
    output: int,
    total: int,
    latency_ms: int = 0,
    status: int = 0,
    error_message: str = "",
    extra: Optional[Dict[str, Any]] = None,
) -> str:
    usage_bid = generate_id(app)
    normalized_usage_scene = normalize_usage_scene(context.usage_scene)
    resolved_billable = _resolve_billable(
        app,
        context=context,
        usage_scene=normalized_usage_scene,
    )
    record = BillUsageRecord(
        usage_bid=usage_bid,
        parent_usage_bid="",
        user_bid=context.user_bid or "",
        shifu_bid=context.shifu_bid or "",
        outline_item_bid=context.outline_item_bid or "",
        progress_record_bid=context.progress_record_bid or "",
        generated_block_bid=context.generated_block_bid or "",
        audio_bid=context.audio_bid or "",
        request_id=context.request_id or "",
        trace_id=context.trace_id or "",
        usage_type=BILL_USAGE_TYPE_LLM,
        record_level=0,
        usage_scene=normalized_usage_scene,
        provider=provider or "",
        model=model or "",
        is_stream=1 if is_stream else 0,
        input=int(input or 0),
        input_cache=int(input_cache or 0),
        output=int(output or 0),
        total=int(total or 0),
        word_count=0,
        duration_ms=0,
        latency_ms=int(latency_ms or 0),
        segment_index=0,
        segment_count=0,
        billable=resolved_billable,
        status=int(status or 0),
        error_message=error_message or "",
        extra=extra or None,
    )
    if _persist_usage_record(app, record):
        if _should_enqueue_usage_settlement(
            billable=resolved_billable,
            status=status,
            record_level=0,
        ):
            _enqueue_usage_settlement(app, usage_bid=usage_bid)
        try:
            usage_source = (
                (extra or {}).get("usage_source", "") if isinstance(extra, dict) else ""
            )
            app.logger.info(
                "record_llm_usage saved usage_bid=%s provider=%s model=%s stream=%s "
                "input=%s input_cache=%s output=%s total=%s latency_ms=%s status=%s "
                "scene=%s billable=%s usage_source=%s user_bid=%s request_id=%s trace_id=%s",
                usage_bid,
                provider or "",
                model or "",
                1 if is_stream else 0,
                int(input or 0),
                int(input_cache or 0),
                int(output or 0),
                int(total or 0),
                int(latency_ms or 0),
                int(status or 0),
                int(normalized_usage_scene or 0),
                int(resolved_billable or 0),
                usage_source,
                context.user_bid or "",
                context.request_id or "",
                context.trace_id or "",
            )
        except Exception:
            # Best-effort logging; ignore failures so they do not mask the result.
            pass
        return usage_bid
    return ""


def record_tts_usage(
    app: Flask,
    context: UsageContext,
    *,
    usage_bid: Optional[str] = None,
    provider: str,
    model: str,
    is_stream: bool,
    input: int,
    output: int,
    total: int,
    word_count: int,
    duration_ms: int,
    latency_ms: int = 0,
    record_level: int = 0,
    parent_usage_bid: str = "",
    segment_index: int = 0,
    segment_count: int = 0,
    status: int = 0,
    error_message: str = "",
    extra: Optional[Dict[str, Any]] = None,
) -> str:
    resolved_usage_bid = usage_bid or generate_id(app)
    normalized_usage_scene = normalize_usage_scene(context.usage_scene)
    resolved_billable = _resolve_billable(
        app,
        context=context,
        usage_scene=normalized_usage_scene,
    )
    record = BillUsageRecord(
        usage_bid=resolved_usage_bid,
        parent_usage_bid=parent_usage_bid or "",
        user_bid=context.user_bid or "",
        shifu_bid=context.shifu_bid or "",
        outline_item_bid=context.outline_item_bid or "",
        progress_record_bid=context.progress_record_bid or "",
        generated_block_bid=context.generated_block_bid or "",
        audio_bid=context.audio_bid or "",
        request_id=context.request_id or "",
        trace_id=context.trace_id or "",
        usage_type=BILL_USAGE_TYPE_TTS,
        record_level=int(record_level or 0),
        usage_scene=normalized_usage_scene,
        provider=provider or "",
        model=model or "",
        is_stream=1 if is_stream else 0,
        input=int(input or 0),
        output=int(output or 0),
        total=int(total or 0),
        word_count=int(word_count or 0),
        duration_ms=int(duration_ms or 0),
        latency_ms=int(latency_ms or 0),
        segment_index=int(segment_index or 0),
        segment_count=int(segment_count or 0),
        billable=resolved_billable,
        status=int(status or 0),
        error_message=error_message or "",
        extra=extra or None,
    )
    if _persist_usage_record(app, record):
        if _should_enqueue_usage_settlement(
            billable=resolved_billable,
            status=status,
            record_level=record_level,
        ):
            _enqueue_usage_settlement(app, usage_bid=resolved_usage_bid)
        return resolved_usage_bid
    return ""
