#!/usr/bin/env python3
"""Shrink a Claude Code transcript slice to fit a token budget.

    transcript-budget.py <in.jsonl> <out.jsonl> [--budget-tokens N]
                         [--report-json PATH]

Escalates a fixed ladder of rungs until the output fits, and exits 1 if even
the last rung cannot. A report of what was elided goes to stderr.
"""
import argparse
import hashlib
import json
import os
import shutil
import sys

# measured across real transcripts, bytes per token ranges 1.8-2.9, so no fixed
# ratio both respects the cap and uses it; count for real via the BPE encoder
ENCODING = "o200k_base"
REEXEC_FLAG = "GPT_ADVISOR_BUDGET_REEXEC"
FALLBACK_BYTES_PER_TOKEN = 1.5


def load_counter():
    """Return (count_fn, exact). Re-execs once under uv when tiktoken is missing."""
    try:
        import tiktoken
    except ImportError:
        uv = shutil.which("uv")
        if uv and not os.environ.get(REEXEC_FLAG):
            os.environ[REEXEC_FLAG] = "1"
            os.execv(uv, [uv, "run", "--quiet", "--with", "tiktoken",
                          "python", os.path.abspath(__file__)] + sys.argv[1:])
        return (lambda text: int(len(text) / FALLBACK_BYTES_PER_TOKEN)), False
    encoder = tiktoken.get_encoding(ENCODING)
    return (lambda text: len(encoder.encode(text, disallowed_special=()))), True

# harness plumbing re-injected every turn; carries no record of what the agent did
PLUMBING_ATTACHMENTS = {
    "hook_success",
    "hook_error",
    "deferred_tools_delta",
    "agent_listing_delta",
    "skill_listing",
    "command_permissions",
    "todo",
    "selected_lines_in_ide",
    "opened_file_in_ide",
    "nested_memory",
    "mcp_resource_listing",
}
NON_CONVERSATIONAL = {
    "file-history-snapshot",
    "last-prompt",
    "ai-title",
    "queue-operation",
    "diagnostics",
}
# newest slice of records treated as recent; they keep the generous cap
RECENT_FRACTION = 0.25

# (recent cap, older cap) in characters; None = no clipping at all
RUNGS = [None, (8000, 2000), (2000, 600), (600, 200), (400, 120)]
DROP_FROM_RUNG = 2
MIN_RETAINED = 0.4
OPAQUE_BLOCK_FIELDS = {"signature", "data"}
# the receiving side frames the attachment, so aim under the caller's ceiling
HEADROOM = 0.99


class Stats:
    def __init__(self):
        self.dropped_records = 0
        self.dropped_attachments = 0
        self.deduped_attachments = 0
        self.clipped_blocks = 0
        self.elided_chars = 0
        self.dropped_middle = 0


def clip(text, cap, stats):
    if cap is None or len(text) <= cap:
        return text
    head = int(cap * 0.6)
    tail = cap - head
    stats.clipped_blocks += 1
    stats.elided_chars += len(text) - cap
    return f"{text[:head]}\n...[{len(text) - cap} chars elided]...\n{text[-tail:]}"


def attachment_digest(record):
    return hashlib.sha1(
        json.dumps(record.get("attachment") or {}, sort_keys=True).encode()
    ).hexdigest()


def clip_payload(value, cap, stats):
    """Clip a block payload, keeping its structure when no clipping is needed."""
    if isinstance(value, str):
        return clip(value, cap, stats)
    serialized = json.dumps(value)
    if cap is None or len(serialized) <= cap:
        return value
    return clip(serialized, cap, stats)


def clip_block(block, cap, stats):
    """Clip the bulky payload of one content block; leave reasoning and prose alone."""
    kind = block.get("type")
    if kind in ("thinking", "redacted_thinking"):
        # the signature is an opaque provider blob: measured 3.7MB against 12KB of
        # actual reasoning text in one transcript
        return {k: v for k, v in block.items() if k not in OPAQUE_BLOCK_FIELDS}
    if kind == "tool_result":
        content = block.get("content")
        if content is None:
            return block
        return {**block, "content": clip_payload(content, cap, stats)}
    if kind == "tool_use":
        return {**block, "input": clip_payload(block.get("input", ""), cap, stats)}
    return block


