from __future__ import annotations

import hashlib
import json
from copy import deepcopy
from datetime import datetime, timezone
from decimal import Decimal
from typing import Any, Mapping

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.enums.model_pricing import ModelPricingRuleStatus, ProviderCostStatus, PricingSnapshotStage
from app.models.credit_record import CreditRecord
from app.models.model_pricing_rule import ModelPricingRule
from app.services.model_pricing.calculator import PricingCalculationError, calculate_pricing
from app.services.model_pricing.rule_service import normalize_provider, resolve_published_rule
from app.services.operation_log_service import log_model_pricing_event


SNAPSHOT_SCHEMA_VERSION = 1
FINAL_STAGES = {
    PricingSnapshotStage.PROVIDER_SYNC_COMPLETED.value,
    PricingSnapshotStage.PROVIDER_ASYNC_COMPLETED.value,
    PricingSnapshotStage.BACKFILL.value,
}


def _json_default(value: Any) -> Any:
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, Decimal):
        return str(value)
    return str(value)


def _canonical_hash(value: Mapping[str, Any]) -> str:
    raw = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=_json_default)
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _reference_at(value: datetime | None) -> datetime:
    value = value or _utcnow()
    return value if value.tzinfo else value.replace(tzinfo=timezone.utc)


def build_refund_pricing_snapshot(charge: Any) -> tuple[dict[str, Any], str]:
    snapshot = {
        "schema_version": SNAPSHOT_SCHEMA_VERSION,
        "refund": {
            "provider_cost": {
                "currency": getattr(charge, "provider_cost_currency", None) or "CNY",
                "amount": "0",
                "status": ProviderCostStatus.NOT_APPLICABLE.value,
                "reason": "user_credit_refund_does_not_reverse_provider_cost",
            },
            "original_charge": {
                "credit_record_id": getattr(charge, "id", None),
                "pricing_rule_id": getattr(charge, "pricing_rule_id", None),
                "pricing_version_code": getattr(charge, "pricing_version_code", None),
                "pricing_snapshot_hash": getattr(charge, "pricing_snapshot_hash", None),
                "provider_cost_amount": str(getattr(charge, "provider_cost_amount", None) or 0),
                "provider_cost_status": getattr(charge, "provider_cost_status", None),
            },
        },
    }
    return snapshot, _canonical_hash(snapshot)


def _rule_snapshot(rule: ModelPricingRule) -> dict[str, Any]:
    return {
        "id": rule.id,
        "provider": rule.provider,
        "model_name": rule.model_name,
        "model_category": rule.model_category,
        "billing_mode": rule.billing_mode,
        "calculator_version": rule.calculator_version,
        "version_code": rule.version_code,
        "effective_from": rule.effective_from.isoformat() if rule.effective_from else None,
        "effective_to": rule.effective_to.isoformat() if rule.effective_to else None,
        "currency": rule.currency,
        "rule_schema_version": rule.rule_schema_version,
        "rule_content_hash": rule.rule_content_hash,
        "rule_json": deepcopy(rule.rule_json or {}),
        "source_url": rule.source_url,
        "source_updated_at": rule.source_updated_at.isoformat() if rule.source_updated_at else None,
    }


def _apply_rule_fields(target: Any, rule: ModelPricingRule, reference_at: datetime) -> None:
    target.pricing_rule_id = rule.id
    target.pricing_version_code = rule.version_code
    target.pricing_billing_mode = rule.billing_mode
    target.pricing_calculator_version = rule.calculator_version
    target.pricing_reference_at = reference_at
    target.pricing_effective_from = rule.effective_from
    target.pricing_effective_to = rule.effective_to
    target.pricing_snapshot_schema_version = SNAPSHOT_SCHEMA_VERSION
    target.provider_cost_currency = rule.currency


