"""Claude Code transcript collectors."""
from __future__ import annotations

import json
import re
from datetime import datetime
from pathlib import Path

from .cursor import MAX_NEEDLES, fresh_cache, load_cache, save_cache
from .paths import claude_project_dir, ts_in_window

# lead 의 단계 체크포인트 라인(prompts/lead/okstra-lead-contract.md "Progress reporting") — phase id 가
# `phase-<digit>` 로 시작하는 라인만 마커로 인정해 일반 대화의 오탐을 줄인다.
_PROGRESS_RE = re.compile(r"^PROGRESS:\s+(phase-\d\S*(?:[ \t].*)?)$", re.MULTILINE)
_PROGRESS_MARKER_MAX_CHARS = 160


def _event_from_record(rec: dict) -> dict | None:
    """jsonl 레코드 1개 → 압축 이벤트. 집계에 기여하지 않으면 None.

    키: t=timestamp, i/o=input/output, c=cache_creation 합, c5/c1=ephemeral
    5m/1h, r=cache_read, u=tool_use 수. 0/부재 필드는 생략(캐시 크기 절약).
    ts-only 레코드도 보존한다 — 임의 윈도우의 first/last ts 산출에 필요.
    """
    msg = rec.get("message")
    if not isinstance(msg, dict):
        msg = {}
    ev: dict = {}
    usage = msg.get("usage")
    if usage:
        for src, key in (("input_tokens", "i"), ("output_tokens", "o"),
                         ("cache_read_input_tokens", "r")):
            v = usage.get(src, 0) or 0
            if v:
                ev[key] = v
        cc_total = usage.get("cache_creation_input_tokens", 0) or 0
        if cc_total:
            ev["c"] = cc_total
        cc_break = usage.get("cache_creation") or {}
        if isinstance(cc_break, dict) and (
                cc_break.get("ephemeral_5m_input_tokens") is not None
                or cc_break.get("ephemeral_1h_input_tokens") is not None):
            v5 = cc_break.get("ephemeral_5m_input_tokens", 0) or 0
            v1 = cc_break.get("ephemeral_1h_input_tokens", 0) or 0
            if v5:
                ev["c5"] = v5
            if v1:
                ev["c1"] = v1
        elif cc_total:
            # API 분해가 없으면 전부 5m 티어로(1.25x — 더 싼 가정, 기존 동작).
            ev["c5"] = cc_total
    if rec.get("type") == "assistant":
        tools = sum(1 for b in (msg.get("content") or [])
                    if isinstance(b, dict) and b.get("type") == "tool_use")
        if tools:
            ev["u"] = tools
        markers = [
            m.group(1).strip()[:_PROGRESS_MARKER_MAX_CHARS]
            for b in (msg.get("content") or [])
            if isinstance(b, dict) and b.get("type") == "text"
            for m in _PROGRESS_RE.finditer(b.get("text") or "")
        ]
        # 서로 다른 phase id 가 한 메시지에 섞이면 계약 인용/회고 요약이라 그
        # 시각은 단계 경계가 아니다(관측: dev-9186 run 의 막판 요약 1건이 전
        # phase 의 lastAt 을 오염). 같은 phase id 의 반복(병렬 dispatch 3줄,
        # intake reading+complete)은 라이브 체크포인트로 인정한다.
        if markers and len({m.split(None, 1)[0] for m in markers}) == 1:
            ev["p"] = markers
    ts = rec.get("timestamp") or msg.get("timestamp")
    if ts:
        ev["t"] = ts
    return ev or None


def _session_meta_from_record(rec: dict) -> tuple[str | None, str | None]:
    """레코드에서 (agentName, model) 후보 추출 — 둘 다 first-non-null 정책."""
    agent = rec.get("agentName") or None
    model = None
    if rec.get("type") == "assistant":
        msg = rec.get("message")
        if isinstance(msg, dict) and msg.get("model"):
            model = msg["model"]
    return agent, model


_SUBAGENT_FILENAME_RE = re.compile(r"^agent-a(?P<agent>.+)-[0-9a-f]{8,}$")


def _agent_name_from_subagent_path(jsonl_path: Path) -> str | None:
    if jsonl_path.parent.name != "subagents":
        return None
    match = _SUBAGENT_FILENAME_RE.match(jsonl_path.stem)
    return match.group("agent") if match else None


