"""Simple CLI for benchmark management.

Usage:
    pentestgpt-benchmark list [--tags TAG ...] [--levels N ...]
    pentestgpt-benchmark start BENCHMARK_ID
    pentestgpt-benchmark stop BENCHMARK_ID
    pentestgpt-benchmark status
    pentestgpt-benchmark run [--range 1-10 | --ids 1,5,10 | --all] [options]
    pentestgpt-benchmark analyze [RUN_DIR]
"""

import argparse
import asyncio
import json
import sys
from pathlib import Path

from pentestgpt.benchmark.analyzer import analyze_run
from pentestgpt.benchmark.config import DEFAULT_BENCHMARKS_DIR, DEFAULT_OUTPUT_DIR, MODEL_NAMES
from pentestgpt.benchmark.docker import (
    get_running_benchmarks,
    start_benchmark,
    stop_benchmark,
)
from pentestgpt.benchmark.models import BenchmarkConfig
from pentestgpt.benchmark.registry import BenchmarkRegistry
from pentestgpt.benchmark.runner import BenchmarkRunner


def cmd_list(args: argparse.Namespace) -> int:
    """List available benchmarks."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    try:
        registry = BenchmarkRegistry(benchmarks_dir)
        registry.load()
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return 1

    # Show tags if requested
    if args.show_tags:
        print("Available tags:")
        for tag in sorted(registry.get_all_tags()):
            count = len(registry.filter(tags=[tag]))
            print(f"  {tag}: {count} benchmarks")
        return 0

    # Filter and list
    benchmarks = registry.filter(tags=args.tags, levels=args.levels)

    print(f"{'ID':<15} {'Level':^5} {'Name':<40} {'Tags'}")
    print("-" * 80)
    for b in benchmarks:
        tags_str = ", ".join(b.tags[:3])
        if len(b.tags) > 3:
            tags_str += "..."
        print(f"{b.id:<15} {b.level:^5} {b.name[:40]:<40} {tags_str}")

    print(f"\nTotal: {len(benchmarks)} benchmarks")
    return 0


def cmd_start(args: argparse.Namespace) -> int:
    """Start a benchmark."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    registry = BenchmarkRegistry(benchmarks_dir)
    benchmark = registry.get(args.benchmark_id)

    if not benchmark:
        print(f"Error: Benchmark '{args.benchmark_id}' not found")
        return 1

    print(f"Starting benchmark: {benchmark.name}")
    print(f"Description: {benchmark.description}")
    print(f"Level: {benchmark.level}, Tags: {', '.join(benchmark.tags)}")
    print()

    result = start_benchmark(benchmark.path)

    if result["success"]:
        target_url = result["target_url"]
        # Extract port from URL for Docker target
        port = target_url.split(":")[-1]
        docker_url = f"http://host.docker.internal:{port}"

        print("\nBenchmark started successfully!")
        print(f"Target URL: {target_url}")
        print()
        print("Run PentestGPT against this target:")
        print(f"  Local:  pentestgpt --target {target_url}")
        print(f"  Docker: pentestgpt --target {docker_url}")
    else:
        print(f"\nFailed to start benchmark: {result['message']}")
        return 1

    return 0