def _build_pricing_snapshot(
    *,
    rule: ModelPricingRule,
    result: Any,
    stage: str,
    audit_metadata: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    # calculated_at 不参与 hash；运行时间单独保存在平铺字段中，保证相同规则/用量快照 hash 稳定。
    snapshot = {
        "schema_version": SNAPSHOT_SCHEMA_VERSION,
        "stage": stage,
        "rule": _rule_snapshot(rule),
        "calculation": deepcopy(result.breakdown),
        "provider_cost": {
            "currency": result.currency,
            "amount": str(result.amount),
            "status": ProviderCostStatus.ESTIMATED.value if result.is_estimated else ProviderCostStatus.CALCULATED.value,
            "is_estimated": bool(result.is_estimated),
            "usage_source": result.usage_source,
        },
    }
    if audit_metadata:
        snapshot["backfill"] = deepcopy(dict(audit_metadata))
    return snapshot


def _build_status_snapshot(
    *,
    status: str,
    stage: str,
    rule: ModelPricingRule | None = None,
    reason: str | None = None,
    audit_metadata: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    snapshot: dict[str, Any] = {
        "schema_version": SNAPSHOT_SCHEMA_VERSION,
        "stage": stage,
        "provider_cost": {"status": status},
    }
    if rule is not None:
        snapshot["rule"] = _rule_snapshot(rule)
        snapshot["provider_cost"]["currency"] = rule.currency
    if reason:
        snapshot["provider_cost"]["reason"] = reason
    if audit_metadata:
        snapshot["backfill"] = deepcopy(dict(audit_metadata))
    return snapshot


def _apply_result(
    target: Any,
    *,
    rule: ModelPricingRule,
    usage: Mapping[str, Any],
    result: Any,
    stage: str,
    audit_metadata: Mapping[str, Any] | None = None,
) -> None:
    now = _utcnow()
    status = ProviderCostStatus.ESTIMATED.value if result.is_estimated else ProviderCostStatus.CALCULATED.value
    pricing_snapshot = _build_pricing_snapshot(
        rule=rule,
        result=result,
        stage=stage,
        audit_metadata=audit_metadata,
    )
    target.provider_cost_amount = result.amount
    target.provider_cost_status = status
    target.provider_cost_is_estimated = bool(result.is_estimated)
    target.provider_cost_calculated_at = now
    target.provider_cost_finalized_at = now if stage in FINAL_STAGES else None
    target.pricing_usage_source = result.usage_source
    target.pricing_snapshot_json = pricing_snapshot
    target.usage_snapshot_json = deepcopy(dict(usage))
    target.pricing_snapshot_hash = _canonical_hash(pricing_snapshot)


def _can_calculate(billing_mode: str, usage: Mapping[str, Any]) -> bool:
    if billing_mode == "text_token_tiered":
        return any(int(usage.get(k) or 0) > 0 for k in ("input_tokens", "output_tokens", "cached_input_tokens", "audio_input_tokens"))
    if billing_mode in {"image_per_output", "image_input_output_tiered"}:
        return int(usage.get("successful_output_count") or usage.get("provider_billed_count") or 0) > 0
    if billing_mode == "video_token_rate":
        if int(usage.get("total_tokens") or 0) > 0:
            return True
        return float(usage.get("output_video_duration_seconds") or 0) > 0
    return False


def _apply_not_applicable(target: Any, usage: Mapping[str, Any], reason: str) -> None:
    target.provider_cost_status = ProviderCostStatus.NOT_APPLICABLE.value
    target.provider_cost_amount = Decimal("0")
    target.provider_cost_is_estimated = False
    target.provider_usage_primary = False
    target.usage_snapshot_json = deepcopy(dict(usage)) or None
    target.pricing_snapshot_json = {
        "schema_version": SNAPSHOT_SCHEMA_VERSION,
        "provider_cost": {
            "status": ProviderCostStatus.NOT_APPLICABLE.value,
            "amount": "0",
            "reason": reason,
        },
    }
    target.pricing_snapshot_hash = _canonical_hash(target.pricing_snapshot_json)


async def enrich_credit_meta_with_pricing(
    db: AsyncSession,
    *,
    meta: Any,
    usage: Mapping[str, Any] | None = None,
    reference_at: datetime | None = None,
    final: bool = False,
) -> Any:
    usage_dict = deepcopy(dict(usage or {}))
    reference = _reference_at(reference_at)

    if getattr(meta, "charge_kind", None) in {"file_parse", "vision_input"}:
        meta.pricing_reference_at = reference
        _apply_not_applicable(meta, usage_dict, "cost_included_in_primary_text_prompt_charge")
        return meta

    provider = getattr(meta, "engine_provider", None)
    model_name = getattr(meta, "engine_model_name", None)
    if not provider or not model_name:
        meta.pricing_reference_at = reference
        meta.provider_cost_status = ProviderCostStatus.PENDING.value if not final else ProviderCostStatus.HISTORICAL_ENGINE_UNAVAILABLE.value
        meta.usage_snapshot_json = usage_dict or None
        return meta

    rule = await resolve_published_rule(db, provider=provider, model_name=model_name, reference_at=reference)
    if not rule:
        meta.pricing_reference_at = reference
        meta.provider_cost_status = ProviderCostStatus.UNMATCHED_RULE.value
        meta.usage_snapshot_json = usage_dict or None
        return meta

    _apply_rule_fields(meta, rule, reference)
    meta.provider_usage_primary = bool(usage_dict.get("provider_usage_primary", True))
    meta.usage_snapshot_json = usage_dict or None

    # 媒体扣费创建阶段只锁定规则与请求快照。图片等待同步生成响应，视频等待异步 Provider 完成；
    # 不能在请求时用预设时长/分辨率提前写入估算成本。
    if not final:
        meta.provider_cost_status = ProviderCostStatus.PENDING.value
        meta.provider_cost_amount = None
        meta.provider_cost_is_estimated = False
        meta.pricing_snapshot_json = {
            "schema_version": SNAPSHOT_SCHEMA_VERSION,
            "stage": PricingSnapshotStage.REQUEST_LOCKED.value,
            "rule": _rule_snapshot(rule),
            "provider_cost": {"currency": rule.currency, "status": meta.provider_cost_status},
        }
        meta.pricing_snapshot_hash = _canonical_hash(meta.pricing_snapshot_json)
        return meta

    if not _can_calculate(rule.billing_mode, usage_dict):
        meta.provider_cost_status = ProviderCostStatus.USAGE_MISSING.value
        meta.pricing_snapshot_json = {
            "schema_version": SNAPSHOT_SCHEMA_VERSION,
            "stage": PricingSnapshotStage.PROVIDER_SYNC_COMPLETED.value,
            "rule": _rule_snapshot(rule),
            "provider_cost": {"currency": rule.currency, "status": meta.provider_cost_status},
        }
        meta.pricing_snapshot_hash = _canonical_hash(meta.pricing_snapshot_json)
        return meta

    try:
        result = calculate_pricing(
            billing_mode=rule.billing_mode,
            calculator_version=rule.calculator_version,
            rule_json=rule.rule_json or {},
            usage=usage_dict,
            currency=rule.currency,
        )
        _apply_result(
            meta,
            rule=rule,
            usage=usage_dict,
            result=result,
            stage=PricingSnapshotStage.PROVIDER_SYNC_COMPLETED.value if final else PricingSnapshotStage.REQUEST_LOCKED.value,
        )
    except PricingCalculationError:
        meta.provider_cost_status = ProviderCostStatus.USAGE_MISSING.value if final else ProviderCostStatus.PENDING.value
    except Exception as exc:
        meta.provider_cost_status = ProviderCostStatus.ERROR.value
        log_model_pricing_event(
            event_type="pricing_cost_calculate",
            event_status="failed",
            owner_type=getattr(meta, "owner_type", None),
            owner_id=getattr(meta, "owner_id", None),
            pricing_rule_id=rule.id,
            pricing_version=rule.version_code,
            provider=provider,
            model_name=model_name,
            billing_mode=rule.billing_mode,
            cost_status=meta.provider_cost_status,
            error=str(exc),
        )
    return meta


async def _load_locked_rule(
    db: AsyncSession,
    charge: CreditRecord,
    *,
    reference_at: datetime,
    use_locked_rule: bool,
) -> ModelPricingRule | None:
    if use_locked_rule and charge.pricing_rule_id:
        return (
            await db.execute(select(ModelPricingRule).where(ModelPricingRule.id == charge.pricing_rule_id).limit(1))
        ).scalar_one_or_none()

    if not use_locked_rule:
        provider = normalize_provider(charge.engine_provider, charge.engine_model_name)
        model_name = str(charge.engine_model_name or "").strip()
        if not provider or not model_name:
            return None
        return (
            await db.execute(
                select(ModelPricingRule)
                .where(ModelPricingRule.provider == provider)
                .where(ModelPricingRule.model_name == model_name)
                .where(ModelPricingRule.publish_status == ModelPricingRuleStatus.PUBLISHED.value)
                .where(ModelPricingRule.effective_from <= reference_at)
                .where(
                    (ModelPricingRule.effective_to.is_(None))
                    | (ModelPricingRule.effective_to > reference_at)
                )
                .order_by(ModelPricingRule.effective_from.desc(), ModelPricingRule.created_at.desc())
                .limit(1)
            )
        ).scalar_one_or_none()

    return await resolve_published_rule(
        db,
        provider=charge.engine_provider,
        model_name=charge.engine_model_name,
        reference_at=reference_at,
    )


async def _missing_rule_status(db: AsyncSession, *, charge: CreditRecord, reference_at: datetime) -> str:
    model_name = str(charge.engine_model_name or "").strip()
    provider = normalize_provider(charge.engine_provider, model_name)
    if not provider or not model_name:
        return ProviderCostStatus.HISTORICAL_ENGINE_UNAVAILABLE.value
    earliest = (
        await db.execute(
            select(func.min(ModelPricingRule.effective_from)).where(
                ModelPricingRule.provider == provider,
                ModelPricingRule.model_name == model_name,
            )
        )
    ).scalar_one_or_none()
    if earliest and _reference_at(reference_at) < _reference_at(earliest):
        return ProviderCostStatus.HISTORICAL_PRICE_UNAVAILABLE.value
    return ProviderCostStatus.UNMATCHED_RULE.value


def _merge_snapshot(existing: Mapping[str, Any] | None, incoming: Mapping[str, Any] | None) -> dict[str, Any] | None:
    if incoming is None:
        return deepcopy(dict(existing or {})) or None
    merged = deepcopy(dict(existing or {}))
    merged.update(deepcopy(dict(incoming)))
    return merged


async def finalize_credit_record_pricing(
    db: AsyncSession,
    *,
    charge: CreditRecord,
    usage: Mapping[str, Any] | None,
    stage: str,
    attachment_snapshot: Mapping[str, Any] | None = None,
    attachment_counts: Mapping[str, Any] | None = None,
    generation_snapshot: Mapping[str, Any] | None = None,
    generation_counts: Mapping[str, Any] | None = None,
    allow_upgrade_estimated: bool = True,
    pricing_reference_at: datetime | None = None,
    use_locked_rule: bool = True,
    force_reprice: bool = False,
    backfill_metadata: Mapping[str, Any] | None = None,
) -> CreditRecord:
    """同一事务内回填，不 commit；JSON 一律构建新对象后整体赋值。

    正常生成链路保持默认行为：使用请求时已锁定的规则，已核算成本不可覆盖。
    历史补录可显式传入统一的当前计价时点、忽略旧规则绑定并强制重算。
    """
    if attachment_snapshot is not None:
        charge.attachment_snapshot_json = deepcopy(dict(attachment_snapshot))
    for key, value in (attachment_counts or {}).items():
        if hasattr(charge, key):
            setattr(charge, key, value)
    if generation_snapshot is not None:
        charge.generation_snapshot_json = _merge_snapshot(charge.generation_snapshot_json, generation_snapshot)
    for key, value in (generation_counts or {}).items():
        if hasattr(charge, key):
            setattr(charge, key, value)

    # 资源下载完成只补资源快照；同步图片/异步视频成本均不得在下载阶段重算。
    if stage == PricingSnapshotStage.RESOURCE_DOWNLOAD_COMPLETED.value:
        return charge

    if charge.type != "consume" or charge.charge_action != "charge":
        _apply_not_applicable(charge, dict(usage or {}), "only_consume_charge_can_be_priced")
        return charge

    current_status = charge.provider_cost_status
    if not force_reprice:
        if current_status == ProviderCostStatus.CALCULATED.value:
            return charge
        if current_status == ProviderCostStatus.ESTIMATED.value and not allow_upgrade_estimated:
            return charge

    usage_dict = deepcopy(dict(usage or {}))
    reference = _reference_at(
        pricing_reference_at
        if pricing_reference_at is not None
        else (charge.pricing_reference_at or charge.created_at)
    )
    rule = await _load_locked_rule(
        db,
        charge,
        reference_at=reference,
        use_locked_rule=use_locked_rule,
    )
    if not rule:
        charge.pricing_reference_at = reference
        if not charge.engine_provider or not charge.engine_model_name:
            status = ProviderCostStatus.HISTORICAL_ENGINE_UNAVAILABLE.value
        elif use_locked_rule:
            status = await _missing_rule_status(db, charge=charge, reference_at=reference)
        else:
            status = ProviderCostStatus.UNMATCHED_RULE.value
        charge.provider_cost_status = status
        charge.provider_cost_amount = None
        charge.provider_cost_is_estimated = False
        charge.provider_cost_calculated_at = None
        charge.provider_cost_finalized_at = None
        charge.usage_snapshot_json = usage_dict or None
        charge.pricing_snapshot_json = _build_status_snapshot(
            status=status,
            stage=stage,
            reason="current_published_rule_not_found" if not use_locked_rule else "pricing_rule_not_found",
            audit_metadata=backfill_metadata,
        )
        charge.pricing_snapshot_hash = _canonical_hash(charge.pricing_snapshot_json)
        return charge

    _apply_rule_fields(charge, rule, reference)
    charge.provider_usage_primary = bool(usage_dict.get("provider_usage_primary", True))
    if not _can_calculate(rule.billing_mode, usage_dict):
        charge.provider_cost_status = (
            ProviderCostStatus.USAGE_MISSING.value
            if stage in FINAL_STAGES
            else ProviderCostStatus.PENDING.value
        )
        charge.provider_cost_amount = None
        charge.provider_cost_is_estimated = False
        charge.provider_cost_calculated_at = None
        charge.provider_cost_finalized_at = None
        charge.usage_snapshot_json = usage_dict or None
        charge.pricing_snapshot_json = _build_status_snapshot(
            status=charge.provider_cost_status,
            stage=stage,
            rule=rule,
            reason="pricing_usage_missing",
            audit_metadata=backfill_metadata,
        )
        charge.pricing_snapshot_hash = _canonical_hash(charge.pricing_snapshot_json)
        return charge

    try:
        result = calculate_pricing(
            billing_mode=rule.billing_mode,
            calculator_version=rule.calculator_version,
            rule_json=rule.rule_json or {},
            usage=usage_dict,
            currency=rule.currency,
        )
        if (
            not force_reprice
            and current_status == ProviderCostStatus.ESTIMATED.value
            and result.is_estimated
        ):
            return charge
        _apply_result(
            charge,
            rule=rule,
            usage=usage_dict,
            result=result,
            stage=stage,
            audit_metadata=backfill_metadata,
        )
        log_model_pricing_event(
            event_type="pricing_snapshot_persist",
            user_id=charge.user_id,
            credit_record_id=charge.id,
            owner_type=charge.owner_type,
            owner_id=charge.owner_id,
            pricing_rule_id=rule.id,
            pricing_version=rule.version_code,
            provider=charge.engine_provider,
            model_name=charge.engine_model_name,
            billing_mode=rule.billing_mode,
            cost_status=charge.provider_cost_status,
            provider_cost=charge.provider_cost_amount,
            is_estimated=charge.provider_cost_is_estimated,
            detail={
                "stage": stage,
                "usage_source": charge.pricing_usage_source,
                "backfill": deepcopy(dict(backfill_metadata or {})) or None,
            },
        )
    except PricingCalculationError as exc:
        charge.provider_cost_status = ProviderCostStatus.USAGE_MISSING.value
        charge.provider_cost_amount = None
        charge.provider_cost_is_estimated = False
        charge.provider_cost_calculated_at = None
        charge.provider_cost_finalized_at = None
        charge.usage_snapshot_json = usage_dict or None
        charge.pricing_snapshot_json = _build_status_snapshot(
            status=charge.provider_cost_status,
            stage=stage,
            rule=rule,
            reason=str(exc),
            audit_metadata=backfill_metadata,
        )
        charge.pricing_snapshot_hash = _canonical_hash(charge.pricing_snapshot_json)
        log_model_pricing_event(
            event_type="pricing_snapshot_failed",
            event_status="warning",
            user_id=charge.user_id,
            credit_record_id=charge.id,
            owner_type=charge.owner_type,
            owner_id=charge.owner_id,
            pricing_rule_id=rule.id,
            pricing_version=rule.version_code,
            provider=charge.engine_provider,
            model_name=charge.engine_model_name,
            billing_mode=rule.billing_mode,
            cost_status=charge.provider_cost_status,
            error=str(exc),
            detail={"stage": stage, "backfill": deepcopy(dict(backfill_metadata or {})) or None},
        )
    except Exception as exc:
        charge.provider_cost_status = ProviderCostStatus.ERROR.value
        charge.provider_cost_amount = None
        charge.provider_cost_is_estimated = False
        charge.provider_cost_calculated_at = None
        charge.provider_cost_finalized_at = None
        charge.usage_snapshot_json = usage_dict or None
        charge.pricing_snapshot_json = _build_status_snapshot(
            status=charge.provider_cost_status,
            stage=stage,
            rule=rule,
            reason=str(exc),
            audit_metadata=backfill_metadata,
        )
        charge.pricing_snapshot_hash = _canonical_hash(charge.pricing_snapshot_json)
        log_model_pricing_event(
            event_type="pricing_snapshot_failed",
            event_status="failed",
            user_id=charge.user_id,
            credit_record_id=charge.id,
            owner_type=charge.owner_type,
            owner_id=charge.owner_id,
            pricing_rule_id=rule.id,
            pricing_version=rule.version_code,
            provider=charge.engine_provider,
            model_name=charge.engine_model_name,
            billing_mode=rule.billing_mode,
            cost_status=charge.provider_cost_status,
            error=str(exc),
            detail={"stage": stage, "backfill": deepcopy(dict(backfill_metadata or {})) or None},
        )
    return charge