def filter_records(lines, caps, stats):
    """Apply one rung to the raw lines, returning the surviving JSON records.

    caps is None for the lossless rung, else (recent_cap, older_cap)."""
    if caps is None:
        out = []
        for line in lines:
            try:
                out.append(json.loads(line))
            except ValueError:
                stats.dropped_records += 1
        return out

    recent_cap, older_cap = caps
    parsed = []
    for line in lines:
        try:
            parsed.append(json.loads(line))
        except ValueError:
            stats.dropped_records += 1
    # recency counted in conversational turns, not raw lines: a burst of hook records
    # would otherwise push real turns out of the recent window
    turns = [i for i, r in enumerate(parsed) if isinstance(r.get("message"), dict)]
    recent_from = turns[len(turns) - int(len(turns) * RECENT_FRACTION)] if turns else 0

    remaining = {}
    for record in parsed:
        if record.get("type") == "attachment":
            remaining[attachment_digest(record)] = remaining.get(attachment_digest(record), 0) + 1

    out = []
    for index, record in enumerate(parsed):
        kind = record.get("type")
        if kind in NON_CONVERSATIONAL:
            stats.dropped_records += 1
            continue
        if kind == "attachment":
            attachment = record.get("attachment") or {}
            if attachment.get("type") in PLUMBING_ATTACHMENTS:
                stats.dropped_attachments += 1
                continue
            # keep the LAST copy of a re-injected attachment: it sits nearest the turns
            # under review and lands in the recent (generous) clipping band
            digest = attachment_digest(record)
            remaining[digest] -= 1
            if remaining[digest] > 0:
                stats.deduped_attachments += 1
                continue

        cap = recent_cap if index >= recent_from else older_cap
        message = record.get("message")
        # per-record harness metadata (uuid, sessionId, cwd, version, ...) is a large
        # constant tax across tens of thousands of records and tells a reviewer nothing
        projected = {"type": kind}
        if isinstance(message, dict):
            content = message.get("content")
            if isinstance(content, list):
                content = [
                    clip_block(b, cap, stats) if isinstance(b, dict) else b
                    for b in content
                ]
            elif isinstance(content, str):
                content = clip(content, cap, stats)
            projected["message"] = {"role": message.get("role"), "content": content}
        elif kind == "attachment":
            projected["attachment"] = clip_payload(record.get("attachment"), cap, stats)
        else:
            projected["record"] = clip_payload(
                {k: v for k, v in record.items() if k != "type"}, cap, stats
            )
        out.append(projected)
    return out


def drop_middle(records, sizes, budget, stats):
    """Last rung: drop whole records from the middle of the filtered stream."""
    total = sum(sizes)
    if total <= budget:
        return records
    head = max(1, int(len(records) * 0.15))
    cut = head
    while cut < len(records) - 1 and total > budget:
        total -= sizes[cut]
        cut += 1
        stats.dropped_middle += 1
    # the opening turns are the oldest context, so they yield after the middle
    while head > 1 and total > budget:
        head -= 1
        total -= sizes[head]
        stats.dropped_middle += 1
    if total > budget:
        return None
    return records[:head] + records[cut:]


def report(stats, rung, tokens, budget, exact):
    measure = "counted" if exact else "estimated"
    return (
        f"rung {rung}: {tokens:,} tok {measure} (budget {budget:,}) | "
        f"records dropped {stats.dropped_records}, "
        f"attachments dropped {stats.dropped_attachments} "
        f"deduped {stats.deduped_attachments}, "
        f"blocks clipped {stats.clipped_blocks} "
        f"({stats.elided_chars:,} chars elided), "
        f"middle records dropped {stats.dropped_middle}"
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("source")
    parser.add_argument("dest")
    parser.add_argument("--budget-tokens", type=int, default=900_000)
    parser.add_argument("--report-json")
    args = parser.parse_args()

    if args.budget_tokens <= 0:
        parser.error("--budget-tokens must be positive")

    with open(args.source, errors="replace") as handle:
        lines = handle.readlines()
    if not lines:
        print("transcript-budget: empty input slice", file=sys.stderr)
        return 1

    count, exact = load_counter()
    budget = int(args.budget_tokens * HEADROOM)

    for rung, caps in enumerate(RUNGS):
        stats = Stats()
        records = filter_records(lines, caps, stats)
        encoded = [json.dumps(r) + "\n" for r in records]
        # per-record counts sum to slightly more than the whole payload's, since
        # they forbid BPE merges across record boundaries: over-count, never under
        sizes = [count(e) for e in encoded]
        last = rung == len(RUNGS) - 1
        # detail on the turns that survive beats uniform 120-char clipping of every
        # turn, so drop old records before tightening the caps further
        if sum(sizes) > budget and rung >= DROP_FROM_RUNG:
            trimmed = drop_middle(records, sizes, budget, stats)
            if trimmed is None and last:
                print(
                    f"transcript-budget: cannot fit {budget:,} tokens "
                    f"even at the last rung",
                    file=sys.stderr,
                )
                return 1
            if trimmed is not None and (last or len(trimmed) >= len(records) * MIN_RETAINED):
                records = trimmed
                encoded = [json.dumps(r) + "\n" for r in records]
                sizes = [count(e) for e in encoded]
        tokens = sum(sizes)
        if tokens <= budget:
            if not records:
                print("transcript-budget: every record was filtered out", file=sys.stderr)
                return 1
            with open(args.dest, "w") as handle:
                handle.write("".join(encoded))
            line = report(stats, rung, tokens, budget, exact)
            print(f"transcript-budget: {line}", file=sys.stderr)
            if args.report_json:
                with open(args.report_json, "w") as handle:
                    json.dump(
                        {
                            "rung": rung,
                            "lossless": rung == 0,
                            "tokens": tokens,
                            "exact": exact,
                            "budget_tokens": budget,
                            "records": len(records),
                            "summary": line,
                            **vars(stats),
                        },
                        handle,
                    )
            return 0
    return 1


if __name__ == "__main__":
    sys.exit(main())