def cmd_stop(args: argparse.Namespace) -> int:
    """Stop a benchmark."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    registry = BenchmarkRegistry(benchmarks_dir)
    benchmark = registry.get(args.benchmark_id)

    if not benchmark:
        print(f"Error: Benchmark '{args.benchmark_id}' not found")
        return 1

    result = stop_benchmark(benchmark.path)

    if result["success"]:
        print("Benchmark stopped successfully")
    else:
        print(f"Failed to stop benchmark: {result['message']}")
        return 1

    return 0


def cmd_status(args: argparse.Namespace) -> int:
    """Show running benchmarks."""
    running = get_running_benchmarks()

    if not running:
        print("No benchmark containers currently running")
        return 0

    print("Running benchmark containers:")
    print(f"{'Name':<40} {'Ports':<30} {'Status'}")
    print("-" * 80)
    for container in running:
        print(f"{container['name']:<40} {container['ports']:<30} {container['status']}")

    return 0


# ============================================================================
# Helper functions for run command
# ============================================================================

def parse_range(range_str: str) -> list[int]:
    """Parse range string to list of IDs."""
    parts = range_str.replace("-", " ").split()
    if len(parts) == 2:
        try:
            start, end = int(parts[0]), int(parts[1])
            if start < 1 or end > 104 or start > end:
                raise ValueError("Range must be between 1-104 and start <= end")
            return list(range(start, end + 1))
        except ValueError as e:
            raise ValueError(f"Invalid range: {e}") from e
    raise ValueError(f"Invalid range format: {range_str}. Use '1-10' or '1 10'")


def parse_ids(ids_str: str) -> list[int]:
    """Parse comma-separated IDs."""
    try:
        ids = [int(id.strip()) for id in ids_str.split(",")]
        for id in ids:
            if id < 1 or id > 104:
                raise ValueError(f"ID {id} out of range (must be 1-104)")
        return ids
    except ValueError as e:
        raise ValueError(f"Invalid IDs: {e}") from e


def find_last_run(output_dir: Path) -> Path:
    """Find the most recent benchmark run directory."""
    if not output_dir.exists():
        raise FileNotFoundError(f"Output directory not found: {output_dir}")

    run_dirs = sorted(output_dir.glob("benchmark_run_*"))
    if not run_dirs:
        raise FileNotFoundError(f"No previous benchmark runs found in {output_dir}")

    return run_dirs[-1]


def load_failed_benchmarks(run_dir: Path) -> list[dict]:
    """Load failed benchmark results from a run directory."""
    summary_file = run_dir / "summary.json"

    if not summary_file.exists():
        raise FileNotFoundError(
            f"summary.json not found in {run_dir}. Not a valid run directory?"
        )

    with open(summary_file) as f:
        data = json.load(f)

    failed = [r for r in data["results"] if not r["success"]]
    return failed


def extract_failed_ids(failed_results: list[dict]) -> list[int]:
    """Extract benchmark IDs from failed results."""
    ids = []
    for result in failed_results:
        bench_id = result["benchmark_id"]
        numeric_id = int(bench_id.split("-")[1])
        ids.append(numeric_id)
    return sorted(ids)


# ============================================================================
# New commands: run and analyze
# ============================================================================

def cmd_run(args: argparse.Namespace) -> int:
    """Run automated benchmarks."""
    benchmarks_dir = Path(args.benchmarks_dir) if args.benchmarks_dir else DEFAULT_BENCHMARKS_DIR

    # Parse benchmark selection
    try:
        if args.range:
            benchmark_ids = parse_range(args.range)
        elif args.ids:
            benchmark_ids = parse_ids(args.ids)
        elif args.all:
            benchmark_ids = list(range(1, 105))
        elif args.retry_failed is not None:
            # Determine run directory
            if args.retry_failed == '':
                run_dir = find_last_run(args.output_dir)
                print(f"Using last run: {run_dir.name}")
            else:
                run_dir = Path(args.retry_failed).resolve()
                if not run_dir.exists():
                    print(f"Error: Run directory not found: {run_dir}")
                    return 1
                print(f"Using specified run: {run_dir}")

            # Load failed benchmarks
            failed_results = load_failed_benchmarks(run_dir)

            if not failed_results:
                print("No failed benchmarks to retry!")
                return 0

            benchmark_ids = extract_failed_ids(failed_results)
            print(f"Found {len(benchmark_ids)} failed benchmarks to retry")
        else:
            print("Error: Must specify --range, --ids, --all, or --retry-failed")
            return 1

    except (ValueError, FileNotFoundError) as e:
        print(f"Error: {e}")
        return 1

    # Map model name
    model = MODEL_NAMES.get(args.model.lower(), args.model) if args.model else None

    # Create config
    config = BenchmarkConfig(
        benchmark_ids=benchmark_ids,
        timeout_seconds=args.timeout,
        benchmarks_dir=benchmarks_dir,
        resume=args.resume,
        output_dir=args.output_dir,
        any_flag=args.any_flag,
        pattern_flag=args.pattern_flag,
        model=model
    )

    # Run benchmarks
    print("\nStarting automated benchmark run")
    print(f"Benchmarks: {len(benchmark_ids)}")
    print(f"Timeout: {args.timeout}s per benchmark")
    if model:
        print(f"Model: {model}")
    print()

    runner = BenchmarkRunner(config)
    asyncio.run(runner.run_all())

    return 0


def cmd_analyze(args: argparse.Namespace) -> int:
    """Analyze benchmark results."""
    run_dir = Path(args.run_dir) if args.run_dir else None

    try:
        analyze_run(run_dir)
        return 0
    except (FileNotFoundError, ValueError, json.JSONDecodeError) as e:
        print(f"Error: {e}")
        return 1


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        prog="pentestgpt-benchmark",
        description="Manage benchmark containers for PentestGPT testing",
    )
    parser.add_argument(
        "--benchmarks-dir",
        "-d",
        help="Path to benchmarks directory",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    # List command
    list_parser = subparsers.add_parser("list", help="List available benchmarks")
    list_parser.add_argument("--tags", "-t", nargs="+", help="Filter by tags")
    list_parser.add_argument(
        "--levels", "-l", nargs="+", type=int, help="Filter by levels (1, 2, 3)"
    )
    list_parser.add_argument("--show-tags", action="store_true", help="Show all tags")

    # Start command
    start_parser = subparsers.add_parser("start", help="Start a benchmark")
    start_parser.add_argument("benchmark_id", help="Benchmark ID (e.g., XBEN-001-24)")

    # Stop command
    stop_parser = subparsers.add_parser("stop", help="Stop a benchmark")
    stop_parser.add_argument("benchmark_id", help="Benchmark ID")

    # Status command
    subparsers.add_parser("status", help="Show running benchmarks")

    # Run command - automated execution
    run_parser = subparsers.add_parser("run", help="Run automated benchmarks")

    # Benchmark selection (mutually exclusive)
    selection = run_parser.add_mutually_exclusive_group(required=True)
    selection.add_argument("--range", type=str, help="Benchmark range (e.g., '1-10')")
    selection.add_argument("--ids", type=str, help="Comma-separated IDs (e.g., '1,5,10')")
    selection.add_argument("--all", action="store_true", help="Run all 104 benchmarks")
    selection.add_argument(
        "--retry-failed",
        nargs='?',
        const='',
        metavar='PATH',
        help="Retry failed benchmarks from last run (or specify run directory)"
    )

    # Execution options
    run_parser.add_argument(
        "--timeout",
        type=int,
        default=900,
        help="Timeout per benchmark in seconds (default: 900)"
    )
    run_parser.add_argument(
        "--model",
        type=str,
        choices=["opus", "sonnet", "haiku"],
        help="Claude model to use (default: sonnet)"
    )
    run_parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume from previous run (skip completed)"
    )
    run_parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Directory for logs and results"
    )

    # Flag validation mode (mutually exclusive)
    flag_mode = run_parser.add_mutually_exclusive_group()
    flag_mode.add_argument(
        "--any-flag",
        action="store_true",
        help="Accept ANY flag as success"
    )
    flag_mode.add_argument(
        "--pattern-flag",
        action="store_true",
        help="Only accept flags matching strict pattern: FLAG{32+chars}"
    )

    # Analyze command
    analyze_parser = subparsers.add_parser("analyze", help="Analyze benchmark results")
    analyze_parser.add_argument(
        "run_dir",
        nargs='?',
        help="Path to run directory (default: most recent)"
    )

    args = parser.parse_args()

    if args.command == "list":
        sys.exit(cmd_list(args))
    elif args.command == "start":
        sys.exit(cmd_start(args))
    elif args.command == "stop":
        sys.exit(cmd_stop(args))
    elif args.command == "status":
        sys.exit(cmd_status(args))
    elif args.command == "run":
        sys.exit(cmd_run(args))
    elif args.command == "analyze":
        sys.exit(cmd_analyze(args))


if __name__ == "__main__":
    main()