"""用量记账(0006 + 0007):一次产生成本的调用 = 一行 usage_events + 双写 messages 列。

chat 类型的入口由 loop.py 在 assistant message 入库后调用;媒体工具(image/video/audio)
在 tool execute 完后由 tool 直接调用对应入口(record_image_usage 等)。

币种(0007):全表统一 CNY(`cost_cny` 列)。chat 路径走 litellm 的 USD cost_map → 内部
×USD_TO_CNY 折算落库;媒体路径价格本身就是 CNY,直接落。units jsonb 里 snapshot 当时
的关键价格参数(chat 没有,media 存 price_cny_per_image 等),便于跨调价对账。
"""
from __future__ import annotations

from datetime import datetime
from decimal import Decimal
from typing import Any, Mapping, Optional
from uuid import UUID

from sqlalchemy import func, select, update

from .engine import session_scope
from .models import Message, UsageEvent


# litellm 的 cost map 给的是 USD,落库前折成 CNY。汇率近似(每年看一次,实质偏差不大);
# 真要精算的话应该按调用时刻的汇率,但开发期/个人用接受。
USD_TO_CNY = Decimal("7.2")


def _safe_chat_cost_usd(response: Any) -> Decimal:
    """litellm.completion_cost(response) 包一层:任何异常都吞掉返 0。

    未知 model / cost map 没收录 / response 结构变都不影响主流程 —— usage_events
    仍写入,只是 cost=0,后续人工补算 OK。返 USD,由 caller 折算。
    """
    try:
        from litellm import completion_cost  # type: ignore[import-not-found]
        cost = completion_cost(completion_response=response)
        if cost is None:
            return Decimal("0")
        return Decimal(str(cost))
    except Exception:
        return Decimal("0")


def _fallback_chat_cost_cny(
    *,
    prompt_tokens: int,
    completion_tokens: int,
    input_cny_per_mtoken: float,
    output_cny_per_mtoken: float,
    cache_hit_tokens: int = 0,
    cache_hit_cny_per_mtoken: float = 0.0,
) -> Decimal:
    """按本地模型档案价格兜底计算 chat 成本(CNY)。

    `prompt_tokens` = 全部输入(含缓存命中部分)。DeepSeek 等对前缀缓存命中按更低
    单价计费,这里把输入拆成「命中 / 未命中」两段:命中段按 `cache_hit_cny_per_mtoken`,
    其余按 `input_cny_per_mtoken`。`cache_hit_cny_per_mtoken<=0` 时不区分,全按 input 价
    (老行为,绝不少记)。
    """
    input_price = Decimal(str(input_cny_per_mtoken or 0))
    output_price = Decimal(str(output_cny_per_mtoken or 0))
    hit_price = Decimal(str(cache_hit_cny_per_mtoken or 0))
    if hit_price <= 0:
        hit_price = input_price  # 未配缓存价 → 命中段按 input 全价(安全兜底)

    pt = int(prompt_tokens)
    hit = max(0, min(int(cache_hit_tokens), pt))  # clamp:命中数不超过总输入
    miss = pt - hit
    ct = int(completion_tokens)

    cost = (
        Decimal(str(miss)) * input_price / Decimal("1000000")
        + Decimal(str(hit)) * hit_price / Decimal("1000000")
        + Decimal(str(ct)) * output_price / Decimal("1000000")
    )
    return cost.quantize(Decimal("0.000001"))


