"""abreka util repair-session-files — shrink bloated pi session JSONs."""

from __future__ import annotations

from pathlib import Path

import click

from abreka.config import find_abreka_dir
from abreka.session_repair import (
    RepairPlan,
    SessionRepairError,
    apply_repair,
    discover_session_files,
    scan_session_file,
)


def _fmt_size(n: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024 or unit == "GB":
            return f"{n:,.1f} {unit}" if unit != "B" else f"{n:,} B"
        n /= 1024
    return f"{n:,.1f} GB"  # unreachable


@click.command("repair-session-files")
@click.argument("paths", nargs=-1, type=click.Path(exists=True, path_type=Path))
@click.option("--apply", "apply_flag", is_flag=True, help="Actually rewrite files (default: dry-run).")
@click.option("--yes", "-y", is_flag=True, help="Skip confirmation when --apply is set.")
@click.option("--keep-backup", is_flag=True, help="Move original to <file>.bak instead of replacing.")
def repair_session_files_cmd(
    paths: tuple[Path, ...],
    apply_flag: bool,
    yes: bool,
    keep_backup: bool,
):
    """Shrink pi_session*.json files by stripping redundant streaming payloads.

    Repairs include:
    - dropping ``message_update`` / ``assistantMessageEvent.type=toolcall_delta``
    - compacting ``message_update.message`` and ``assistantMessageEvent.partial``

    By default this is a DRY RUN: it reports what would change but does not
    touch any file. Pass --apply to actually rewrite files.

    With no PATHS, scans the current abreka project's experiments directory.
    """
    if paths:
        targets = [Path(p) for p in paths]
    else:
        abreka_dir = find_abreka_dir()
        targets = discover_session_files(abreka_dir)

    if not targets:
        click.echo("No session files found.")
        return

    # Pass 1: counts-only summary so we never hold more than one file's
    # decoded events in memory at a time. Multi-GB sessions on this project
    # can produce >1 GB of Python dicts per file when fully decoded.
    summaries: list[RepairPlan] = []
    repairable_paths: list[Path] = []
    errored = 0
    unchanged = 0

    for path in targets:
        summary = scan_session_file(path, counts_only=True)
        summaries.append(summary)
        if not summary.ok:
            errored += 1
            continue
        if summary.needs_repair:
            repairable_paths.append(path)
        else:
            unchanged += 1

    click.echo(f"Scanned {len(summaries)} session file(s):")
    click.echo("")
    for summary in summaries:
        rel = str(summary.path)
        if not summary.ok:
            click.echo(f"  ERROR  {rel}")
            click.echo(f"         {summary.error}")
            continue
        if not summary.needs_repair:
            click.echo(
                f"  ok     {rel}  ({_fmt_size(summary.original_size)}, "
                f"{summary.total_events:,} events, nothing to repair)"
            )
            continue

        est_after = max(0, summary.original_size - summary.estimated_saved_bytes)
        click.echo(f"  FIX    {rel}")
        click.echo(
            f"         before: {_fmt_size(summary.original_size)}  "
            f"events: {summary.total_events:,}"
        )
        click.echo(
            f"         after:  ~{_fmt_size(est_after)}  "
            f"drop: {summary.dropped_events:,} events, "
            f"compact: {summary.compacted_events:,} events"
        )

    click.echo("")
    click.echo(
        f"Summary: {len(repairable_paths)} file(s) to repair, "
        f"{unchanged} already clean, {errored} errored."
    )

    if errored:
        click.echo(
            f"Errors encountered on {errored} file(s); those will be skipped.",
            err=True,
        )

    if not repairable_paths:
        if not apply_flag:
            click.echo("")
            click.echo("Nothing to repair.")
        return

    if not apply_flag:
        click.echo("")
        click.echo("This was a DRY RUN. Re-run with --apply to rewrite the files above.")
        return

    if not yes:
        click.confirm(
            f"Rewrite {len(repairable_paths)} session file(s)?",
            abort=True,
        )

    # Pass 2: rescan each file with kept_events, apply, then drop. Peak
    # memory during apply is bounded to one file's decoded events.
    click.echo("")
    failures = 0
    for path in repairable_paths:
        plan = scan_session_file(path)
        if not plan.ok:
            failures += 1
            click.echo(f"  FAILED    {path}: {plan.error}", err=True)
            continue
        try:
            apply_repair(plan, keep_backup=keep_backup)
            new_size = path.stat().st_size
            saved = plan.original_size - new_size
            click.echo(
                f"  repaired  {path}  "
                f"{_fmt_size(plan.original_size)} → {_fmt_size(new_size)} "
                f"(-{_fmt_size(saved)})"
            )
        except SessionRepairError as e:
            failures += 1
            click.echo(f"  FAILED    {path}: {e}", err=True)
        except Exception as e:
            failures += 1
            click.echo(f"  FAILED    {path}: {e}", err=True)
        finally:
            # Drop the decoded event list as soon as possible so the next
            # iteration starts from a clean slate.
            plan.kept_events = None
            del plan

    click.echo("")
    if failures:
        raise click.ClickException(
            f"{failures} file(s) failed to repair. See messages above."
        )
    click.echo(f"Repaired {len(repairable_paths)} file(s).")
