cli.py

#!/usr/bin/env python3
"""Terminal entry point. The only module that does terminal I/O; the aicoe
package itself stays UI-free."""

from __future__ import annotations

import argparse
import getpass
import itertools
import logging
import os
import re
import sys
import threading
import time

import anthropic

from aicoe.progress import LoggingReporter as _LoggingReporter


class TerminalReporter(_LoggingReporter):
    """LoggingReporter plus an in-place spinner (elapsed + tokens) for long
    calls, redrawn on a timer so it stays alive during the silent thinking
    phase; log lines clear it first."""

    _FRAMES = "✻✢✳∗·"

    def __init__(self):
        super().__init__()
        self._lock = threading.Lock()
        self._active = False
        self._thread = None
        self._label = ""
        self._start = 0.0
        self._tokens = 0
        self._tty = sys.stdout.isatty()

    def _clear_line(self):
        if self._tty:
            sys.stdout.write("\r\033[K")
            sys.stdout.flush()

    def _emit(self, fn, message):
        with self._lock:
            self._clear_line()
            fn(message)

    def info(self, message):
        self._emit(super().info, message)

    def warn(self, message):
        self._emit(super().warn, message)

    def stage(self, message):
        self._emit(super().stage, message)

    def debug(self, message):
        self._emit(super().debug, message)

    def start_activity(self, label):
        if not self._tty or self._active:
            return
        self._label, self._tokens, self._start, self._active = label, 0, time.time(), True
        self._thread = threading.Thread(target=self._spin, daemon=True)
        self._thread.start()

    def update_activity(self, output_tokens=None, label=None):
        if output_tokens is not None:
            self._tokens = output_tokens
        if label is not None:
            self._label = label

    def stop_activity(self):
        if not self._active:
            return
        self._active = False
        if self._thread is not None:
            self._thread.join(timeout=1)
        with self._lock:
            self._clear_line()

    def _spin(self):
        frames = itertools.cycle(self._FRAMES)
        while self._active:
            with self._lock:
                if not self._active:
                    break
                elapsed = int(time.time() - self._start)
                tok = self._tokens
                tokstr = f"{tok / 1000:.1f}k" if tok >= 1000 else str(tok)
                sys.stdout.write(
                    f"\r\033[K{next(frames)} {self._label}… ({elapsed}s · ↓ {tokstr} tokens)"
                )
                sys.stdout.flush()
            time.sleep(0.15)


# Bracketed paste / focus reporting / stray CSI+SS3 sequences land verbatim in
# getpass/input; an invisible control byte in a pasted API key corrupts the
# x-api-key header (bare 400 from the API edge), so all prompts strip them.
_BRACKETED_PASTE_RE = re.compile(r"\x1b?\[20[01]~")
_ANSI_ESCAPE_RE = re.compile(r"\x1b(?:\[[0-9;?]*[ -/]*[@-~]|O[@-~]|[@-Z\\\]^_])")


def _strip_terminal_noise(raw: str | None) -> str:
    if not raw:
        return ""
    return _ANSI_ESCAPE_RE.sub("", _BRACKETED_PASTE_RE.sub("", raw))


def _clean_key(raw: str | None) -> str:
    return "".join(ch for ch in _strip_terminal_noise(raw) if ch.isprintable() and not ch.isspace())


def _clean_line(raw: str | None) -> str:
    return "".join(ch for ch in _strip_terminal_noise(raw) if ch.isprintable()).strip()


def _init_terminal():
    """Disable bracketed paste (?2004) and focus reporting (?1004), which
    inject escape bytes into interactive reads."""
    if _is_tty():
        sys.stdout.write("\x1b[?2004l\x1b[?1004l")
        sys.stdout.flush()

from aicoe import runs as runs_mod
from aicoe.claude_client import ClaudeClient, MODEL_CHOICES, FABLE, OPUS, SONNET
from aicoe.config import RunConfig, RESEARCH_CONTEXT_FILE
from aicoe.pipeline import run, stage_status
from aicoe.progress import LoggingReporter

DEFAULT_INPUT = os.path.join(os.path.dirname(__file__), "examples", "sample_interviews")


