"""Per-source usage block constructors (consumed by collect())."""
from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from .paths import utc_now
from .pricing import (
    claude_billable_equivalent,
    claude_cost_usd,
)

_CACHE_INCLUDED_SOURCES = frozenset({"codex-cli", "grok-cli", "kimi-cli"})


def accounting_workers(state: Mapping[str, Any]) -> list[dict]:
    """초기 명부와 명부 밖 실행의 사용량 행을 같은 소비 경로로 전달한다."""
    rows = []
    for key in ("workers", "additionalWorkerUsage"):
        collection = state.get(key)
        if isinstance(collection, list):
            rows.extend(row for row in collection if isinstance(row, dict))
    return rows


def usage_blocks(state: dict) -> list[dict]:
    blocks = [state.get("leadUsage") or {}]
    blocks.extend(worker.get("usage") or {} for worker in accounting_workers(state))
    unattributed = (state.get("usageSummary") or {}).get("unattributedWorkerUsage")
    if isinstance(unattributed, dict):
        blocks.append(unattributed)
    return [block for block in blocks if isinstance(block, dict)]


def normalize_usage_block(block: dict) -> dict:
    """과거 캐시 포함 합계를 정규화하고 게시 당시 합계를 별도로 보존한다."""
    normalized = dict(block)
    cached = block.get("cachedInputTokens") or 0
    if (
        block.get("source") not in _CACHE_INCLUDED_SOURCES
        or "cacheReadTokens" in block
        or block.get("accountingBasis") == "cache-read-excluded"
        or not isinstance(cached, int)
        or isinstance(cached, bool)
        or cached <= 0
    ):
        return normalized
    total = block.get("totalTokens")
    if not isinstance(total, int) or total < cached:
        return normalized
    normalized["reportedTotalTokens"] = total
    normalized["totalTokens"] = total - cached
    normalized["cacheReadTokens"] = cached
    normalized["accountingBasis"] = "cache-read-excluded"
    if isinstance(block.get("cliTotalTokens"), int):
        normalized["reportedCliTotalTokens"] = block["cliTotalTokens"]
        normalized["cliTotalTokens"] = max(0, block["cliTotalTokens"] - cached)
    return normalized


def usage_block(totals: dict, source: str, note: str | None = None) -> dict:
    block = {
        "totalTokens": totals.get("totalTokens", 0) or 0,
        "inputTokens": totals.get("inputTokens", 0) or 0,
        "outputTokens": totals.get("outputTokens", 0) or 0,
        "toolUses": totals.get("toolUses", 0) or 0,
        "durationMs": totals.get("durationMs", 0) or 0,
        "startedAt": totals.get("startedAt"),
        "endedAt": totals.get("endedAt"),
        "source": source,
        "collectedAt": utc_now(),
    }
    for key in ("cacheCreationTokens", "cacheCreation5mTokens", "cacheCreation1hTokens",
                "cacheReadTokens", "cachedInputTokens",
                "reasoningOutputTokens", "cachedTokens", "thoughtsTokens", "toolTokens"):
        if totals.get(key):
            block[key] = totals[key]

    # Billable-equivalent + cost.
    if source == "claude-jsonl":
        cc_1h = totals.get("cacheCreation1hTokens", 0) or 0
        be = claude_billable_equivalent(
            totals.get("inputTokens", 0) or 0,
            totals.get("cacheCreationTokens", 0) or 0,
            totals.get("cacheReadTokens", 0) or 0,
            totals.get("outputTokens", 0) or 0,
            cache_create_1h_t=cc_1h,
        )
        block["billableEquivalentTokens"] = be
        cost = claude_cost_usd(
            totals.get("model"),
            totals.get("inputTokens", 0) or 0,
            totals.get("cacheCreationTokens", 0) or 0,
            totals.get("cacheReadTokens", 0) or 0,
            totals.get("outputTokens", 0) or 0,
            cache_create_1h_t=cc_1h,
        )
        if cost is not None:
            block["estimatedCostUsd"] = cost
        if totals.get("model"):
            block["model"] = totals["model"]
    if note:
        block["note"] = note
    return block


def na_block(reason: str) -> dict:
    return {
        "totalTokens": 0,
        "toolUses": 0,
        "durationMs": 0,
        "startedAt": None,
        "endedAt": None,
        "source": "unavailable",
        "collectedAt": utc_now(),
        "note": reason,
    }