#!/usr/bin/env python3
"""Progressive coding harness — the stage-gated interview simulator.

The reported onsite format presents a problem in stages. Each stage must have a
working solution before the next one opens, and you do not get to see stage N+1
while designing stage N. That constraint is the whole difficulty: you must pick a
representation that survives requirements you have not been told yet.

This tool reproduces it. It shows you one gate at a time, refuses to show the next
until your code passes the current one, and records the wall-clock time at which
each gate opened.

Usage:
    ./progressive.py list                      # the catalog
    ./progressive.py start token-stream-differ # begin; prints gate 1 only
    ./progressive.py brief token-stream-differ # re-print the current gate
    ./progressive.py test  token-stream-differ # run current gate; unlock on pass
    ./progressive.py status                    # all in-flight problems
    ./progressive.py log                       # timing history
    ./progressive.py chart                     # time-to-first-gate trend
    ./progressive.py reset token-stream-differ # start over

The metric that matters most is TIME-TO-FIRST-PASSING-GATE. In a gated format an
unopened gate scores zero no matter how elegant the unwritten design was, so the
compounding skill is getting something correct fast and extending it under
pressure. `chart` plots that number over time.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import shutil
import sys
import time
import traceback
from typing import Any

HERE = os.path.dirname(os.path.abspath(__file__))
WORK = os.path.join(HERE, "work")
LOG = os.path.join(HERE, "timing-log.jsonl")

sys.path.insert(0, HERE)

from problems import CATALOG, load_problem  # noqa: E402


# ---------------------------------------------------------------------------
# terminal helpers
# ---------------------------------------------------------------------------

BOLD = "\033[1m" if sys.stdout.isatty() else ""
DIM = "\033[2m" if sys.stdout.isatty() else ""
GREEN = "\033[32m" if sys.stdout.isatty() else ""
RED = "\033[31m" if sys.stdout.isatty() else ""
YELLOW = "\033[33m" if sys.stdout.isatty() else ""
OFF = "\033[0m" if sys.stdout.isatty() else ""


def rule(char: str = "─", width: int = 74) -> str:
    return char * width


def mmss(seconds: float) -> str:
    seconds = int(seconds)
    return f"{seconds // 60}m{seconds % 60:02d}s"


# ---------------------------------------------------------------------------
# state
# ---------------------------------------------------------------------------


def work_dir(problem_id: str) -> str:
    return os.path.join(WORK, problem_id)


def state_path(problem_id: str) -> str:
    return os.path.join(work_dir(problem_id), "state.json")


def read_state(problem_id: str) -> dict | None:
    path = state_path(problem_id)
    if not os.path.exists(path):
        return None
    with open(path) as handle:
        return json.load(handle)


def write_state(problem_id: str, state: dict) -> None:
    os.makedirs(work_dir(problem_id), exist_ok=True)
    with open(state_path(problem_id), "w") as handle:
        json.dump(state, handle, indent=2)


def append_log(record: dict) -> None:
    with open(LOG, "a") as handle:
        handle.write(json.dumps(record) + "\n")


def read_log() -> list[dict]:
    if not os.path.exists(LOG):
        return []
    out = []
    with open(LOG) as handle:
        for line in handle:
            line = line.strip()
            if line:
                out.append(json.loads(line))
    return out


# ---------------------------------------------------------------------------
# commands
# ---------------------------------------------------------------------------


def cmd_list(args: argparse.Namespace) -> int:
    themes = {t.strip() for t in args.theme.split(",")} if args.theme else None
    print(f"\n{BOLD}Progressive problem catalog{OFF}  ({len(CATALOG)} problems)\n")
    print(f"  {'id':28} {'gates':>5}  {'budget':>7}  themes")
    print(f"  {rule('-', 70)}")
    for entry in CATALOG:
        if themes and not themes & set(entry["themes"]):
            continue
        state = read_state(entry["id"])
        if state and state.get("completed_at"):
            mark = f"{GREEN}✓{OFF}"
        elif state:
            mark = f"{YELLOW}·{OFF}"
        else:
            mark = " "
        print(
            f"{mark} {entry['id']:28} {len(entry['gates']):>5}  "
            f"{entry['budget_min']:>5} min  {', '.join(entry['themes'])}"
        )
    print(f"\n  {GREEN}✓{OFF} complete   {YELLOW}·{OFF} in progress   "
          f"{DIM}all 15 have automated gate tests{OFF}\n")
    print(f"  {DIM}Themes: state, concurrency, memory, streaming, parsing, scheduling{OFF}\n")
    return 0


def _print_gate(entry: dict, index: int) -> None:
    gate = entry["gates"][index]
    print(f"\n{rule('═')}")
    print(f"{BOLD}GATE {index + 1} of {len(entry['gates'])}{OFF} — {gate['title']}")
    print(rule("═"))
    print(gate["brief"].strip())
    print(rule("─"))
    print(
        f"{DIM}Narrate out loud: restate → clarify → approach → complexity → "
        f"test the tricky invariant → code.{OFF}"
    )
    print(f"Run: ./progressive.py test {entry['id']}\n")


def cmd_start(args: argparse.Namespace) -> int:
    entry = _entry(args.problem)
    if entry is None:
        return 1
    existing = read_state(entry["id"])
    if existing and not args.force:
        print(
            f"{YELLOW}Already started.{OFF} Use `brief` to re-read the current gate, "
            f"or `start --force` to restart."
        )
        return 1

    os.makedirs(work_dir(entry["id"]), exist_ok=True)
    attempt = os.path.join(work_dir(entry["id"]), "attempt.py")
    starter = os.path.join(entry["path"], "starter.py") if entry["automated"] else None
    if starter and os.path.exists(starter) and (not os.path.exists(attempt) or args.force):
        shutil.copy(starter, attempt)

    write_state(
        entry["id"],
        {
            "problem": entry["id"],
            "started_at": time.time(),
            "unlocked": 1,
            "gate_times": {},
            "attempts": {},
            "completed_at": None,
        },
    )

    print(f"\n{BOLD}{entry['title']}{OFF}")
    print(f"{DIM}{', '.join(entry['themes'])} · budget {entry['budget_min']} min · "
          f"{len(entry['gates'])} gates{OFF}")
    print(f"\nEdit: {attempt}")
    print(f"\n{BOLD}Timer started.{OFF} Gate 2 stays closed until gate 1 passes.")
    _print_gate(entry, 0)
    return 0


def cmd_brief(args: argparse.Namespace) -> int:
    entry = _entry(args.problem)
    if entry is None:
        return 1
    state = read_state(entry["id"])
    if state is None:
        print(f"Not started. Run: ./progressive.py start {entry['id']}")
        return 1
    if args.all:
        if not state.get("completed_at"):
            print(f"{RED}--all is only available after every gate passes.{OFF}")
            return 1
        for index in range(len(entry["gates"])):
            _print_gate(entry, index)
        return 0
    _print_gate(entry, state["unlocked"] - 1)
    return 0


def _load_attempt(entry: dict) -> Any:
    path = os.path.join(work_dir(entry["id"]), "attempt.py")
    if not os.path.exists(path):
        raise SystemExit(f"no attempt file at {path}")
    spec = importlib.util.spec_from_file_location(f"attempt_{entry['id']}", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def cmd_test(args: argparse.Namespace) -> int:
    entry = _entry(args.problem)
    if entry is None:
        return 1
    state = read_state(entry["id"])
    if state is None:
        print(f"Not started. Run: ./progressive.py start {entry['id']}")
        return 1
    if state.get("completed_at"):
        print(f"{GREEN}All gates already passed.{OFF} `reset` to run it again.")
        return 0

    index = state["unlocked"] - 1
    key = str(state["unlocked"])
    state["attempts"][key] = state["attempts"].get(key, 0) + 1

    if True:
        module = _load_attempt(entry)
        gate = entry["gates"][index]
        try:
            gate["test"](module)
        except NotImplementedError as exc:
            write_state(entry["id"], state)
            print(f"{YELLOW}GATE {state['unlocked']}: not implemented{OFF} ({exc})")
            return 1
        except AssertionError as exc:
            write_state(entry["id"], state)
            print(f"{RED}GATE {state['unlocked']}: FAIL{OFF} — {exc}")
            return 1
        except Exception:
            write_state(entry["id"], state)
            print(f"{RED}GATE {state['unlocked']}: ERROR{OFF}")
            traceback.print_exc(limit=8)
            return 1

    elapsed = time.time() - state["started_at"]
    state["gate_times"][key] = elapsed
    print(f"\n{GREEN}GATE {state['unlocked']}: PASS{OFF} at {mmss(elapsed)} "
          f"({state['attempts'][key]} run(s))")

    append_log(
        {
            "ts": time.time(),
            "problem": entry["id"],
            "gate": state["unlocked"],
            "elapsed_s": round(elapsed, 1),
            "attempts": state["attempts"][key],
        }
    )

    if state["unlocked"] >= len(entry["gates"]):
        state["completed_at"] = time.time()
        write_state(entry["id"], state)
        budget = entry["budget_min"] * 60
        verdict = (
            f"{GREEN}inside budget{OFF}" if elapsed <= budget else f"{RED}over budget{OFF}"
        )
        print(f"\n{BOLD}All {len(entry['gates'])} gates passed{OFF} in {mmss(elapsed)} "
              f"(budget {entry['budget_min']}m) — {verdict}")
        first = state["gate_times"].get("1")
        if first:
            print(f"Time-to-first-gate: {BOLD}{mmss(first)}{OFF}")
        chapter = entry.get("chapter") or "tracks/coding/WARMUP.md"
        print(f"\n{DIM}Now read the pattern from first principles:{OFF}")
        print(f"{DIM}  swe-interview-prep/{chapter}{OFF}")
        print(f"{DIM}Anything you got wrong goes into ../../../review/ at 1 day.{OFF}\n")
        return 0

    state["unlocked"] += 1
    write_state(entry["id"], state)
    _print_gate(entry, state["unlocked"] - 1)
    return 0


def cmd_status(args: argparse.Namespace) -> int:
    ids = [args.problem] if args.problem else [e["id"] for e in CATALOG]
    rows = []
    for problem_id in ids:
        state = read_state(problem_id)
        if state is None:
            continue
        entry = _entry(problem_id, quiet=True)
        gates = len(entry["gates"]) if entry else 0
        passed = len(state["gate_times"])
        first = state["gate_times"].get("1")
        total = (state.get("completed_at") or time.time()) - state["started_at"]
        rows.append((problem_id, f"{passed}/{gates}", mmss(first) if first else "—",
                     mmss(total), "done" if state.get("completed_at") else "open"))
    if not rows:
        print("Nothing started yet.")
        return 0
    print(f"\n  {'problem':28} {'gates':>6} {'to-G1':>7} {'total':>8}  state")
    print(f"  {rule('-', 68)}")
    for row in rows:
        print(f"  {row[0]:28} {row[1]:>6} {row[2]:>7} {row[3]:>8}  {row[4]}")
    print()
    return 0


def cmd_log(args: argparse.Namespace) -> int:
    records = read_log()
    if not records:
        print("No timing log yet.")
        return 0
    for record in records[-args.tail :]:
        stamp = time.strftime("%Y-%m-%d %H:%M", time.localtime(record["ts"]))
        print(
            f"{stamp}  {record['problem']:28} gate {record['gate']}  "
            f"{mmss(record['elapsed_s']):>7}  {record['attempts']} run(s)"
        )
    return 0


def cmd_chart(args: argparse.Namespace) -> int:
    """ASCII trend of time-to-first-passing-gate, in chronological order."""
    firsts = [r for r in read_log() if r["gate"] == 1]
    if not firsts:
        print("No gate-1 passes logged yet. The chart needs at least one.")
        return 0
    firsts.sort(key=lambda r: r["ts"])
    values = [r["elapsed_s"] / 60 for r in firsts]
    peak = max(values)
    width = 44
    print(f"\n{BOLD}Time-to-first-passing-gate{OFF}  {DIM}(minutes, chronological){OFF}\n")
    for record, value in zip(firsts, values):
        bar = "█" * max(1, int(round(value / peak * width)))
        stamp = time.strftime("%m-%d", time.localtime(record["ts"]))
        print(f"  {stamp}  {record['problem'][:22]:22} {bar} {value:.1f}")
    recent = values[-5:]
    early = values[: max(1, len(values) // 3)]
    print(f"\n  first {len(early)} avg: {sum(early)/len(early):.1f} min")
    print(f"  last  {len(recent)} avg: {sum(recent)/len(recent):.1f} min")
    if sum(recent) / len(recent) < sum(early) / len(early):
        print(f"  {GREEN}trending down — this is the number that predicts the round{OFF}\n")
    else:
        print(f"  {YELLOW}not trending down yet — see tracks/coding/README.md "
              f"§ Time-to-first-correct{OFF}\n")
    return 0


def cmd_reset(args: argparse.Namespace) -> int:
    entry = _entry(args.problem)
    if entry is None:
        return 1
    target = work_dir(entry["id"])
    if os.path.exists(target):
        if not args.keep_code:
            shutil.rmtree(target)
        else:
            os.remove(state_path(entry["id"]))
        print(f"Reset {entry['id']}.")
    else:
        print("Nothing to reset.")
    return 0


# ---------------------------------------------------------------------------


def _entry(problem_id: str, quiet: bool = False) -> dict | None:
    for entry in CATALOG:
        if entry["id"] == problem_id:
            return load_problem(entry) if entry["automated"] else entry
    if not quiet:
        print(f"{RED}unknown problem: {problem_id}{OFF}")
        print("Run `./progressive.py list` for the catalog.")
    return None


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="progressive.py", description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("list", help="show the catalog")
    p.add_argument("--theme", help="filter by comma-separated themes")
    p.set_defaults(func=cmd_list)

    p = sub.add_parser("start", help="begin a problem and open gate 1")
    p.add_argument("problem")
    p.add_argument("--force", action="store_true", help="restart, overwriting attempt.py")
    p.set_defaults(func=cmd_start)

    p = sub.add_parser("brief", help="re-print the current gate")
    p.add_argument("problem")
    p.add_argument("--all", action="store_true", help="all gates (only once complete)")
    p.set_defaults(func=cmd_brief)

    p = sub.add_parser("test", help="run the current gate; unlock the next on pass")
    p.add_argument("problem")
    p.set_defaults(func=cmd_test)

    p = sub.add_parser("status", help="in-flight problems")
    p.add_argument("problem", nargs="?")
    p.set_defaults(func=cmd_status)

    p = sub.add_parser("log", help="timing history")
    p.add_argument("--tail", type=int, default=40)
    p.set_defaults(func=cmd_log)

    p = sub.add_parser("chart", help="time-to-first-gate trend")
    p.set_defaults(func=cmd_chart)

    p = sub.add_parser("reset", help="start a problem over")
    p.add_argument("problem")
    p.add_argument("--keep-code", action="store_true", help="keep attempt.py")
    p.set_defaults(func=cmd_reset)

    args = parser.parse_args(argv)
    return args.func(args)


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