td_lang/cli.py

"""
TD Lang CLI — Command-line interface for .td files.

Usage:
    python -m td_lang run examples/demo_merge.td       # Compile + execute
    python -m td_lang compile examples/demo_merge.td   # Compile only (outputs .py)
    python -m td_lang check examples/demo_merge.td     # Syntax check only
    python -m td_lang info examples/demo_merge.td      # Show plan without compiling
    python -m td_lang --version                        # Show version
"""

import argparse
import sys

from . import __version__
from .executor import TDExecutor
from .errors import TDLangError
from .grammar import parse_td_file
from .ast_nodes import (
    LoadCmd, MergeCmd, HealCmd, EvalCmd, CommitCmd,
    SynthCmd, TrainCmd, DebateCmd, DiagnoseCmd,
    ForkCmd, ResetCmd, PruneCmd, EditCmd,
    FuseCmd, AbsorbCmd, RepeatBlock, IfBlock,
    NotifyCmd, SaveCmd, ScheduleCmd,
    DownloadCmd, LogBlock, CompareCmd, VerifyCmd,
    VoteCmd, PromptBlock, DistillCmd, RollbackCmd,
    CurriculumCmd, StarCmd, BestOfCmd, ExploitCmd, ArenaCmd, ResearchArenaCmd,
    SnapshotCmd, ReportCmd,
)


