import argparse
import uuid

from sage.common.config.ports import SagePorts

_DEFAULT_GATEWAY_URL = f"http://localhost:{SagePorts.GATEWAY_DEFAULT}"

from experiments.common.cli_args import (
    add_common_benchmark_args,
    build_run_config,
    validate_benchmark_args,
)
from experiments.q2_controlmix import ControlPlaneExperiment


class BurstTownExperiment(ControlPlaneExperiment):
    """
    Q6 BurstTown workload.

    Transaction-like burst and priority pressure over mixed request classes.
    """

    pass


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Q6 – BurstTown: benchmark bursty mixed-priority transactional "
            "workloads against different backends."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    add_common_benchmark_args(parser)
    parser.add_argument("--rate", type=float, default=10.0, metavar="REQ_S")
    parser.add_argument("--total-requests", type=int, default=100, metavar="N")
    parser.add_argument("--llm-ratio", type=float, default=0.7, metavar="RATIO")
    parser.add_argument("--gateway", type=str, default=_DEFAULT_GATEWAY_URL, metavar="URL")
    return parser


def main() -> None:  # pragma: no cover
    import sys
    from pathlib import Path

    from config.config_loader import ConfigLoader
    from experiments.config import WorkloadConfig

    parser = _build_parser()
    args = parser.parse_args()
    validate_benchmark_args(args)
    run_cfg = build_run_config(
        args,
        workload="Q6",
        rate=args.rate,
        total_requests=args.total_requests,
        llm_ratio=args.llm_ratio,
        gateway=args.gateway,
    )

    if args.verbose:
        import json

        print("Run config:")
        print(json.dumps(run_cfg, indent=2))

    if args.dry_run:
        print("[dry-run] Q6 BurstTown – configuration validated. Exiting without running.")
        sys.exit(0)

    loader = ConfigLoader()
    config = loader.get_default_config("Q6")
    config.experiment_section = "Q6"
    if args.quick:
        config = loader.apply_quick_mode(config)

    config.workload = WorkloadConfig(
        total_requests=args.total_requests,
        request_rate=args.rate,
        seed=args.seed,
        llm_ratio=args.llm_ratio,
        warmup_requests=config.workload.warmup_requests,
        input_tokens_min=config.workload.input_tokens_min,
        input_tokens_max=config.workload.input_tokens_max,
        output_tokens_min=config.workload.output_tokens_min,
        output_tokens_max=config.workload.output_tokens_max,
        arrival_pattern=config.workload.arrival_pattern,
    )
    config.gateway_url = args.gateway

    output_dir = Path(args.output_dir) / "q6"
    for rep in range(1, args.repeat + 1):
        rep_label = f" (repetition {rep}/{args.repeat})" if args.repeat > 1 else ""
        print(f"Running Q6 BurstTown{rep_label} …")
        rep_output = output_dir / (f"rep{rep}" if args.repeat > 1 else "")
        experiment = BurstTownExperiment(config=config, output_dir=rep_output, verbose=args.verbose)
        experiment.backend = args.backend
        experiment.nodes = int(args.nodes)
        experiment.parallelism = int(args.parallelism)
        experiment.run_id = f"q6-{args.backend}-{rep}-{uuid.uuid4().hex[:8]}"
        experiment.setup()
        experiment.run()
        experiment.teardown()
        print(f"  Results saved to {rep_output.absolute()}")


if __name__ == "__main__":  # pragma: no cover
    main()
