"""Antigravity (`agy`) CLI usage collector.

공식 headless 문서는 ``--output-format json`` / ``stream-json`` 의 최종
``usage`` 객체에 토큰을 둔다. okstra 는 ``stream-json`` 으로 호출하고, 러너가
스트림의 마지막 ``usage`` 스냅샷을 래퍼 status 사이드카의 ``usage`` 에 적는다
(`worker_runner`) — 그것이 정본이다. 워커 로그는 러너가 사람이 읽는 줄로 옮겨
적은 것이라 stream-json 이 아니다. raw stream-json 로그를 읽는 경로는 러너
이전의 로그를 위해 남긴다.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Mapping

from .jsonl_io import iter_jsonl


def _usage_total(usage: Mapping[str, Any], model: str | None) -> dict:
    return {
        "totalTokens": usage.get("total_tokens", 0) or 0,
        "inputTokens": usage.get("input_tokens", 0) or 0,
        "outputTokens": usage.get("output_tokens", 0) or 0,
        "thoughtsTokens": usage.get("thinking_tokens", 0) or 0,
        "cacheReadTokens": usage.get("cache_read_tokens", 0) or 0,
        "model": model,
        "available": True,
    }


def _read_status(status_path: Path) -> dict | None:
    try:
        data = json.loads(status_path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    return data if isinstance(data, dict) else None


def status_carries_usage(path: Path) -> bool:
    """이 경로가 ``usage`` 스냅샷을 실은 래퍼 status 사이드카인가."""
    if not path.name.endswith(".status.json"):
        return False
    status = _read_status(path)
    return isinstance((status or {}).get("usage"), dict) and bool(status["usage"])


def antigravity_status_total(status_path: Path) -> dict:
    """래퍼 status 사이드카의 ``usage`` 스냅샷. 모델은 served-model 관측값이다."""
    status = _read_status(status_path) or {}
    usage = status.get("usage")
    if not isinstance(usage, dict) or not usage:
        return {"totalTokens": 0, "available": False}
    attestation = status.get("servedModelAttestation")
    model = (
        attestation.get("observedModel") if isinstance(attestation, dict) else None
    )
    return _usage_total(usage, model if isinstance(model, str) and model else None)


def antigravity_session_total(json_path: Path) -> dict:
    """raw stream-json 로그에서 마지막 result.usage 를 읽는다."""
    result_usage: dict | None = None
    step_usage: dict | None = None
    model: str | None = None
    for record in iter_jsonl(json_path):
        kind = record.get("event")
        if kind == "init":
            init = record.get("init")
            if isinstance(init, dict) and isinstance(init.get("model"), str):
                model = init["model"]
            continue
        if kind == "result":
            result = record.get("result")
            if isinstance(result, dict) and isinstance(result.get("usage"), dict):
                result_usage = result["usage"]
            continue
        if kind == "step_update":
            step = record.get("step_update")
            if isinstance(step, dict) and isinstance(step.get("usage"), dict):
                step_usage = step["usage"]
    usage = result_usage or step_usage
    if usage is None:
        return {"totalTokens": 0, "available": False}
    return _usage_total(usage, model)


def find_antigravity_session(project_root: Path, started_at: str, ended_at: str) -> Path | None:
    sessions = find_antigravity_sessions(project_root, started_at, ended_at)
    return sessions[-1] if sessions else None


def find_antigravity_sessions(project_root: Path, started_at: str, ended_at: str) -> list[Path]:
    # 세션 파일은 홈 디렉터리가 아니라 워커 로그다. 경로는 호출부가 prompt 에서 붙인다.
    return []