# Phase labels for info command
_PHASE_MAP = {
    LoadCmd: ("1", "load"),
    MergeCmd: ("1", "merge"),
    HealCmd: ("1", "heal"),
    EvalCmd: ("1", "eval"),
    CommitCmd: ("1", "commit"),
    SynthCmd: ("2", "synth"),
    TrainCmd: ("2", "train"),
    DebateCmd: ("2", "debate"),
    DiagnoseCmd: ("2", "diagnose"),
    ForkCmd: ("3", "fork"),
    ResetCmd: ("3", "reset"),
    PruneCmd: ("3", "prune"),
    EditCmd: ("3", "edit"),
    FuseCmd: ("6", "fuse"),
    AbsorbCmd: ("6", "absorb"),
    RepeatBlock: ("7", "repeat"),
    IfBlock: ("7", "if"),
    NotifyCmd: ("8", "notify"),
    SaveCmd: ("8", "save"),
    SnapshotCmd: ("4", "snapshot"),
    ReportCmd: ("4", "report"),
    ScheduleCmd: ("9", "schedule"),
    DownloadCmd: ("10", "download"),
    CompareCmd: ("10", "compare"),
    VerifyCmd: ("10", "verify"),
    VoteCmd: ("11", "vote"),
    PromptBlock: ("11", "prompt"),
    DistillCmd: ("11", "distill"),
    RollbackCmd: ("11", "rollback"),
    CurriculumCmd: ("12", "curriculum"),
    StarCmd: ("12", "star"),
    BestOfCmd: ("12", "best_of"),
    ExploitCmd: ("12", "exploit"),
    ArenaCmd: ("13", "arena"),
    ResearchArenaCmd: ("13", "research_arena"),
}


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="TD Lang — compile and run .td files for Time Dilation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python -m td_lang check examples/demo_merge.td      # Check syntax
  python -m td_lang compile examples/demo_merge.td    # Compile to .py
  python -m td_lang run examples/demo_merge.td        # Compile + run
  python -m td_lang run examples/demo_merge.td --dry   # Compile only
  python -m td_lang info examples/demo_merge.td        # Show plan summary
        """,
    )

    parser.add_argument(
        "--version",
        action="version",
        version=f"td_lang {__version__}",
    )

    parser.add_argument(
        "action",
        choices=["check", "compile", "run", "info"],
        help="What to do: check (syntax), compile (.py), run (compile+execute), info (show plan)",
    )

    parser.add_argument(
        "file",
        type=str,
        help="Path to the .td file",
    )

    parser.add_argument(
        "--output",
        type=str,
        default="td_lang_outputs",
        help="Output directory (default: td_lang_outputs)",
    )

    parser.add_argument(
        "--dry",
        action="store_true",
        help="With 'run': compile but don't execute",
    )

    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show extra detail (compiled Python, full AST, etc.)",
    )

    return parser.parse_args()


def print_banner():
    """Print the td_lang banner."""
    banner = f"""
    ╔═══════════════════════════════════════╗
    ║                                       ║
    ║   ████████╗██████╗    ██╗      ██████╗║
    ║   ╚══██╔══╝██╔══██╗   ██║     ██╔════╝║
    ║      ██║   ██║  ██║   ██║     ██║  ███║
    ║      ██║   ██║  ██║   ██║     ██║   ██║
    ║      ██║   ██████╔╝   ██████╗ ╚██████╔╝║
    ║      ╚═╝   ╚═════╝    ╚═════╝  ╚═════╝║
    ║                                       ║
    ║   TD Lang v{__version__} — .td file compiler   ║
    ║                                       ║
    ╚═══════════════════════════════════════╝
    """
    print(banner)


def print_info(filepath: str) -> None:
    """Show what a .td file does without compiling — human-readable plan summary."""
    program = parse_td_file(filepath)

    print(f"\n  File: {filepath}")
    print(f"  Commands: {len(program.commands)}")

    if program.gates:
        print(f"  Gates: {', '.join(program.gates.must_pass)}")
    if program.budget:
        parts = []
        if program.budget.max_gpu_hours is not None:
            parts.append(f"{program.budget.max_gpu_hours} GPU hrs")
        if program.budget.max_cost is not None:
            parts.append(f"${program.budget.max_cost}")
        print(f"  Budget: {', '.join(parts)}")
    if program.data_contract:
        print(f"  Data contract: fields={program.data_contract.required_fields}")
    if program.reward_contract:
        print(f"  Reward contract: verifiers={program.reward_contract.verifiers}")

    print("\n  Plan:")
    for i, cmd in enumerate(program.commands, 1):
        phase, name = _PHASE_MAP.get(type(cmd), ("?", type(cmd).__name__))
        target = getattr(cmd, 'target', getattr(cmd, 'alias', ''))
        detail = ""
        if hasattr(cmd, 'method'):
            detail += f" method={cmd.method}"
        if hasattr(cmd, 'source') and name in ("merge", "synth"):
            detail += f" from={cmd.source}"
        if hasattr(cmd, 'layers') and cmd.layers != "all":
            detail += f" layers={cmd.layers}"
        if hasattr(cmd, 'output') and cmd.output:
            detail += f" -> {cmd.output}"
        print(f"    {i}. [P{phase}] {name} {target}{detail}")

    print()


def main():
    """Main entry point for td_lang CLI."""
    args = parse_args()
    print_banner()

    executor = TDExecutor(output_dir=args.output)

    try:
        if args.action == "info":
            print_info(args.file)

        elif args.action == "check":
            program = executor.check(args.file)
            print("\n[td_lang] File is valid!")

        elif args.action == "compile":
            py_path = executor.compile(args.file)
            print(f"\n[td_lang] Generated: {py_path}")
            print("[td_lang] You can run it with: python", py_path)
            if args.verbose:
                print("\n--- Generated Python ---")
                print(py_path.read_text())
                print("--- End ---")

        elif args.action == "run":
            result = executor.run(args.file, dry_run=args.dry)
            if result["status"] == "success":
                sys.exit(0)
            elif result["status"] == "dry_run":
                sys.exit(0)
            else:
                sys.exit(1)

    except TDLangError as e:
        print(f"\n[td_lang] ERROR: {e}")
        sys.exit(1)

    except FileNotFoundError:
        print(f"\n[td_lang] ERROR: File not found: {args.file}")
        print("[td_lang] Check the path and try again.")
        sys.exit(1)

    except KeyboardInterrupt:
        print("\n[td_lang] Interrupted.")
        sys.exit(130)