from __future__ import annotations

from dataclasses import dataclass, asdict
from typing import Any, Mapping

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

from app.models.credit_record import CreditRecord
from app.models.generation_record import GenerationRecord
from app.models.system_config import SystemConfig
from app.models.user import User
from app.services.credits import calc_image_credits, calc_text_credits, calc_video_credits
from app.utils.exceptions import InsufficientCreditsError
from app.utils.id_gen import generate_id


CHARGE_TEXT_PROMPT = "CHATAPI_TEXT_PROMPT"
CHARGE_FILE_PARSE = "CHATAPI_FILE_PARSE"
CHARGE_VISION_INPUT = "CHATAPI_VISION_INPUT"
CHARGE_MEDIA_IMAGE = "CHATAPI_MEDIA_IMAGE"
CHARGE_MEDIA_VIDEO = "CHATAPI_MEDIA_VIDEO"


@dataclass
class BillingItem:
    charge_key: str
    amount: float
    charged: bool
    skipped_reason: str | None = None


@dataclass
class BillingSummary:
    record_id: str
    user_id: str
    items: list[BillingItem]

    @property
    def total_charged(self) -> float:
        return round(sum(item.amount for item in self.items if item.charged), 2)

    def get_amount(self, charge_key: str) -> float:
        return round(sum(item.amount for item in self.items if item.charge_key == charge_key and item.charged), 2)

    def to_dict(self) -> dict[str, Any]:
        data = asdict(self)
        data["total_charged"] = self.total_charged
        return data


def _round2(value: float | int | None) -> float:
    return round(float(value or 0), 2)


def _safe_int(value: Any, default: int = 0) -> int:
    try:
        if value is None:
            return default
        return int(value)
    except Exception:
        return default


async def _get_config_float_or_none(db: AsyncSession, key: str) -> float | None:
    result = await db.execute(select(SystemConfig).where(SystemConfig.key == key).limit(1))
    config = result.scalar_one_or_none()
    if not config:
        return None
    try:
        return float(config.value)
    except Exception:
        return None


async def _calc_optional_token_credits(db: AsyncSession, tokens: int, config_key: str) -> float:
    """Calculate optional token billing. Missing config means do not charge.

    This prevents double-charging existing projects where uploaded file/OCR/vision
    content is already included in the LLM provider's input_tokens.
    """
    tokens = _safe_int(tokens)
    if tokens <= 0:
        return 0.0
    rate = await _get_config_float_or_none(db, config_key)
    if rate is None:
        return 0.0
    return round(tokens * rate / 1000, 2)


def _legacy_description_keywords(charge_key: str) -> list[str]:
    # Compatibility with old patch/original project records that were inserted
    # before this safe billing service added [CHARGE_KEY] prefixes.
    if charge_key == CHARGE_TEXT_PROMPT:
        return ["ChatAPI提示词整理", "提示词优化"]
    if charge_key == CHARGE_MEDIA_IMAGE:
        return ["ChatAPI异步图片生成", "图片生成"]
    if charge_key == CHARGE_MEDIA_VIDEO:
        return ["ChatAPI异步视频生成", "视频生成"]
    if charge_key == CHARGE_FILE_PARSE:
        return ["文件解析Token"]
    if charge_key == CHARGE_VISION_INPUT:
        return ["图片理解Token"]
    return []


async def _find_existing_charge(db: AsyncSession, user_id: str, related_id: str, charge_key: str) -> CreditRecord | None:
    base = (
        select(CreditRecord)
        .where(CreditRecord.user_id == user_id)
        .where(CreditRecord.related_id == related_id)
        .where(CreditRecord.type == "consume")
    )

    result = await db.execute(base.where(CreditRecord.description.like(f"[{charge_key}]%")).limit(1))
    existing = result.scalar_one_or_none()
    if existing:
        return existing

    for keyword in _legacy_description_keywords(charge_key):
        result = await db.execute(base.where(CreditRecord.description.like(f"%{keyword}%")).limit(1))
        existing = result.scalar_one_or_none()
        if existing:
            return existing
    return None


async def deduct_credits_locked_once(
    db: AsyncSession,
    *,
    user_id: str,
    amount: float,
    description: str,
    related_id: str,
    charge_key: str,
) -> BillingItem:
    """Deduct credits with row lock and idempotency.

    - User row is locked by SELECT ... FOR UPDATE, so concurrent deductions for
      the same user are serialized in PostgreSQL/MySQL.
    - CreditRecord description prefix + related_id is used as an idempotency key
      without changing existing table structures.
    """
    amount = _round2(amount)
    if amount <= 0:
        return BillingItem(charge_key=charge_key, amount=0.0, charged=False, skipped_reason="amount_lte_zero")

    result = await db.execute(select(User).where(User.id == user_id).with_for_update())
    user = result.scalar_one_or_none()
    if not user:
        raise ValueError("User not found")

    existing_charge = await _find_existing_charge(db, user_id, related_id, charge_key)
    if existing_charge:
        return BillingItem(
            charge_key=charge_key,
            amount=abs(_round2(existing_charge.amount)),
            charged=False,
            skipped_reason="already_charged",
        )

    if float(user.credits or 0) < amount:
        raise InsufficientCreditsError()

    user.credits = round(float(user.credits or 0) - amount, 2)
    db.add(
        CreditRecord(
            id=generate_id(),
            user_id=user_id,
            type="consume",
            amount=-amount,
            balance_after=user.credits,
            description=description,
            related_id=related_id,
        )
    )
    await db.flush()
    return BillingItem(charge_key=charge_key, amount=amount, charged=True)