def record_chat_usage(
    *,
    task_id: UUID,
    user_id: UUID,
    message_id: Optional[UUID],
    model_profile: str,
    prompt_tokens: int,
    completion_tokens: int,
    input_cny_per_mtoken: float = 0.0,
    output_cny_per_mtoken: float = 0.0,
    cache_hit_tokens: int = 0,
    cache_hit_cny_per_mtoken: float = 0.0,
    extra_units: Optional[Mapping[str, Any]] = None,
    response: Any = None,
    kind: str = "chat",
) -> Decimal:
    """记一次 chat 调用:写 usage_events 行 + 回填 messages.model_profile/tokens_in/out。

    `message_id` 来自 `Session.append` 的返回值;若为 None(系统消息 / 旧路径未拿到)
    则 usage_events 仍写但 message_id=NULL,messages 列不回填。
    `model_profile` 形如 `"deepseek_v4.pro"`(family.variant)。
    `kind` 默认 "chat";chat 形态的辅助调用可传别的标签(如 "context_fold"),
    按 kind GROUP BY 对账 + 派生统计(折叠次数 = count)—— prompt_optimize 同范式。
    返回算出的 cost_cny(已落库),调用方可用作 SSE 显示。
    """
    cost_usd = _safe_chat_cost_usd(response)
    cost_cny = (cost_usd * USD_TO_CNY).quantize(Decimal("0.000001"))
    if cost_cny == 0 and (input_cny_per_mtoken or output_cny_per_mtoken):
        cost_cny = _fallback_chat_cost_cny(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            input_cny_per_mtoken=input_cny_per_mtoken,
            output_cny_per_mtoken=output_cny_per_mtoken,
            cache_hit_tokens=cache_hit_tokens,
            cache_hit_cny_per_mtoken=cache_hit_cny_per_mtoken,
        )
    units = {
        "tokens_in": int(prompt_tokens),
        "tokens_out": int(completion_tokens),
        # snapshot 折算系数,便于历史对账(汇率/价格涨跌后仍能还原当时折算逻辑)
        "usd_to_cny": float(USD_TO_CNY),
    }
    if input_cny_per_mtoken or output_cny_per_mtoken:
        units.update({
            "input_cny_per_mtoken": float(input_cny_per_mtoken or 0),
            "output_cny_per_mtoken": float(output_cny_per_mtoken or 0),
        })
    if extra_units:
        units.update(extra_units)

    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=message_id,
            kind=kind,
            model_profile=model_profile,
            units=units,
            cost_cny=cost_cny,
        ))
        if message_id is not None:
            s.execute(
                update(Message)
                .where(Message.message_id == message_id)
                .values(
                    tokens_in=int(prompt_tokens),
                    tokens_out=int(completion_tokens),
                    model_profile=model_profile,
                )
            )
    return cost_cny


def record_malformed_tool_call(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    tool: str,
    arg_len: int,
    error: str,
    head: str,
    tail: str,
    tokens_in: int = 0,
    tokens_out: int = 0,
) -> None:
    """记一次被丢弃的畸形 tool_call(kind=tool_malformed)。

    这不是记账而是留痕:畸形轮整轮丢弃、messages 无痕,此行是 admin「工具失败聚集」
    面板与巡检邮件唯一的数据源(core/toolfail.py 第二段扫描按 units->>'tool'/'err'
    聚合)。cost_cny 恒 0 —— cost 统计全 kind 合计且随任务展示,provider 抖动的浪费
    不该算成用户花销;该轮真实 token 快照进 units,将来要算浪费成本可从 units 反推
    (token 汇总只算 kind=chat,此处 tokens_in/out 不会混进 token 统计)。
    """
    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # 畸形轮不入 messages,无可关联
            kind="tool_malformed",
            model_profile=model_profile,
            units={
                "tool": tool,
                "len": int(arg_len),
                "err": str(error)[:200],
                "head": head,
                "tail": tail,
                "tokens_in": int(tokens_in),
                "tokens_out": int(tokens_out),
            },
            cost_cny=Decimal("0"),
        ))


def record_salvaged_tool_call(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    tool: str,
    arg_len: int,
    salvaged_len: int,
    head: str,
) -> None:
    """记一次从畸形 arguments 里抢救成功的 tool_call(kind=tool_salvaged)。

    与 record_malformed_tool_call 是一对:畸形轮本会整轮丢弃 + 非流式重 roll,salvage
    命中时改成就地抠出尾部完好 JSON、当轮照常执行,省掉一次重试。此行是留痕(cost 恒 0,
    tokens 不入 kind=chat 汇总),供「工具失败聚集」面板 / 事后核对区分「真丢弃」与「抢救回」:
    arg_len=损坏前缀+JSON 的总长,salvaged_len=抠出 JSON 的长,差值即被丢弃的 wire 垃圾前缀长。
    head 存前 300 字(过 ascii 转义,消费端编码不可控)供核验前缀形态是否仍是 char-0 型。
    """
    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # 抢救出的 tool_call 正常执行并入 messages,但此留痕行不与之关联
            kind="tool_salvaged",
            model_profile=model_profile,
            units={
                "tool": tool,
                "len": int(arg_len),
                "salvaged_len": int(salvaged_len),
                "head": head,
            },
            cost_cny=Decimal("0"),
        ))