def _iter_claude_session_jsonls(
    proj_dir: Path,
    subagent_parent_sids: list[str] | None = None,
):
    yield from sorted(proj_dir.glob("*.jsonl"))
    if subagent_parent_sids is None:
        yield from sorted(proj_dir.glob("*/subagents/*.jsonl"))
        return
    for sid in sorted({sid for sid in subagent_parent_sids if sid}):
        yield from sorted((proj_dir / sid / "subagents").glob("*.jsonl"))


def _advance_usage_scan(jsonl_path: Path, usage_state: dict) -> dict:
    """`usage_state['offset']` 이후의 완결 라인을 읽어 이벤트를 커밋하고,
    개행 없는 마지막 라인은 transient 로만 반영한 view 를 돌려준다.

    transient tail: 아직 쓰는 중일 수 있는 라인 — 이번 집계에는 포함하되
    커서를 전진시키지 않아, 다음 호출이 완결본으로 다시 읽는다(이중 집계도
    누락도 없음). 깨진 utf-8 / JSON / 비-dict 라인은 건너뛰되 커서는 전진
    (구버전은 text-mode 디코드 실패 시 collect 전체가 죽었다 — fail-open 개선).
    """
    events = list(usage_state.get("events") or [])
    agent_name = usage_state.get("agentName")
    model = usage_state.get("model")
    offset = usage_state.get("offset", 0) or 0
    try:
        size = jsonl_path.stat().st_size
    except OSError:
        size = 0
    if offset > size:
        # 식별자 가드를 통과했더라도 truncate 방어 — 처음부터 재스캔.
        events, agent_name, model, offset = [], None, None, 0
    tail_events: list[dict] = []
    tail_agent: str | None = None
    tail_model: str | None = None
    try:
        with jsonl_path.open("rb") as fh:
            fh.seek(offset)
            while True:
                raw = fh.readline()
                if not raw:
                    break
                rec = None
                stripped = raw.strip()
                if stripped:
                    try:
                        parsed = json.loads(stripped.decode("utf-8"))
                        rec = parsed if isinstance(parsed, dict) else None
                    except (UnicodeDecodeError, json.JSONDecodeError):
                        rec = None
                ev = _event_from_record(rec) if rec else None
                rec_agent, rec_model = _session_meta_from_record(rec) if rec else (None, None)
                if raw.endswith(b"\n"):
                    offset = fh.tell()
                    if agent_name is None and rec_agent:
                        agent_name = rec_agent
                    if model is None and rec_model:
                        model = rec_model
                    if ev:
                        events.append(ev)
                else:
                    tail_agent, tail_model = rec_agent, rec_model
                    if ev:
                        tail_events.append(ev)
                    break
    except OSError:
        pass
    usage_state.update(offset=offset, events=events,
                       agentName=agent_name, model=model)
    return {"events": events + tail_events,
            "agentName": agent_name if agent_name is not None else tail_agent,
            "model": model if model is not None else tail_model}


def _totals_from_events(events: list[dict], agent_name: str | None,
                        model: str | None,
                        since: str | None, until: str | None) -> dict:
    input_t = output_t = cache_create_t = cache_read_t = 0
    cache_create_5m_t = cache_create_1h_t = 0
    tool_uses = 0
    progress_markers: list[dict] = []
    first_ts: str | None = None
    last_ts: str | None = None
    for ev in events:
        ts = ev.get("t")
        if ts and not ts_in_window(ts, since, until):
            continue
        input_t += ev.get("i", 0)
        output_t += ev.get("o", 0)
        cache_create_t += ev.get("c", 0)
        cache_create_5m_t += ev.get("c5", 0)
        cache_create_1h_t += ev.get("c1", 0)
        cache_read_t += ev.get("r", 0)
        tool_uses += ev.get("u", 0)
        for marker in ev.get("p", ()):
            progress_markers.append({"at": ts, "marker": marker})
        if ts:
            if first_ts is None or ts < first_ts:
                first_ts = ts
            if last_ts is None or ts > last_ts:
                last_ts = ts
    duration_ms = 0
    if first_ts and last_ts:
        try:
            a = datetime.fromisoformat(first_ts.replace("Z", "+00:00"))
            b = datetime.fromisoformat(last_ts.replace("Z", "+00:00"))
            duration_ms = max(0, int((b - a).total_seconds() * 1000))
        except ValueError:
            duration_ms = 0
    # '처리 토큰' total 에서 cache_read 는 제외한다. claude 는 매 턴 직전까지의
    # 컨텍스트 전체를 캐시에서 재읽기(cache_read)하므로, 단순 합산하면 같은 토큰을
    # 턴 수만큼 중복 카운트해 처리량이 비현실적으로 부풀려진다(예: in-session
    # lead 가 1.7억으로 표시됨). cache_read 는 cacheReadTokens 로 따로 노출되고,
    # 비용은 pricing 이 0.1x 단가로 별도 반영하므로 total 에서 빼도 비용은 불변.
    total = input_t + output_t + cache_create_t
    return {
        "totalTokens": total,
        "inputTokens": input_t,
        "outputTokens": output_t,
        "cacheCreationTokens": cache_create_t,
        "cacheCreation5mTokens": cache_create_5m_t,
        "cacheCreation1hTokens": cache_create_1h_t,
        "cacheReadTokens": cache_read_t,
        "toolUses": tool_uses,
        "durationMs": duration_ms,
        "agentName": agent_name,
        "model": model,
        "startedAt": first_ts,
        "endedAt": last_ts,
        "progressMarkers": progress_markers,
    }