async def charge_chatapi_prompt_usage(
    db: AsyncSession,
    *,
    record: GenerationRecord,
    usage: Mapping[str, Any],
    project_name: str | None = None,
) -> BillingSummary:
    """Charge ChatAPI prompt optimization and optional uploaded-file/vision tokens.

    file_parse_credits / vision_input_credits are optional and disabled unless
    SystemConfig contains these keys:
    - file_parse_credits_per_1000_tokens
    - vision_input_credits_per_1000_tokens
    """
    project_name = project_name or "AI生成任务"
    items: list[BillingItem] = []

    input_tokens = _safe_int(usage.get("input_tokens"))
    output_tokens = _safe_int(usage.get("output_tokens"))
    text_credits = await calc_text_credits(db, input_tokens, output_tokens)
    items.append(
        await deduct_credits_locked_once(
            db,
            user_id=record.user_id,
            amount=text_credits,
            description=f"ChatAPI提示词整理",
            related_id=record.id,
            charge_key=CHARGE_TEXT_PROMPT,
        )
    )

    file_tokens = usage.get("file_parse_tokens") or usage.get("file_tokens") or usage.get("document_tokens") or 0
    file_parse_credits = await _calc_optional_token_credits(
        db,
        _safe_int(file_tokens),
        "file_parse_credits_per_1000_tokens",
    )
    items.append(
        await deduct_credits_locked_once(
            db,
            user_id=record.user_id,
            amount=file_parse_credits,
            description=f"文件解析Token",
            related_id=record.id,
            charge_key=CHARGE_FILE_PARSE,
        )
    )

    vision_tokens = usage.get("vision_input_tokens") or usage.get("image_input_tokens") or usage.get("image_tokens") or 0
    vision_input_credits = await _calc_optional_token_credits(
        db,
        _safe_int(vision_tokens),
        "vision_input_credits_per_1000_tokens",
    )
    items.append(
        await deduct_credits_locked_once(
            db,
            user_id=record.user_id,
            amount=vision_input_credits,
            description=f"图片理解Token",
            related_id=record.id,
            charge_key=CHARGE_VISION_INPUT,
        )
    )

    if hasattr(record, "text_credits_cost"):
        # Store expected text-side cost, even when this task is a retry and the
        # actual CreditRecord was already written by an earlier attempt.
        record.text_credits_cost = round(text_credits + file_parse_credits + vision_input_credits, 2)
    if hasattr(record, "text_tokens_used"):
        record.text_tokens_used = _safe_int(usage.get("total_tokens"), input_tokens + output_tokens)

    return BillingSummary(record_id=record.id, user_id=record.user_id, items=items)


async def charge_generation_media_by_params(
    db: AsyncSession,
    *,
    user_id: str,
    record_id: str,
    gen_type: str,
    image_size: str | None = None,
    duration: int | None = None,
    resolution: str | None = None,
    project_name: str | None = None,
    description_prefix: str = "ChatAPI异步",
) -> BillingSummary:
    """Charge image/video generation fee safely before creating provider task."""
    project_name = project_name or "AI生成任务"
    gen_type = (gen_type or "").lower().strip()
    items: list[BillingItem] = []

    if gen_type == "image":
        size = image_size or "2K"
        amount = await calc_image_credits(db, size)
        items.append(
            await deduct_credits_locked_once(
                db,
                user_id=user_id,
                amount=amount,
                description=f"{description_prefix}图片生成",
                related_id=record_id,
                charge_key=CHARGE_MEDIA_IMAGE,
            )
        )
    elif gen_type == "video":
        amount = await calc_video_credits(db, duration or 5, resolution or "720p")
        items.append(
            await deduct_credits_locked_once(
                db,
                user_id=user_id,
                amount=amount,
                description=f"{description_prefix}视频生成",
                related_id=record_id,
                charge_key=CHARGE_MEDIA_VIDEO,
            )
        )
    else:
        raise ValueError(f"不支持的生成类型: {gen_type}")

    return BillingSummary(record_id=record_id, user_id=user_id, items=items)


async def charge_generation_media_for_record(
    db: AsyncSession,
    *,
    record: GenerationRecord,
    project_name: str | None = None,
    description_prefix: str = "ChatAPI异步",
) -> BillingSummary:
    return await charge_generation_media_by_params(
        db,
        user_id=record.user_id,
        record_id=record.id,
        gen_type=record.gen_type,
        image_size=record.image_size,
        duration=record.duration,
        resolution=record.resolution,
        project_name=project_name,
        description_prefix=description_prefix,
    )