"""Bandito CLI entry point."""

from __future__ import annotations

import argparse


def get_client():
    """Create a BanditoClient or print signup message and return None."""
    from .api import BanditoClient

    try:
        return BanditoClient()
    except RuntimeError:
        print("\n  Bandito API key required.")
        print("  Sign up (free):  bandito signup")
        print("  Then set:        export BANDITO_API_KEY=bnd_...\n")
        return None


# -- Help text -----------------------------------------------------------------

_HELP = """\

  bandito — the AI engineer for LLM applications

    signup        Create an account
    setup         Configure local settings and org storage
    observe       See your projects, traces, and stats
    tui           Interactive grading workbench
    judge         LLM-as-judge (calibrate, run)
    diagnostics   System health check
    audit         Data integrity checks
    analyze       Deep LLM-powered analysis
    assert        Run assertions against traces
    dataset       Curate evaluation datasets
    improve       Build datasets and test changes
    skill         Install Claude Code skills for Bandito

  Run bandito <command> to see available options.
"""

_SETUP_HELP = """\

  bandito setup

    config          Configure local settings and org trace storage
"""

_OBSERVE_HELP = """\

  bandito observe

    projects    List projects and stats
    traces      List recent traces
    trace <id>  Inspect a single trace

  For system health: bandito diagnostics --project NAME
"""

_JUDGE_HELP = """\

  bandito judge

    calibrate   Create rubric (if needed) and calibrate judge
    run         Score ungraded traces with calibrated judge
    overview    Calibration state and judge usage at a glance
    runs        Judge run history
"""

_IMPROVE_HELP = """\

  bandito improve

    dataset     Curate evaluation datasets from production traces
    replay      Test model swaps against graded data

  These build on findings from bandito analyze.
"""

_SKILL_HELP = """\

  bandito skill

    install    Install Claude Code skill for Bandito

  The skill teaches Claude Code to use Bandito for debugging and analyzing
  your LLM applications. After install, invoke it with /bandito in Claude Code.
"""

_DATASET_HELP = """\

  bandito dataset

    create       Create a new dataset from traces
    list         List datasets for a project
    inspect      Inspect a dataset's contents
    compare      Compare dataset vs production (drift detection)
    regression   Run diagnostics on dataset
    export       Export dataset to JSONL, pytest, or pairs

  Examples:
    bandito dataset create golden-set --project my-bot --human-score-eq 1
    bandito dataset list --project my-bot
    bandito dataset compare golden-set --project my-bot
    bandito dataset export golden-set --project my-bot --format pytest --output test.py
"""

_ANALYZE_DESCRIPTIONS = {
    "tradeoffs": (
        "Cost, quality, and latency optimization.\n"
        "  Find where you're overpaying, which segments can use cheaper models,\n"
        "  and where latency is bottlenecked — all with quality impact estimates.\n"
        "  Example: \"gpt-4o-mini handles 62% of queries at 91% quality for 12% cost.\""
    ),
    "drift": (
        "Detect what changed and when.\n"
        "  Changepoint detection on quality, cost, and latency. Correlate shifts\n"
        "  with deploys, model updates, or traffic pattern changes.\n"
    ),
    "coverage": (
        "Find where you're blind.\n"
        "  Eval coverage gaps, unscored segments, judge calibration staleness,\n"
        "  and areas where you lack data to make confident decisions.\n"
        "  Example: \"0 human grades on code queries (30% of traffic). Judge uncalibrated for 2 weeks.\""
    ),
}




# -- Shared parser helpers -----------------------------------------------------

def _add_action_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--project", default=None, metavar="NAME", help="Project name")
    parser.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")
    parser.add_argument("--last", type=int, default=None, metavar="N", help="Only use the last N traces")
    parser.add_argument("--days", type=int, default=14, metavar="N", help="Analyze last N days (default: 14)")
    parser.add_argument("--all", action="store_true", dest="all_traces", help="Analyze all traces (ignore --days)")


