"""Codex CLI rollout collectors."""
from __future__ import annotations

import json
import os
from pathlib import Path
from .jsonl_io import iter_jsonl
from .paths import CODEX_SESSIONS, codex_session_roots, ts_in_window


_DEFAULT_CODEX_SESSIONS = CODEX_SESSIONS


def codex_session_total(jsonl_path: Path) -> dict:
    """Return last token_count snapshot from a codex rollout jsonl."""
    last: dict | None = None
    cwd_val: str | None = None
    model_val: str | None = None
    started: str | None = None
    ended: str | None = None
    for rec in iter_jsonl(jsonl_path):
        if rec.get("type") == "session_meta":
            payload = rec.get("payload") or {}
            cwd_val = payload.get("cwd")
            started = payload.get("timestamp")
        elif rec.get("type") == "turn_context":
            payload = rec.get("payload") or {}
            if model_val is None and payload.get("model"):
                model_val = payload["model"]
        elif rec.get("type") == "event_msg":
            payload = rec.get("payload") or {}
            if payload.get("type") == "token_count":
                last = payload
                ended = rec.get("timestamp")
    if last is None:
        return {"totalTokens": 0, "cwd": cwd_val, "model": model_val, "available": False}
    info = last.get("info") or {}
    total = info.get("total_token_usage") or {}
    return {
        "totalTokens": total.get("total_tokens", 0) or 0,
        "inputTokens": total.get("input_tokens", 0) or 0,
        "cachedInputTokens": total.get("cached_input_tokens", 0) or 0,
        "outputTokens": total.get("output_tokens", 0) or 0,
        "reasoningOutputTokens": total.get("reasoning_output_tokens", 0) or 0,
        "cwd": cwd_val,
        "model": model_val,
        "startedAt": started,
        "endedAt": ended,
        "available": True,
    }


def find_codex_session(cwd: Path, started_at: str, ended_at: str) -> Path | None:
    """Find the latest codex rollout jsonl in the requested window."""
    sessions = find_codex_sessions(cwd, started_at, ended_at)
    return sessions[-1] if sessions else None


def _session_meta_payload(path: Path) -> dict | None:
    try:
        with path.open() as fh:
            first = fh.readline()
    except OSError:
        return None
    if not first:
        return None
    try:
        record = json.loads(first)
    except json.JSONDecodeError:
        return None
    if record.get("type") != "session_meta":
        return None
    payload = record.get("payload")
    if not isinstance(payload, dict):
        payload = {}
    timestamp = payload.get("timestamp") or record.get("timestamp") or ""
    return {**payload, "timestamp": timestamp}


def _session_metadata(path: Path) -> tuple[str, str] | None:
    payload = _session_meta_payload(path)
    if payload is None:
        return None
    return str(payload.get("cwd") or ""), str(payload.get("timestamp") or "")


def codex_session_is_worker(path: Path) -> bool:
    """exec 래퍼 세션. 대화형 리드는 originator=codex-tui / source=cli 이다."""
    payload = _session_meta_payload(path)
    if not payload:
        return False
    originator = str(payload.get("originator") or "").strip()
    source = str(payload.get("source") or "").strip()
    return originator == "codex_exec" or source == "exec"


def codex_session_ids(path: Path) -> set[str]:
    ids = {path.name, path.stem}
    payload = _session_meta_payload(path) or {}
    for key in ("session_id", "id"):
        value = str(payload.get(key) or "").strip()
        if value:
            ids.add(value)
    return {item for item in ids if item}


def find_codex_sessions(
    cwd: Path,
    started_at: str,
    ended_at: str,
    *,
    session_roots: tuple[Path, ...] | None = None,
) -> list[Path]:
    """Find codex rollout jsonls whose meta.cwd matches the window."""
    if not started_at or not ended_at:
        return []
    if session_roots is None:
        if CODEX_SESSIONS != _DEFAULT_CODEX_SESSIONS:
            session_roots = (CODEX_SESSIONS,)
        else:
            session_roots = codex_session_roots(Path.home(), os.environ)
    target_cwd = str(cwd)
    candidates: list[tuple[str, Path]] = []
    seen: set[Path] = set()
    for root in session_roots:
        if not root.is_dir():
            continue
        for p in root.rglob("rollout-*.jsonl"):
            identity = p.resolve()
            if identity in seen:
                continue
            seen.add(identity)
            metadata = _session_metadata(p)
            if metadata is None:
                continue
            session_cwd, ts = metadata
            if session_cwd != target_cwd:
                continue
            if not ts_in_window(ts, started_at, ended_at):
                continue
            candidates.append((ts, p))
    if not candidates:
        return []
    candidates.sort()
    return [path for _ts, path in candidates]