#!/usr/bin/env python3
"""CLI for threshold calibration.

This script provides a command-line interface for running threshold
calibration experiments and generating reports.

Usage:
    # Basic calibration with default settings
    python calibrate_thresholds.py --data calibration_data.csv

    # Optimize for quality with 100 iterations
    python calibrate_thresholds.py --data data.csv --target quality --iterations 100

    # Use specific search strategy
    python calibrate_thresholds.py --data data.csv --strategy bayesian

    # Save report and config
    python calibrate_thresholds.py --data data.csv --output report.json --config thresholds.yaml

Example data format (CSV):
    query_id,task_id,mean_logit,max_prob,logit_margin,logit_std,entropy,top1_model,top1_prob,effect_modelA,effect_modelB,...
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

# Add parent directory to path for local development
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from embergraphrouter.calibration import (
    CalibrationConfig,
    CalibrationReport,
    OfflineDataset,
    ThresholdCalibrator,
)
from embergraphrouter.calibration.config import (
    BALANCED_CONFIG,
    COST_FOCUSED_CONFIG,
    LATENCY_FOCUSED_CONFIG,
    OptimizationTarget,
    PARETO_CONFIG,
    QUALITY_FOCUSED_CONFIG,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Calibrate tiered routing thresholds from offline data.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    # Data input
    parser.add_argument(
        "--data",
        type=str,
        required=True,
        help="Path to calibration data (CSV or Parquet)",
    )
    parser.add_argument(
        "--effect-prefix",
        type=str,
        default="effect_",
        help="Prefix for effect columns (default: effect_)",
    )

    # Optimization settings
    parser.add_argument(
        "--target",
        type=str,
        choices=["quality", "cost", "latency", "balanced", "pareto"],
        default="balanced",
        help="Optimization target (default: balanced)",
    )
    parser.add_argument(
        "--quality-weight",
        type=float,
        default=0.5,
        help="Weight for quality in balanced optimization (default: 0.5)",
    )
    parser.add_argument(
        "--cost-weight",
        type=float,
        default=0.3,
        help="Weight for cost in balanced optimization (default: 0.3)",
    )
    parser.add_argument(
        "--latency-weight",
        type=float,
        default=0.2,
        help="Weight for latency in balanced optimization (default: 0.2)",
    )

    # Search settings
    parser.add_argument(
        "--strategy",
        type=str,
        choices=["grid", "random", "bayesian"],
        default="bayesian",
        help="Search strategy (default: bayesian)",
    )
    parser.add_argument(
        "--iterations",
        type=int,
        default=100,
        help="Number of search iterations (default: 100)",
    )
    parser.add_argument(
        "--cv-folds",
        type=int,
        default=5,
        help="Number of cross-validation folds (default: 5, 0 to disable)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility (default: 42)",
    )

    # Output settings
    parser.add_argument(
        "--output",
        type=str,
        help="Path to save report (JSON or TXT)",
    )
    parser.add_argument(
        "--config",
        type=str,
        help="Path to save threshold configuration (YAML)",
    )
    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="Enable verbose logging",
    )

    return parser.parse_args()


def load_data(args: argparse.Namespace) -> OfflineDataset:
    """Load calibration data."""
    data_path = Path(args.data)

    if not data_path.exists():
        logger.error("Data file not found: %s", data_path)
        sys.exit(1)

    logger.info("Loading data from %s", data_path)

    if data_path.suffix == ".parquet":
        dataset = OfflineDataset.from_parquet(
            data_path,
            effect_prefix=args.effect_prefix,
        )
    else:
        dataset = OfflineDataset.from_csv(
            data_path,
            effect_prefix=args.effect_prefix,
        )

    summary = dataset.summary()
    logger.info(
        "Loaded %d samples, %d models, %d tasks",
        summary["n_samples"],
        summary["n_models"],
        summary["n_tasks"],
    )
    logger.info(
        "Difficulty distribution: %.1f%% trivial, %.1f%% hard, %.1f%% medium",
        summary["difficulty_distribution"]["trivial_pct"],
        summary["difficulty_distribution"]["hard_pct"],
        summary["difficulty_distribution"]["medium_pct"],
    )

    return dataset


def create_config(args: argparse.Namespace) -> CalibrationConfig:
    """Create calibration configuration from arguments."""
    # Map target to OptimizationTarget
    target_map = {
        "quality": OptimizationTarget.QUALITY,
        "cost": OptimizationTarget.COST,
        "latency": OptimizationTarget.LATENCY,
        "balanced": OptimizationTarget.BALANCED,
        "pareto": OptimizationTarget.PARETO,
    }

    return CalibrationConfig(
        optimization_target=target_map[args.target],
        quality_weight=args.quality_weight,
        cost_weight=args.cost_weight,
        latency_weight=args.latency_weight,
        n_search_iterations=args.iterations,
        n_cv_folds=args.cv_folds,
        random_seed=args.seed,
    )


def main() -> int:
    """Main entry point."""
    args = parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Load data
    dataset = load_data(args)

    # Create config
    config = create_config(args)
    logger.info(
        "Configuration: target=%s, iterations=%d, strategy=%s, cv_folds=%d",
        config.optimization_target.value,
        config.n_search_iterations,
        args.strategy,
        config.n_cv_folds,
    )

    # Run calibration
    logger.info("Starting calibration...")
    calibrator = ThresholdCalibrator(config, search_strategy=args.strategy)
    result = calibrator.calibrate(dataset)

    # Generate report
    report = CalibrationReport.from_result(result)

    # Print report
    print("\n")
    print(report.to_text())

    # Save outputs
    if args.output:
        output_path = Path(args.output)
        report.save(output_path)
        logger.info("Report saved to %s", output_path)

    if args.config:
        config_path = Path(args.config)
        with open(config_path, "w") as f:
            f.write(report.to_yaml_config())
        logger.info("Configuration saved to %s", config_path)

    # Print optimal thresholds for easy copy-paste
    print("\nOptimal thresholds (Python format):")
    print("TieredThresholds(")
    for key, value in report.optimal_thresholds.items():
        print(f"    {key}={value},")
    print(")")

    return 0


if __name__ == "__main__":
    sys.exit(main())