def _is_tty() -> bool:
    return sys.stdin.isatty() and sys.stdout.isatty()


def resolve_api_key(model: str) -> str:
    """ANTHROPIC_API_KEY or hidden prompt, validated with a 1-token request.
    Only a clearly-bad key re-prompts; other errors proceed (the run surfaces
    them anyway)."""
    key = _clean_key(os.environ.get("ANTHROPIC_API_KEY"))
    while True:
        if not key:
            if not _is_tty():
                sys.exit("ANTHROPIC_API_KEY not set and no TTY to prompt. Set it and retry.")
            key = _clean_key(getpass.getpass("Anthropic API key (input hidden): "))
            if not key:
                print("No key read; try again.")
                continue
        try:
            anthropic.Anthropic(api_key=key).messages.create(
                model=model, max_tokens=1, messages=[{"role": "user", "content": "ping"}]
            )
            return key
        except anthropic.AuthenticationError:
            print("Invalid API key; try again.")
            key = None
        except anthropic.PermissionDeniedError:
            print(f"This key cannot access {model}; proceeding (the run may fail).")
            return key
        except Exception as e:  # network/endpoint quirk; don't block the run
            print(f"Could not validate key ({type(e).__name__}); proceeding anyway.")
            return key


def choose_model(args) -> str:
    if args.model:
        return MODEL_CHOICES.get(args.model.lower(), args.model)
    if not _is_tty():
        return SONNET
    print("\nChoose a model:")
    print(f"  1) Sonnet ({SONNET}) — faster, cheaper [default]")
    print(f"  2) Opus   ({OPUS}) — most capable")
    print(f"  3) Fable   ({FABLE}) — 1M context")
    choice = _clean_line(input("Selection [1/2/3]: "))
    return {"2": OPUS, "3": FABLE}.get(choice, SONNET)


def choose_input_folder(args) -> str:
    if args.input:
        return args.input
    if not _is_tty():
        return DEFAULT_INPUT
    print(
        "\nInput folder: ALL files in it must be the SAME format "
        "(.vtt parsed natively; .txt/.srt/.pdf supported; any other single\n"
        "format is handled by a Claude-generated processor). Results are written "
        "into per-stage subfolders of this same folder."
    )
    folder = _clean_line(input(f"Input folder [{DEFAULT_INPUT}]: "))
    return folder or DEFAULT_INPUT


def resolve_research_context(results_dir, args):
    """Load/confirm/prompt for the run's RESEARCH_CONTEXT.md (new runs are
    seeded from the input folder's file). Returns the context string or None."""
    path = os.path.join(results_dir, RESEARCH_CONTEXT_FILE)

    if os.path.exists(path):
        try:
            text = open(path, encoding="utf-8").read().strip()
        except Exception as e:  # noqa: BLE001
            print(f"(could not read {RESEARCH_CONTEXT_FILE}: {e})")
            return None
        print(f"\nFound existing {RESEARCH_CONTEXT_FILE}:\n{'-' * 40}\n{text}\n{'-' * 40}")
        if args.yes or not _is_tty():
            return text or None
        answer = _clean_line(input("Use this research context? [Y]es / [n]o (re-enter): ")).lower()
        if answer in ("", "y", "yes", "c", "continue"):
            return text or None
        # Re-enter; keep the existing text if they enter nothing.
        new = _clean_line(
            input(
                f"\nEnter the new research context/goals "
                f"(Enter alone keeps the current {RESEARCH_CONTEXT_FILE}): "
            )
        )
        if not new:
            return text or None
        _write_research_context(path, new)
        return new

    # No file yet: optional single-line prompt.
    if args.yes or not _is_tty():
        return None
    text = _clean_line(
        input(
            "\nOptionally describe your research context/goals (used in the "
            "coding prompts). Press Enter to skip: "
        )
    )
    if not text:
        return None
    _write_research_context(path, text)
    return text


def _write_research_context(path, text):
    try:
        os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            f.write(text.strip() + "\n")
        print(f"Saved research context to {path}")
    except Exception as e:  # noqa: BLE001
        print(f"(could not save {RESEARCH_CONTEXT_FILE}: {e})")