def claude_session_totals(
    jsonl_path: Path, *, since: str | None = None, until: str | None = None,
    incremental: bool = False,
) -> dict:
    """Return totals + agentName + assistant model + time window for a Claude session jsonl.

    ``since`` / ``until`` are ISO-8601 timestamp strings (UTC ``...Z``). When
    given, only records whose ``timestamp`` falls within ``[since, until]`` are
    counted toward tokens / tool_uses / duration. This is the run-scoping seam:
    an **in-session** lead writes its run into the user's whole-session jsonl,
    so without a window the totals swallow every unrelated turn (observed:
    lead billed 1.7억 tokens / $416 / 3h for a single requirements-discovery
    run). ``agentName`` / ``model`` are session metadata and are read from the
    whole file regardless of the window. Records without a timestamp are kept
    (conservative — never silently drop usage when we can't place it in time).

    ``incremental=True`` 면 $OKSTRA_HOME 캐시의 byte cursor 이후만 읽는다.
    캐시에는 윈도우 적용 전 이벤트가 저장되므로 호출마다 다른 since/until
    에도 결과는 전체 스캔과 동일하다 (P6 plan 참조).
    """
    if incremental:
        cache = load_cache(jsonl_path)
        view = _advance_usage_scan(jsonl_path, cache["usage"])
        save_cache(jsonl_path, cache)
    else:
        view = _advance_usage_scan(jsonl_path, fresh_cache()["usage"])
    agent_name = view["agentName"] or _agent_name_from_subagent_path(jsonl_path)
    return _totals_from_events(view["events"], agent_name,
                               view["model"], since, until)


def _needle_scan(jsonl_path: Path, entry: dict, needle_lower: str) -> bool:
    """entry({'offset','found'}) 를 전진시키며 needle 존재 여부 반환.

    미완결 tail 라인도 검사한다 — 부분 문자열 매칭은 라인 완결 후에도 유효
    하므로 found=True 는 그대로 커밋해도 안전하다. 단 offset 은 완결 라인
    까지만 전진해, 미완결 tail 은 다음 호출이 다시 본다.
    """
    if entry.get("found"):
        return True
    offset = entry.get("offset", 0) or 0
    try:
        if offset > jsonl_path.stat().st_size:
            offset = 0  # truncate/교체 방어
        with jsonl_path.open("rb") as fh:
            fh.seek(offset)
            while True:
                raw = fh.readline()
                if not raw:
                    break
                if needle_lower in raw.decode("utf-8", errors="replace").lower():
                    entry["found"] = True
                    entry["offset"] = offset
                    return True
                if raw.endswith(b"\n"):
                    offset = fh.tell()
    except OSError:
        return False
    entry["offset"] = offset
    return False


def _cached_needle_scan(jsonl_path: Path, cache: dict, needle_lower: str) -> bool:
    """`cache['needles']` 의 per-needle cursor 를 유지하며 `_needle_scan` 수행.
    파일당 MAX_NEEDLES 개까지 오래된 순으로 교체 보존한다."""
    needles = cache.setdefault("needles", {})
    entry = needles.get(needle_lower)
    if entry is None:
        entry = {"offset": 0, "found": False}
        while len(needles) >= MAX_NEEDLES:
            needles.pop(next(iter(needles)))
        needles[needle_lower] = entry
    return _needle_scan(jsonl_path, entry, needle_lower)