def _stub(name: str, description: str) -> None:
    print(f"\n  bandito {name}\n")
    if description:
        for line in description.split("\n"):
            print(f"  {line}")
        print()
    print("  Coming soon.\n")
    print("  Available now:")
    print("    bandito observe projects    See your projects")
    print("    bandito observe traces      List recent traces")
    print("    bandito tui                 Grade traces interactively")
    print("    bandito judge overview      Judge health at a glance")
    print()


def _analyze_help() -> None:
    print("\n  bandito analyze <surface>\n")
    print("  LLM-powered analysis that produces actionable recommendations.")
    print("  Built on foundation layer: segmentation, span breakdown, clustering.\n")
    for name, desc in _ANALYZE_DESCRIPTIONS.items():
        first_line = desc.split("\n")[0].strip()
        print(f"    {name:<14s} {first_line}")
    print()
    print("  Requires BANDITO_API_KEY.")
    print("  Run bandito diagnostics first for a quick system health check.")
    print()


# -- Main entry point ----------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(prog="bandito", add_help=False)
    parser.add_argument("-h", "--help", action="store_true")
    subs = parser.add_subparsers(dest="group")

    # -- signup (top-level) ----------------------------------------------------

    subs.add_parser("signup", help="Create an account")

    # -- setup -----------------------------------------------------------------

    g_setup = subs.add_parser("setup", add_help=False)
    g_setup.add_argument("-h", "--help", action="store_true")
    setup_subs = g_setup.add_subparsers(dest="command")

    setup_subs.add_parser("config", help="Configure local settings and org trace storage")

    # -- observe ---------------------------------------------------------------

    g_observe = subs.add_parser("observe", add_help=False)
    g_observe.add_argument("-h", "--help", action="store_true")
    observe_subs = g_observe.add_subparsers(dest="command")

    projects_cmd = observe_subs.add_parser("projects", help="List projects and their stats")
    projects_cmd.add_argument("--project", default=None, metavar="NAME", help="Resolve storage from this project's config")

    basics = observe_subs.add_parser(
        "basics",
        help="Quick stats and detected patterns",
        description="Aggregate stats, model breakdown, and pattern detection. No LLM required.",
    )
    _add_action_args(basics)

    traces_cmd = observe_subs.add_parser("traces", help="List recent traces")
    traces_cmd.add_argument("--project", default=None, metavar="NAME")
    traces_cmd.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")
    traces_cmd.add_argument("--last", type=int, default=10, metavar="N", help="Number of traces to show (default: 10, 0 for all)")
    traces_cmd.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")
    traces_cmd.add_argument("--no-cache", action="store_true", help="Force fresh data, bypass cache")

    trace_cmd = observe_subs.add_parser("trace", help="Inspect a single trace")
    trace_cmd.add_argument("trace_id", help="Trace ID to inspect")
    trace_cmd.add_argument("--project", default=None, metavar="NAME")
    trace_cmd.add_argument("--no-cache", action="store_true", help="Force fresh data, bypass cache")

    # -- tui -------------------------------------------------------------------

    g_tui = subs.add_parser("tui", help="Launch interactive TUI (dashboard & grading)")
    g_tui.add_argument("--project", default=None, metavar="NAME", help="Start with this project selected")
    g_tui.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")

    # -- judge -----------------------------------------------------------------

    g_judge = subs.add_parser("judge", add_help=False)
    g_judge.add_argument("-h", "--help", action="store_true")
    judge_subs = g_judge.add_subparsers(dest="command")
    judge_cal = judge_subs.add_parser("calibrate", help="Create rubric (if needed) and calibrate judge")
    judge_cal.add_argument("--project", required=True, metavar="NAME")
    judge_cal.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")
    judge_cal.add_argument("-n", type=int, default=None, metavar="COUNT", help="Max traces to calibrate against")
    judge_cal.add_argument("--dry-run", action="store_true", help="Show prompt for first trace, don't call API")
    judge_cal.add_argument("--regenerate", action="store_true", help="Re-generate rubric from traces + LLM even if one exists")
    judge_run = judge_subs.add_parser("run", help="Score ungraded traces with calibrated judge")
    judge_run.add_argument("--project", required=True, metavar="NAME")
    judge_run.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")
    judge_run.add_argument("--sample", type=int, default=None, metavar="N", help="Randomly sample N traces to judge")
    judge_run.add_argument("--rejudge", action="store_true", help="Re-score traces that already have judge evals")
    judge_run.add_argument("--dry-run", action="store_true", help="Show prompt for first trace, don't call API")
    judge_run.add_argument("--dataset", default=None, metavar="NAME", help="Run on traces from a dataset")
    judge_overview = judge_subs.add_parser("overview", help="Calibration state and judge usage at a glance")
    judge_overview.add_argument("--project", required=True, metavar="NAME")
    judge_overview.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")
    judge_runs_cmd = judge_subs.add_parser("runs", help="Judge run history")
    judge_runs_cmd.add_argument("--project", required=True, metavar="NAME")
    judge_runs_cmd.add_argument("--last", type=int, default=None, metavar="N", help="Show only the last N runs")
    judge_runs_cmd.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")

    # -- assert ----------------------------------------------------------------

    g_assert = subs.add_parser("assert", add_help=False)
    g_assert.add_argument("-h", "--help", action="store_true")
    g_assert.add_argument("--project", required=True, metavar="NAME", help="Project name")
    g_assert.add_argument("--file", required=True, metavar="FILE", help="Python file with assertion functions")
    g_assert.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag (e.g. prod)")
    g_assert.add_argument("--last", type=int, default=None, metavar="N", help="Only check the last N traces")
    g_assert.add_argument("--store", action="store_true", help="Store results on Bandito backend")

    # -- dataset -----------------------------------------------------------------

    g_dataset = subs.add_parser("dataset", add_help=False)
    g_dataset.add_argument("-h", "--help", action="store_true")
    dataset_subs = g_dataset.add_subparsers(dest="command")

    # dataset create
    p_create = dataset_subs.add_parser("create", help="Create a new dataset")
    p_create.add_argument("name", help="Dataset name")
    p_create.add_argument("--project", required=True, metavar="NAME", help="Project name")
    p_create.add_argument("--description", default=None, metavar="TEXT", help="Dataset description")
    p_create.add_argument("--human-score-eq", type=int, choices=[0, 1], metavar="N", help="Filter by human score (0 or 1)")
    p_create.add_argument("--judge-score-eq", type=int, choices=[0, 1], metavar="N", help="Filter by judge score (0 or 1)")
    p_create.add_argument("--judge-score-lt", type=float, metavar="N", help="Filter by judge score < N")
    p_create.add_argument("--limit", type=int, default=100, metavar="N", help="Max traces to include")

    # dataset list
    p_list = dataset_subs.add_parser("list", help="List datasets for a project")
    p_list.add_argument("--project", required=True, metavar="NAME", help="Project name")

    # dataset inspect
    p_inspect = dataset_subs.add_parser("inspect", help="Inspect a dataset")
    p_inspect.add_argument("name", help="Dataset name or ID")
    p_inspect.add_argument("--project", required=True, metavar="NAME", help="Project name")
    p_inspect.add_argument("--limit", type=int, default=10, metavar="N", help="Max items to show")

    # dataset compare
    p_compare = dataset_subs.add_parser("compare", help="Compare dataset vs production")
    p_compare.add_argument("name", help="Dataset name or ID")
    p_compare.add_argument("--project", required=True, metavar="NAME", help="Project name")

    # dataset regression
    p_regression = dataset_subs.add_parser("regression", help="Run diagnostics on dataset")
    p_regression.add_argument("name", help="Dataset name or ID")
    p_regression.add_argument("--project", required=True, metavar="NAME", help="Project name")

    # dataset export
    p_export = dataset_subs.add_parser("export", help="Export dataset")
    p_export.add_argument("name", help="Dataset name or ID")
    p_export.add_argument("--project", required=True, metavar="NAME", help="Project name")
    p_export.add_argument("--format", choices=["jsonl", "pytest", "pairs"], default="jsonl", help="Export format")
    p_export.add_argument("--output", required=True, metavar="FILE", help="Output file")

    # -- diagnostics -----------------------------------------------------------

    g_diag = subs.add_parser("diagnostics", add_help=False)
    g_diag.add_argument("-h", "--help", action="store_true")
    g_diag.add_argument("--project", required=True, metavar="NAME", help="Project name")
    g_diag.add_argument("--tag", default=None, metavar="TAG", help="Filter traces by tag")
    g_diag.add_argument("--days", type=int, default=14, metavar="N", help="Analyze last N days (default: 14)")
    g_diag.add_argument("--all", action="store_true", dest="all_traces", help="Analyze all traces (ignore --days)")
    g_diag.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")

    # -- refresh ----------------------------------------------------------------

    refresh_cmd = subs.add_parser("refresh", help="Refresh cache for projects and traces")
    refresh_cmd.add_argument("--project", default=None, metavar="NAME", help="Project to refresh (default: all projects)")
    refresh_cmd.add_argument("--lite", action="store_true", help="Refresh only layers 1+2 (projects + summaries, faster)")
    refresh_cmd.add_argument("--no-cache", action="store_true", help="Force fresh data, bypass cache")

    # -- audit -----------------------------------------------------------------

    g_audit = subs.add_parser("audit", add_help=False)
    g_audit.add_argument("-h", "--help", action="store_true")
    g_audit.add_argument("--project", required=True, metavar="NAME", help="Project name")
    g_audit.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")

    # -- analyze ---------------------------------------------------------------

    g_analyze = subs.add_parser("analyze", add_help=False)
    g_analyze.add_argument("-h", "--help", action="store_true")
    analyze_subs = g_analyze.add_subparsers(dest="command")

    for name in _ANALYZE_DESCRIPTIONS:
        p = analyze_subs.add_parser(name, help=_ANALYZE_DESCRIPTIONS[name].split("\n")[0].strip())
        _add_action_args(p)
        p.add_argument("--json", action="store_true", dest="json_output", help="Output as JSON")

    # -- improve ---------------------------------------------------------------

    g_improve = subs.add_parser("improve", add_help=False)
    g_improve.add_argument("-h", "--help", action="store_true")
    improve_subs = g_improve.add_subparsers(dest="command")

    improve_subs.add_parser("dataset", help="Manage evaluation datasets")
    improve_subs.add_parser("replay", help="Test a model swap against graded data")

    # -- skill -----------------------------------------------------------------

    g_skill = subs.add_parser("skill", add_help=False)
    g_skill.add_argument("-h", "--help", action="store_true")
    skill_subs = g_skill.add_subparsers(dest="command")

    skill_subs.add_parser("install", help="Install Claude Code skill for Bandito")

    # -- Parse & Route ---------------------------------------------------------

    args = parser.parse_args()

    if args.help or args.group is None:
        print(_HELP)
        return

    group = args.group
    command = getattr(args, "command", None)

    # -- signup --

    if group == "signup":
        from .signup import run_signup
        run_signup()
        return

    # -- setup --

    if group == "setup":
        if args.help or command is None:
            print(_SETUP_HELP)
            return
        if command == "config":
            from .setup import run_setup_config
            run_setup_config()
            return

    def _resolve_project(project: str | None, use_cache: bool = True) -> str | None:
        """Resolve project name from args or list of projects."""
        from .retrieval import list_projects as rl_list_projects
        
        if project:
            return project
        
        projects = rl_list_projects(use_cache=use_cache)
        if not projects:
            print("\n  No projects found.\n")
            return None
        return projects[0]

    # -- observe --

    if group == "observe":
        if args.help or command is None:
            print(_OBSERVE_HELP)
            return

        from ..storage import get_store
        from .retrieval import get_retrieval, list_projects as rl_list_projects
        from .observe import list_projects as cli_list_projects, list_traces, show_trace

        if command == "projects":
            store = get_store(project=args.project if hasattr(args, "project") else None)
            cli_list_projects(store)
            return
        if command == "basics":
            print("\n  `observe basics` has moved to `bandito diagnostics`.")
            print(f"  Run:  bandito diagnostics --project {args.project or 'NAME'}\n")
            return
        if command == "traces":
            no_cache = getattr(args, "no_cache", False)
            project = _resolve_project(args.project, no_cache)
            if not project:
                return
            
            if no_cache:
                r = get_retrieval(project)
                r.invalidate()
            
            store = get_store(project=project)
            client = get_client()
            list_traces(
                store, project,
                tag=getattr(args, "tag", None),
                last=getattr(args, "last", 10),
                client=client,
                json_output=getattr(args, "json_output", False),
            )
            return
        if command == "trace":
            no_cache = getattr(args, "no_cache", False)
            project = _resolve_project(args.project, no_cache)
            if not project:
                return
            
            if no_cache:
                r = get_retrieval(project)
                r.invalidate()
            
            store = get_store(project=project)
            client = get_client()
            show_trace(store, args.trace_id, client=client)
            return

    # -- tui --

    if group == "tui":
        from ..storage import get_store
        from .tui import launch

        client = get_client()
        if client is None:
            return
        store = get_store(project=args.project)
        launch(screen="dashboard", store=store, client=client, project=args.project)
        return

    # -- judge --

    if group == "judge":
        if args.help or command is None:
            print(_JUDGE_HELP)
            return

        from .judge_cmd import run_calibrate, run_judge_cmd, run_overview, run_runs

        if command == "calibrate":
            run_calibrate(
                args.project,
                max_traces=getattr(args, "n", None),
                dry_run=getattr(args, "dry_run", False),
                regenerate=getattr(args, "regenerate", False),
            )
            return
        if command == "overview":
            run_overview(
                args.project,
                json_output=getattr(args, "json_output", False),
            )
            return
        if command == "runs":
            run_runs(
                args.project,
                last=getattr(args, "last", None),
                json_output=getattr(args, "json_output", False),
            )
            return
        if command == "run":
            run_judge_cmd(
                args.project,
                sample=getattr(args, "sample", None),
                rejudge=getattr(args, "rejudge", False),
                dry_run=getattr(args, "dry_run", False),
                dataset=getattr(args, "dataset", None),
            )
            return

    # -- assert --

    if group == "assert":
        if args.help:
            print("\n  bandito assert\n")
            print("  Run Python assertion functions against traces.")
            print("  Each function receives a trace dict and returns True (pass) or False (fail).\n")
            print("  Options:")
            print("    --project NAME   Project name (required)")
            print("    --file FILE      Python file with assertion functions (required)")
            print("    --last N         Only check the last N traces")
            print("    --store          Store results on Bandito backend\n")
            return
        from .assert_ import run_assert
        run_assert(
            args.project,
            args.file,
            last=getattr(args, "last", None),
            store=getattr(args, "store", False),
        )
        return

    # -- diagnostics --

    if group == "diagnostics":
        if args.help:
            print("\n  bandito diagnostics --project NAME\n")
            print("  System health check: reliability, performance, cost, and attention items.")
            print("  Deterministic — no LLM required.\n")
            print("  Options:")
            print("    --project NAME   Project name (required)")
            print("    --tag TAG        Filter traces by tag")
            print("    --json           Output as JSON\n")
            return
        from .diagnostics import run_diagnostics_cmd
        run_diagnostics_cmd(
            args.project,
            tag=getattr(args, "tag", None),
            days=getattr(args, "days", 14),
            all_traces=getattr(args, "all_traces", False),
            json_output=getattr(args, "json_output", False),
        )
        return

    # -- audit --

    if group == "audit":
        if args.help:
            print("\n  bandito audit --project NAME\n")
            print("  Data integrity checks: storage connectivity, backend connectivity,")
            print("  project name consistency, trace-evaluation linkage.\n")
            print("  Options:")
            print("    --project NAME   Project name (required)")
            print("    --json           Output as JSON\n")
            return
        from .audit import run_audit
        run_audit(
            args.project,
            json_output=getattr(args, "json_output", False),
        )
        return

    # -- analyze --

    if group == "analyze":
        if args.help or command is None:
            _analyze_help()
            return
        if command == "tradeoffs":
            from .analyze_cmd import run_tradeoffs_cmd
            run_tradeoffs_cmd(
                args.project,
                tag=getattr(args, "tag", None),
                last=getattr(args, "last", None),
                days=getattr(args, "days", 14),
                all_traces=getattr(args, "all_traces", False),
                json_output=getattr(args, "json_output", False),
            )
            return
        if command == "drift":
            from .analyze_cmd import run_drift_cmd
            run_drift_cmd(
                args.project,
                tag=getattr(args, "tag", None),
                last=getattr(args, "last", None),
                days=getattr(args, "days", 14),
                all_traces=getattr(args, "all_traces", False),
                json_output=getattr(args, "json_output", False),
            )
            return
        if command == "coverage":
            from .analyze_cmd import run_coverage_cmd
            run_coverage_cmd(
                args.project,
                tag=getattr(args, "tag", None),
                last=getattr(args, "last", None),
                days=getattr(args, "days", 14),
                all_traces=getattr(args, "all_traces", False),
                json_output=getattr(args, "json_output", False),
            )
            return
        _stub(
            f"analyze {command}",
            _ANALYZE_DESCRIPTIONS.get(command, ""),
        )
        return

    # -- refresh --

    if group == "refresh":
        if args.help:
            print("\n  bandito refresh\n")
            print("  Refresh cache for projects and traces.")
            print("  By default, refreshes all layers (projects, summaries, full traces).\n")
            print("  Options:")
            print("    --project NAME   Project to refresh (default: all)")
            print("    --lite          Refresh only layers 1+2 (projects + summaries, faster)")
            print("    --no-cache      Force fresh data, bypass cache\n")
            return

        from .retrieval import get_retrieval, list_projects

        lite = getattr(args, "lite", False)
        no_cache = getattr(args, "no_cache", False)
        project = args.project

        if project:
            r = get_retrieval(project)
            counts = r.refresh(lite=lite)
            if no_cache:
                r.invalidate()
                counts = r.refresh(lite=lite)
            print(f"\n  ✓ Refreshed project '{project}'")
            if counts.get("traces"):
                print(f"    {counts['traces']} traces")
            if counts.get("full_traces"):
                print(f"    {counts['full_traces']} full traces")
        else:
            projects = list_projects(use_cache=not no_cache)
            total_traces = 0
            total_full = 0
            for p in projects:
                r = get_retrieval(p)
                counts = r.refresh(lite=lite)
                total_traces += counts.get("traces", 0)
                total_full += counts.get("full_traces", 0)
            print(f"\n  ✓ Refreshed {len(projects)} projects")
            print(f"    {total_traces} traces")
            print(f"    {total_full} full traces")
        return

    # -- dataset --

    if group == "dataset":
        if args.help or command is None:
            print(_DATASET_HELP)
            return

        from .dataset_cmd import run_dataset_cmd
        run_dataset_cmd(group, command, args)
        return

    # -- improve --

    if group == "improve":
        if args.help or command is None:
            print(_IMPROVE_HELP)
            return
        if command == "dataset":
            _stub(
                "improve dataset",
                "Curate evaluation datasets from production traces.\n"
                "  Filter by score, segment, recency. Export for CI or replay.",
            )
            return
        if command == "replay":
            _stub(
                "improve replay",
                "Test a model swap against graded data offline.\n"
                "  Compare quality, cost, and latency before deploying.\n"
                "\n"
                "  bandito improve replay --model gpt-4o-mini --dataset my-cases.json",
            )
            return

    # -- skill --

    if group == "skill":
        if args.help or command is None:
            print(_SKILL_HELP)
            return
        if command == "install":
            from .skill_cmd import run_skill_install
            run_skill_install()
            return