def print_runs(input_folder):
    """List all runs (and the legacy pseudo-run) for an input folder."""
    entries = runs_mod.list_runs(input_folder)
    if not entries:
        print(f"No runs in {input_folder} yet.")
        return
    current = runs_mod.get_current(input_folder)
    print(f"\nRuns in {input_folder}:")
    for m in entries:
        marks = []
        if m["run_id"] == current:
            marks.append("current")
        if m.get("parent_run"):
            marks.append(f"forked from {m['parent_run']}")
        suffix = f"  [{', '.join(marks)}]" if marks else ""
        label = f" — {m['label']}" if m.get("label") else ""
        created = f" ({m['created'][:10]})" if m.get("created") else ""
        print(f"  {m['run_id']}{label}{created}{suffix}")


def choose_run(args, input_folder):
    """Resolve which run this invocation works in; may create or prompt."""
    if args.new_run is not None:
        run_id = runs_mod.create_run(input_folder, label=args.new_run)
        print(f"Created run '{run_id}'.")
        return run_id
    if args.run:
        return runs_mod.resolve_run(input_folder, args.run)

    default = runs_mod.resolve_run(input_folder)
    existing = runs_mod.list_runs(input_folder)
    if args.yes or not _is_tty() or len(existing) <= 1:
        return default

    print("\nRuns in this folder:")
    for i, m in enumerate(existing, 1):
        mark = " (current)" if m["run_id"] == default else ""
        label = f" — {m['label']}" if m.get("label") else ""
        print(f"  {i}) {m['run_id']}{label}{mark}")
    print("  n) new run")
    choice = _clean_line(input(f"Select run [{default}]: ")).lower()
    if not choice:
        return default
    if choice == "n":
        label = _clean_line(input("Label for the new run (optional): "))
        run_id = runs_mod.create_run(input_folder, label=label)
        print(f"Created run '{run_id}'.")
        return run_id
    if choice.isdigit() and 1 <= int(choice) <= len(existing):
        run_id = existing[int(choice) - 1]["run_id"]
    else:
        run_id = choice
    run_id = runs_mod.resolve_run(input_folder, run_id)
    if run_id != runs_mod.LEGACY_RUN_ID:
        runs_mod.set_current(input_folder, run_id)
    return run_id


def warn_if_stale(input_folder, run_id):
    """Print a warning when the raw inputs drifted since the run was created."""
    info = runs_mod.staleness(input_folder, runs_mod.load_manifest(input_folder, run_id))
    if not runs_mod.is_stale(info):
        return
    bits = [f"{kind} {', '.join(names)}" for kind, names in info.items() if names]
    print(
        f"\nWARNING: input files changed since run '{run_id}' was created "
        f"({'; '.join(bits)}). This run keeps coding its frozen snapshots; "
        f"create a new run to ingest the current files."
    )


def print_stage_status(results_dir, use_description, input_dir=None):
    """Show what's already on disk for this run, before picking stages."""
    try:
        statuses = stage_status(results_dir, use_description=use_description, input_dir=input_dir)
    except Exception as e:  # never block the run on a status hiccup
        print(f"(could not read existing progress: {e})")
        return
    stage_models = (runs_mod.load_manifest_at(results_dir) or {}).get("stage_models") or {}
    width = max(len(f"Stage {n} ({name})") for n, name, _ in statuses)
    print(f"\nPipeline status for {results_dir}:")
    for n, name, status in statuses:
        label = f"Stage {n} ({name})"
        model = stage_models.get(str(n), "")
        suffix = f"  [{model}]" if model else ""
        print(f"  {label.ljust(width)}  - {status}{suffix}")


def choose_stages(args):
    if args.stages:
        return tuple(args.stages)
    if not _is_tty():
        return (1, 2, 3, 4, 5)
    raw = _clean_line(input("\nStages to run [1 2 3 4 5]: "))
    if not raw:
        return (1, 2, 3, 4, 5)
    return tuple(int(x) for x in raw.replace(",", " ").split())