def record_empty_response(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    attempt: int,
    tokens_in: int = 0,
    tokens_out: int = 0,
) -> None:
    """记一次 provider 吐空(kind=empty_response):assistant 轮既无 tool_calls 又无正文。

    背景(task 2a1bc25d 案):unifyllm 网关对某档 Claude 偶发把 tool_use 漏成正文 / 直接吐空,
    run loop 见 tool_calls 空即当「模型答完」静默 done —— 无报错、run_status=idle,只能人肉
    挖 DB 才发现。此行是留痕(cost 恒 0,tokens 不入 kind=chat 汇总):即使当轮非流式重试救回,
    也留一条供「工具失败聚集」面板看到「哪个模型档在吐空」(第四段扫描,tool 名固定 "(empty)")。
    attempt=第几次尝试(1=首个流式轮),tokens 快照供估算浪费。单条转瞬即逝不触发面板(阈值
    min_count/min_tasks),跨 task 系统性吐空才冒头 —— 正是要抓的网关级抖动。
    """
    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # 空响应轮不入 messages(丢弃重试),无可关联
            kind="empty_response",
            model_profile=model_profile,
            units={
                "attempt": int(attempt),
                "tokens_in": int(tokens_in),
                "tokens_out": int(tokens_out),
            },
            cost_cny=Decimal("0"),
        ))


def record_run_error(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    error: str,
) -> None:
    """记一次 run 级终态错误(kind=run_error)。

    留痕而非记账(cost 恒 0):run 在 LLM 请求层/构建期直接抛异常(RateLimitError
    余额不足、认证失败等)时,整轮没有 tool 消息,错误只写 tasks.run_error 一列 ——
    而该列只留最后一次,下次成功即被清掉。此行是「工具失败聚集」面板 + 巡检邮件
    看到 run 级错误的唯一持久数据源(core/toolfail.py 第三段扫描按 units->>'err'
    聚合)。task 2a1bc25d 案:Zai 余额不足连挂 3 次续跑,前端无提示、admin 无感知。
    """
    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # 错误轮通常无 assistant message 落库,无可关联
            kind="run_error",
            model_profile=model_profile,
            units={"err": str(error)[:500]},
            cost_cny=Decimal("0"),
        ))


def record_image_usage(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    n_images: int,
    size: str,
    price_cny_per_image: float,
    search: bool = False,
    extra_units: Optional[Mapping[str, Any]] = None,
) -> Decimal:
    """记一次图像生成:写 usage_events(kind=image)。

    单价(CNY/张)由 caller 从配置文件读出后传入,**同步 snapshot 进 units jsonb** ——
    将来豆包调价改 YAML 即可,历史记录不动且仍能完整还原当时单价。
    `model_profile` 形如 `"doubao.seedream_5"`(family.variant 风格,跟 chat 对齐)。
    `extra_units` 给将来扩展(如 quality / style 等额外加价维度)预留。
    返回 cost_cny(已落库,可作 SSE / tool 返回串显示用)。
    """
    price = Decimal(str(price_cny_per_image))
    cost_cny = (price * n_images).quantize(Decimal("0.000001"))
    units: dict[str, Any] = {
        "n_images": int(n_images),
        "size": size,
        "search": bool(search),
        "price_cny_per_image": float(price_cny_per_image),
    }
    if extra_units:
        units.update(extra_units)

    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # image tool 在 tool execute 时调用,message 还未落库
            kind="image",
            model_profile=model_profile,
            units=units,
            cost_cny=cost_cny,
        ))
    return cost_cny


def record_video_usage(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    resolution: str,
    ratio: str,
    duration_s: int,
    fps: int,
    width: int,
    height: int,
    tokens: int,
    price_cny_per_mtoken: float,
    has_video_input: bool = False,
    watermark: bool = False,
    extra_units: Optional[Mapping[str, Any]] = None,
) -> Decimal:
    """记一次视频生成:写 usage_events(kind=video)。

    成本算法:`cost_cny = tokens / 1_000_000 * price_cny_per_mtoken`。tokens 由 caller
    传入(响应里若有官方 usage 字段就用,否则按公式 `(in_dur+out_dur)*W*H*fps/1024`
    估算),`price_cny_per_mtoken` 同步 snapshot 进 units jsonb,日后调价不影响历史对账。

    `model_profile` 形如 `"doubao.seedance_2_fast"`(family.variant 风格)。
    `has_video_input=True` 表示图生/视频编辑路径(单价 22 元/Mtok);False = 文生视频(37 元/Mtok)。
    caller 自行根据请求 body 是否含 image_url/video_url 传对应单价 + flag。

    **失败任务不要走这里** —— Volcengine 失败不计费,失败的 tool 调用直接返 [Error] 不写 usage。
    """
    price = Decimal(str(price_cny_per_mtoken))
    tok = Decimal(str(int(tokens)))
    cost_cny = (price * tok / Decimal("1000000")).quantize(Decimal("0.000001"))
    units: dict[str, Any] = {
        "resolution": resolution,
        "ratio": ratio,
        "duration_s": int(duration_s),
        "fps": int(fps),
        "width": int(width),
        "height": int(height),
        "tokens": int(tokens),
        "price_cny_per_mtoken": float(price_cny_per_mtoken),
        "has_video_input": bool(has_video_input),
        "watermark": bool(watermark),
    }
    if extra_units:
        units.update(extra_units)

    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,
            kind="video",
            model_profile=model_profile,
            units=units,
            cost_cny=cost_cny,
        ))
    return cost_cny


