__main__.py

"""
CLI entry point for LLM API Relay Authentication Benchmark.

Usage:
    # Official mode: run against official API to record baseline
    python -m llm_benchmark official --api-type openai --api-key $KEY --model gpt-4o -o baseline.json

    # Test mode: run against relay service and compare
    python -m llm_benchmark test --api-type openai --api-key $KEY --model gpt-4o \\
        --base-url https://relay.test/v1 --baseline baseline.json -o results.json

    # With config file:
    python -m llm_benchmark --config config.json official
"""

import argparse
import json
import os
import sys
from typing import Optional

from llm_benchmark.runner import BenchmarkRunner
from llm_benchmark.reporter import BenchmarkReporter


def load_config(path: str) -> dict:
    """Load configuration from a JSON file."""
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="LLM API Relay Authentication Benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Official baseline recording
  python -m llm_benchmark official --api-type openai --api-key sk-xxx --model gpt-4o -o baseline.json

  # Test a relay service
  python -m llm_benchmark test --api-type openai --api-key sk-yyy --model gpt-4o \\
      --base-url https://relay.example.com/v1 --baseline baseline.json -o results.json

  # Using config file
  python -m llm_benchmark --config config.json official
        """,
    )

    parser.add_argument(
        "--config",
        type=str,
        help="Path to JSON config file. When provided, config values are used as defaults.",
    )

    # Shared arguments
    parser.add_argument(
        "--api-type",
        type=str,
        choices=["openai", "anthropic"],
        default=None,
        help='API type: "openai" (OpenAI-compatible) or "anthropic"',
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=None,
        help="API key. Can also be set via OPENAI_API_KEY or ANTHROPIC_API_KEY env vars.",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="Model name (e.g., gpt-4o, claude-3-5-sonnet-20241022)",
    )
    parser.add_argument(
        "--base-url",
        type=str,
        default=None,
        help="Base URL for the API endpoint (required for relay services)",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Temperature setting (default: 0.0)",
    )
    parser.add_argument(
        "--top-p",
        type=float,
        default=1.0,
        help="Top-p setting (default: 1.0)",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=4096,
        help="Max tokens for generation (default: 4096)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=120.0,
        help="API timeout in seconds (default: 120)",
    )
    parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output file path for results",
    )

    # Subcommands
    subparsers = parser.add_subparsers(dest="mode", help="Operation mode")

    # Official mode
    official_parser = subparsers.add_parser(
        "official",
        help="Run in official mode to record baseline results",
    )
    official_parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output file for baseline results",
    )

    # Test mode
    test_parser = subparsers.add_parser(
        "test",
        help="Run in test mode against a relay service",
    )
    test_parser.add_argument(
        "--baseline", "-b",
        type=str,
        required=True,
        help="Path to baseline JSON from official mode",
    )
    test_parser.add_argument(
        "--output", "-o",
        type=str,
        default=None,
        help="Output file for test results",
    )

    return parser


def resolve_args(args, config: Optional[dict] = None):
    """Resolve arguments, using config file values as defaults, then CLI args as overrides."""
    env = os.environ

    resolved = {}

    if config:
        resolved.update(config)

    # Override with CLI args (only if explicitly set)
    cli_args = vars(args)

    # Map CLI arg names to expected keys
    key_map = {
        "api_type": "api_type",
        "api_key": "api_key",
        "model": "model",
        "base_url": "base_url",
        "temperature": "temperature",
        "top_p": "top_p",
        "max_tokens": "max_tokens",
        "timeout": "timeout",
    }

    for cli_key, config_key in key_map.items():
        val = cli_args.get(cli_key)
        if val is not None:
            resolved[config_key] = val

    # Resolve API key from env if not set
    if "api_key" not in resolved or not resolved["api_key"]:
        api_type = resolved.get("api_type", "openai")
        if api_type == "openai":
            resolved["api_key"] = env.get("OPENAI_API_KEY", "")
        elif api_type == "anthropic":
            resolved["api_key"] = env.get("ANTHROPIC_API_KEY", "")

    return resolved


def main():
    parser = build_parser()
    cli_args = parser.parse_args()

    if not cli_args.mode:
        parser.print_help()
        sys.exit(1)

    # Load config if provided
    config = None
    if cli_args.config:
        if not os.path.exists(cli_args.config):
            print(f"Error: Config file not found: {cli_args.config}")
            sys.exit(1)
        config = load_config(cli_args.config)
        print(f"  Loaded config from: {cli_args.config}")

    resolved = resolve_args(cli_args, config)

    # Validate required fields
    api_type = resolved.get("api_type", "")
    if not api_type:
        print("Error: --api-type is required (openai or anthropic)")
        sys.exit(1)

    api_key = resolved.get("api_key", "")
    if not api_key:
        print(
            "Error: --api-key is required. Set it directly or via "
            "OPENAI_API_KEY / ANTHROPIC_API_KEY environment variable."
        )
        sys.exit(1)

    model = resolved.get("model", "")
    if not model:
        print("Error: --model is required")
        sys.exit(1)

    base_url = resolved.get("base_url")
    temperature = resolved.get("temperature", 0.0)
    top_p = resolved.get("top_p", 1.0)
    max_tokens = resolved.get("max_tokens", 4096)
    timeout = resolved.get("timeout", 120.0)

    # Determine output file
    output_file = cli_args.output
    if not output_file:
        output_file = resolved.get("output")
    if not output_file:
        suffix = "baseline" if cli_args.mode == "official" else "results"
        output_file = f"llm_benchmark_{suffix}_{model.replace('/', '_')}_{api_type}.json"

    # Create runner
    runner = BenchmarkRunner(
        api_type=api_type,
        api_key=api_key,
        model=model,
        base_url=base_url,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
        timeout=timeout,
    )

    try:
        if cli_args.mode == "official":
            # Official mode: record baseline
            run = runner.run_all()
            BenchmarkReporter.print_summary(run)
            BenchmarkReporter.save_json(run.to_dict(), output_file)

        elif cli_args.mode == "test":
            # Test mode: compare with baseline
            baseline_file = cli_args.baseline
            if not os.path.exists(baseline_file):
                print(f"Error: Baseline file not found: {baseline_file}")
                sys.exit(1)

            run = runner.run_all(baseline_file=baseline_file)
            BenchmarkReporter.print_summary(run)

            # Load official run for comparison
            with open(baseline_file, "r", encoding="utf-8") as f:
                official_data = json.load(f)

            # Reconstruct official run
            from llm_benchmark.runner import BenchmarkRun, TestResult

            official_results = []
            for r_data in official_data.get("results", []):
                official_results.append(TestResult(**r_data))

            official_run = BenchmarkRun(
                mode="official",
                api_type=official_data.get("api_type", api_type),
                model=official_data.get("model", model),
                base_url=official_data.get("base_url"),
                timestamp=official_data.get("timestamp", ""),
                results=official_results,
                temperature=official_data.get("temperature", temperature),
            )
            official_run.calculate_summary()

            # Generate comparison
            comparison = BenchmarkReporter.compare_runs(official_run, run)
            BenchmarkReporter.print_comparison(comparison)

            # Save results
            output_data = {
                "test_run": run.to_dict(),
                "comparison": comparison,
            }
            BenchmarkReporter.save_json(output_data, output_file)
    finally:
        runner.close()


if __name__ == "__main__":
    main()