def make_confirm(assume_yes: bool):
    """Return a confirm(step) -> bool gate. Aborts cleanly on 'a'/no."""
    def confirm(step: str) -> bool:
        if assume_yes or not _is_tty():
            print(f"-> {step}")
            return True
        answer = _clean_line(input(f"\n-> Next: {step}.  [C]ontinue / [A]bort? ")).lower()
        return answer in ("", "c", "continue", "y", "yes")

    return confirm


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="AICoE: qualitative coding via the Claude API.")
    p.add_argument("--model", help="'sonnet', 'opus', or an explicit model id")
    p.add_argument("--input", help="folder of same-format transcripts (results land in its subfolders)")
    p.add_argument(
        "--stages",
        nargs="*",
        type=int,
        help="stages to run, e.g. --stages 1 2 3 4 5 (5 = findings synthesis)",
    )
    p.add_argument("--codes-file", help="skip stages 1-2; load comprehensive codes from this JSON")
    p.add_argument("--run", help="work in this run id (see --list-runs); 'legacy' = flat pre-runs layout")
    p.add_argument(
        "--new-run",
        nargs="?",
        const="",
        metavar="LABEL",
        help="create a fresh run (optionally labelled) and work in it",
    )
    p.add_argument("--list-runs", action="store_true", help="list runs for the input folder and exit")
    p.add_argument("--export-qdpx", metavar="PATH", help="write the run as a REFI-QDA .qdpx and exit")
    p.add_argument("--use-description", action="store_true", help="use code descriptions")
    p.add_argument(
        "--no-chunk",
        action="store_true",
        help="disable utterance/turn chunking preprocessing (keep raw lines)",
    )
    p.add_argument(
        "--workers",
        type=int,
        default=5,
        help="max concurrent API requests within a stage (default 5; 1 = sequential)",
    )
    p.add_argument("--yes", action="store_true", help="auto-continue every step (no prompts)")
    p.add_argument("--verbose", action="store_true", help="debug logging (shows token usage)")
    return p.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    # Log to stdout so the activity spinner (also stdout) and log lines share one
    # stream and the in-place line clearing works cleanly.
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(message)s",
        stream=sys.stdout,
    )
    # --verbose raises aicoe's own verbosity only; keep HTTP libraries quiet so
    # the terminal isn't flooded with request/connection internals.
    for noisy in ("anthropic", "httpx", "httpcore", "urllib3"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    _init_terminal()
    reporter = TerminalReporter()

    if args.list_runs:
        print_runs(args.input or choose_input_folder(args))
        return

    if args.export_qdpx:  # local, no API key needed
        from aicoe import qdpx
        input_folder = choose_input_folder(args)
        run_id = runs_mod.resolve_run(input_folder, args.run)
        data = qdpx.build_qdpx(runs_mod.run_dir(input_folder, run_id), project_name=f"AICoE {run_id}")
        with open(args.export_qdpx, "wb") as f:
            f.write(data)
        print(f"Wrote {args.export_qdpx} (run '{run_id}').")
        return

    model = choose_model(args)
    api_key = resolve_api_key(model)
    input_folder = choose_input_folder(args)
    try:
        run_id = choose_run(args, input_folder)
    except ValueError as e:
        sys.exit(str(e))
    results_dir = runs_mod.run_dir(input_folder, run_id)
    research_context = resolve_research_context(results_dir, args)
    warn_if_stale(input_folder, run_id)
    print_stage_status(results_dir, args.use_description, input_dir=input_folder)
    stages = choose_stages(args)

    config = RunConfig(
        results_dir=results_dir,
        input_dir=input_folder,
        use_description=args.use_description,
        stages=stages,
        codes_file=args.codes_file,
        chunk_utterances=not args.no_chunk,
        context=research_context,
        max_workers=max(1, args.workers),
    )
    client = ClaudeClient(api_key=api_key, model=model, reporter=reporter)
    confirm = make_confirm(args.yes)

    print(f"\nModel: {model}\nInput folder: {input_folder}\nRun: {run_id} → {results_dir}\nStages: {stages}\n")
    try:
        run(config, client, reporter=reporter, confirm=confirm)
    except KeyboardInterrupt:
        print("\nInterrupted.")
    except Exception as e:  # noqa: BLE001
        reporter.warn(f"Run failed: {e}")
        raise


if __name__ == "__main__":
    main()