def record_vision_usage(
    *,
    task_id: UUID,
    user_id: UUID,
    model_profile: str,
    prompt_tokens: int,
    completion_tokens: int,
    input_cny_per_mtoken: float,
    output_cny_per_mtoken: float,
    extra_units: Optional[Mapping[str, Any]] = None,
) -> Decimal:
    """记一次图像理解(看图 / OCR):写 usage_events(kind=vision)。

    vision 是 token 计费(同 chat,不是 per-call),`cost_cny = in*in_price/1e6 + out*out_price/1e6`。
    tokens 由 caller 从 ARK /chat/completions 响应的 usage 字段取(没有则传 0,cost=0 不阻塞)。
    单价(CNY/Mtok)同步 snapshot 进 units jsonb,日后调价不影响历史对账。
    `model_profile` 形如 `"doubao.seed_2_lite"`(family.variant 风格,跟 chat/image 对齐)。

    **失败任务不要走这里** —— ARK 失败不计费,失败的 tool 调用直接返 [Error] 不写 usage。
    """
    inp = Decimal(str(input_cny_per_mtoken or 0))
    out = Decimal(str(output_cny_per_mtoken or 0))
    cost_cny = (
        Decimal(str(int(prompt_tokens))) * inp / Decimal("1000000")
        + Decimal(str(int(completion_tokens))) * out / Decimal("1000000")
    ).quantize(Decimal("0.000001"))
    units: dict[str, Any] = {
        "tokens_in": int(prompt_tokens),
        "tokens_out": int(completion_tokens),
        "input_cny_per_mtoken": float(input_cny_per_mtoken or 0),
        "output_cny_per_mtoken": float(output_cny_per_mtoken or 0),
    }
    if extra_units:
        units.update(extra_units)

    with session_scope() as s:
        s.add(UsageEvent(
            user_id=user_id,
            task_id=task_id,
            message_id=None,  # vision tool 在 tool execute 时调用,message 还未落库
            kind="vision",
            model_profile=model_profile,
            units=units,
            cost_cny=cost_cny,
        ))
    return cost_cny


def check_daily_quota(*, user_id: UUID, kind: str, limit: int) -> tuple[int, bool]:
    """每账号每日 kind=image/video 调用配额检查。返回 (今日已用次数, 是否超额)。

    `limit <= 0` 视为不限,直接返回 (0, False) 不查 DB。
    "一天"=服务器本地 00:00 起算(用户直觉日历日;非滑动 24h 窗口)。
    失败任务不计 —— `record_*_usage` 只在成功+下载完才落库 → 失败 retry 不烧配额。
    跨 task 跨 variant 全口径合计 —— 配额是账号级,与具体调用哪个 variant 无关。

    并发 race:同 user 两次调用同时通过 check 会两个都跑成功(off-by-one),
    可接受 —— 单 task 单活 run gate 已经挡了同 task 并发;跨 task 罕见,
    且这是软上限(非计费 hard cap),日级偶尔多 1 张不影响保护意图。
    """
    if limit <= 0:
        return 0, False
    now_local = datetime.now().astimezone()
    today_start = now_local.replace(hour=0, minute=0, second=0, microsecond=0)
    with session_scope() as s:
        used = s.execute(
            select(func.count()).select_from(UsageEvent).where(
                UsageEvent.user_id == user_id,
                UsageEvent.kind == kind,
                UsageEvent.created_at >= today_start,
            )
        ).scalar_one()
    used_int = int(used)
    return used_int, used_int >= int(limit)