def find_claude_team_sessions(
    cwd: Path,
    team_needles,
    lead_sid: str | None = None,
    projects_root: Path | None = None,
    *,
    incremental: bool = False,
) -> dict[str, Path]:
    """Map sessionId -> jsonl path for all jsonls tagged with any of `team_needles`.

    `team_needles` is an ``Iterable[str]`` of teamName generations; a jsonl is
    included when it matches any one of them (union). The lead re-issue
    fragments an implicit-team session into multiple `session-<leadSid>`
    generations, so a single string is no longer sufficient.

    Matching is case-insensitive on the teamName needle to tolerate runs where
    the lead recorded `team.teamName` with a different case than the harness
    serialised into the transcript (e.g. `okstra-DEV-6827` vs `okstra-dev-6827`).

    If `lead_sid` is provided and exists in the project dir, it is always
    included even when no teamName needle matches — this lets us recover lead
    usage in fallback runs that never wrote `team.teamName` into team-state.

    `projects_root` 는 테스트/진단용 주입 시드 — 기본은 실제 ~/.claude/projects.

    ``incremental=True`` 면 파일별 needle cursor 이후의 신규 byte 만 검사한다.
    needle(=team 이름)은 run 마다 다르므로 파일당 MAX_NEEDLES 개까지 오래된
    순으로 교체 보존한다.
    """
    proj_dir = claude_project_dir(cwd, projects_root)
    out: dict[str, Path] = {}
    if not proj_dir.is_dir():
        return out
    needles_lower = [f'"teamname":"{n.lower()}"' for n in team_needles if n]
    for p in _iter_claude_session_jsonls(proj_dir):
        for needle_lower in needles_lower:
            if incremental:
                cache = load_cache(p)
                hit = _cached_needle_scan(p, cache, needle_lower)
                save_cache(p, cache)
            else:
                hit = _needle_scan(p, {"offset": 0, "found": False}, needle_lower)
            if hit:
                out[p.stem] = p
                break
    if lead_sid:
        direct = proj_dir / f"{lead_sid}.jsonl"
        if direct.is_file():
            out.setdefault(lead_sid, direct)
    return out


_TEAM_TAG_NEEDLE = '"teamname":"'


def _agent_name_matches(agent_name: str | None, prefixes: list[str]) -> bool:
    if not agent_name:
        return False
    for prefix in prefixes:
        if agent_name == prefix or agent_name.startswith(f"{prefix}-"):
            return True
    return False


def find_claude_agent_sessions(
    cwd: Path,
    agent_prefixes: list[str],
    projects_root: Path | None = None,
    *,
    incremental: bool = False,
    subagent_parent_sids: list[str] | None = None,
) -> dict[str, Path]:
    """Map worker-like Claude jsonls by agent name or nested subagent filename.

    Legacy top-level worker jsonls record `agentName` in the JSONL body. Newer
    in-process subagents can instead live under
    `<lead-session>/subagents/agent-a<name>-<hash>.jsonl` with no body
    `agentName`, so this finder also matches that filename form.

    Callers can pass `subagent_parent_sids` to limit nested subagents to the
    current lead session directory. Top-level fallback discovery remains broad
    and must still be filtered by the run window after totals are read.
    """
    proj_dir = claude_project_dir(cwd, projects_root)
    out: dict[str, Path] = {}
    if not proj_dir.is_dir():
        return out
    agent_needles = [f'"agentname":"{p.lower()}' for p in agent_prefixes if p]
    if not agent_needles:
        return out
    for p in _iter_claude_session_jsonls(proj_dir, subagent_parent_sids):
        path_agent = _agent_name_from_subagent_path(p)
        filename_matched = _agent_name_matches(path_agent, agent_prefixes)
        scoped_subagent = (
            filename_matched
            and path_agent is not None
            and subagent_parent_sids is not None
        )
        if incremental:
            cache = load_cache(p)
            matched = filename_matched or any(
                _cached_needle_scan(p, cache, n) for n in agent_needles
            )
            team_tagged = matched and _cached_needle_scan(p, cache, _TEAM_TAG_NEEDLE)
            save_cache(p, cache)
        else:
            matched = filename_matched or any(
                _needle_scan(p, {"offset": 0, "found": False}, n)
                for n in agent_needles
            )
            team_tagged = matched and _needle_scan(
                p, {"offset": 0, "found": False}, _TEAM_TAG_NEEDLE
            )
        if matched and (scoped_subagent or not team_tagged):
            out[p.stem] = p
    return out