from __future__ import annotations

import asyncio
import json
import signal
from dataclasses import asdict, dataclass, replace
from datetime import UTC, datetime
from pathlib import Path

from agentic_pd_hybrid.replay import ReplayConfig, replay_trace
from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions
from agentic_pd_hybrid.stack import ManagedPdStack, launch_pd_stack
from agentic_pd_hybrid.topology import SingleNodeTopology


@dataclass(frozen=True)
class BenchmarkConfig:
    trace_path: Path
    output_root: Path
    topology: SingleNodeTopology
    policy_name: str
    mechanism_name: str = "pd-disaggregation"
    target_duration_s: float = 600.0
    start_time_s: float = 0.0
    session_sample_rate: float = 0.01
    min_turns: int = 1
    time_scale: float = 1.0
    concurrency_limit: int = 32
    timeout_s: float = 1200.0
    request_timeout_s: float | None = None
    stream: bool = True
    stream_idle_timeout_s: float | None = 900.0
    kvcache_direct_max_uncached_tokens: int = 2048
    kvcache_admission_mode: str = "router"
    kvcache_seed_max_resident_tokens: int | None = None
    kvcache_seed_max_output_tokens: int | None = None
    kvcache_seed_min_turn_id: int = 1
    kvcache_seed_only_multiturn_sessions: bool = False
    kvcache_prefill_backup_policy: str = "release-after-transfer"
    kvcache_seed_max_inflight_decode: int | None = 3
    kvcache_seed_max_decode_transfer_queue_reqs: int | None = None
    kvcache_direct_max_decode_transfer_queue_reqs: int | None = None
    kvcache_prefill_priority_eviction: bool = False
    kvcache_prefill_direct_priority: int = -100
    kvcache_prefill_normal_priority: int = 100
    pool_poll_interval_s: float = 0.0
    pool_poll_include_sessions: bool = True
    enable_backpressure: bool = False
    backpressure_max_pause_s: float = 2.0
    kvcache_migration_reject_threshold: int = 3
    sample_profile: str = "default"
    min_initial_input_tokens: int | None = None
    max_initial_input_tokens: int | None = None
    max_append_input_tokens: int | None = None
    max_output_tokens: int | None = None
    min_overlap_ratio: float | None = None
    launch_stack: bool = True


@dataclass(frozen=True)
class BenchmarkArtifacts:
    run_dir: Path
    sampled_trace_path: Path
    metrics_path: Path
    summary_path: Path
    benchmark_config_path: Path


def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
    run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    run_label = f"{config.mechanism_name}-{config.policy_name}"
    if config.mechanism_name == "kvcache-centric":
        run_label = f"{run_label}-{config.kvcache_admission_mode}-admission"
    run_dir = config.output_root / f"{run_label}-{run_id}"
    run_dir.mkdir(parents=True, exist_ok=True)

    topology = config.topology
    if config.mechanism_name == "kvcache-centric":
        prefill_extra_server_args = topology.prefill_extra_server_args + (
            "--enable-streaming-session",
        )
        if config.kvcache_prefill_priority_eviction:
            prefill_extra_server_args = prefill_extra_server_args + (
                "--radix-eviction-policy",
                "priority",
            )
        topology = replace(
            topology,
            prefill_extra_server_args=prefill_extra_server_args,
            decode_extra_server_args=topology.decode_extra_server_args
            + (
                "--enable-streaming-session",
                "--disaggregation-decode-allow-local-prefill",
            ),
        )

    sampled_trace_path = run_dir / "sampled-trace.jsonl"
    sample_summary = sample_trace_sessions(
        SessionSampleConfig(
            trace_path=config.trace_path,
            output_path=sampled_trace_path,
            target_duration_s=config.target_duration_s,
            start_time_s=config.start_time_s,
            session_sample_rate=config.session_sample_rate,
            min_turns=config.min_turns,
            profile=config.sample_profile,  # type: ignore[arg-type]
            min_initial_input_tokens=config.min_initial_input_tokens,
            max_initial_input_tokens=config.max_initial_input_tokens,
            max_append_input_tokens=config.max_append_input_tokens,
            max_output_tokens=config.max_output_tokens,
            min_overlap_ratio=config.min_overlap_ratio,
        )
    )

    stack: ManagedPdStack | None = None
    previous_sigint = signal.getsignal(signal.SIGINT)
    previous_sigterm = signal.getsignal(signal.SIGTERM)

    def _handle_termination(signum, _frame) -> None:
        if stack is not None:
            stack.stop()
        raise SystemExit(128 + signum)

    try:
        signal.signal(signal.SIGINT, _handle_termination)
        signal.signal(signal.SIGTERM, _handle_termination)
        _mechanisms_with_router = {"pd-disaggregation", "kvcache-centric", "pd-colo"}
        _naive_dp = config.mechanism_name == "pd-colo"
        if config.launch_stack:
            stack = launch_pd_stack(
                topology=topology,
                run_dir=run_dir,
                prefill_policy="round_robin",
                decode_policy=_decode_policy_for(config.policy_name),
                timeout_s=config.timeout_s,
                router_request_timeout_s=(
                    config.request_timeout_s
                    if config.request_timeout_s is not None
                    else config.timeout_s
                ),
                include_router=(
                    config.mechanism_name in _mechanisms_with_router
                ),
                naive_dp=_naive_dp,
            )
            router_url = (
                stack.router_url
                if config.mechanism_name in _mechanisms_with_router
                else None
            )
        else:
            router_url = (
                topology.router_url
                if config.mechanism_name in _mechanisms_with_router
                else None
            )

        metrics_path = run_dir / "request-metrics.jsonl"
        replay_config = ReplayConfig(
            trace_path=sampled_trace_path,
            output_path=metrics_path,
            policy_name=config.policy_name,
            mechanism_name=config.mechanism_name,
            topology=topology,
            router_url=router_url,
            model_name=topology.model_name,
            pace=True,
            time_scale=config.time_scale,
            request_limit=None,
            concurrency_limit=config.concurrency_limit,
            header_mode=_header_mode_for(config.policy_name),
            timeout_s=config.timeout_s,
            stream=config.stream,
            stream_idle_timeout_s=config.stream_idle_timeout_s,
            kvcache_direct_max_uncached_tokens=config.kvcache_direct_max_uncached_tokens,
            kvcache_admission_mode=config.kvcache_admission_mode,  # type: ignore[arg-type]
            kvcache_seed_max_resident_tokens=config.kvcache_seed_max_resident_tokens,
            kvcache_seed_max_output_tokens=config.kvcache_seed_max_output_tokens,
            kvcache_seed_min_turn_id=config.kvcache_seed_min_turn_id,
            kvcache_seed_only_multiturn_sessions=(
                config.kvcache_seed_only_multiturn_sessions
            ),
            kvcache_prefill_backup_policy=config.kvcache_prefill_backup_policy,  # type: ignore[arg-type]
            kvcache_seed_max_inflight_decode=(
                config.kvcache_seed_max_inflight_decode
            ),
            kvcache_seed_max_decode_transfer_queue_reqs=(
                config.kvcache_seed_max_decode_transfer_queue_reqs
            ),
            kvcache_direct_max_decode_transfer_queue_reqs=(
                config.kvcache_direct_max_decode_transfer_queue_reqs
            ),
            kvcache_prefill_priority_eviction=(
                config.kvcache_prefill_priority_eviction
            ),
            kvcache_prefill_direct_priority=config.kvcache_prefill_direct_priority,
            kvcache_prefill_normal_priority=config.kvcache_prefill_normal_priority,
            pool_poll_interval_s=config.pool_poll_interval_s,
            pool_poll_include_sessions=config.pool_poll_include_sessions,
            enable_backpressure=config.enable_backpressure,
            backpressure_max_pause_s=config.backpressure_max_pause_s,
            kvcache_migration_reject_threshold=config.kvcache_migration_reject_threshold,
        )
        if config.request_timeout_s is not None:
            replay_config = replace(
                replay_config,
                timeout_s=config.request_timeout_s,
            )
        asyncio.run(replay_trace(replay_config))
    finally:
        signal.signal(signal.SIGINT, previous_sigint)
        signal.signal(signal.SIGTERM, previous_sigterm)
        if stack is not None:
            stack.stop()

    benchmark_config_path = run_dir / "benchmark-config.json"
    with benchmark_config_path.open("w", encoding="utf-8") as handle:
        json.dump(
            {
                "policy_name": config.policy_name,
                "mechanism_name": config.mechanism_name,
                "target_duration_s": config.target_duration_s,
                "start_time_s": config.start_time_s,
                "session_sample_rate": config.session_sample_rate,
                "min_turns": config.min_turns,
                "time_scale": config.time_scale,
                "concurrency_limit": config.concurrency_limit,
                "timeout_s": config.timeout_s,
                "request_timeout_s": config.request_timeout_s,
                "stream": config.stream,
                "stream_idle_timeout_s": config.stream_idle_timeout_s,
                "kvcache_direct_max_uncached_tokens": config.kvcache_direct_max_uncached_tokens,
                "kvcache_admission_mode": config.kvcache_admission_mode,
                "kvcache_seed_max_resident_tokens": config.kvcache_seed_max_resident_tokens,
                "kvcache_seed_max_output_tokens": config.kvcache_seed_max_output_tokens,
                "kvcache_seed_min_turn_id": config.kvcache_seed_min_turn_id,
                "kvcache_seed_only_multiturn_sessions": (
                    config.kvcache_seed_only_multiturn_sessions
                ),
                "kvcache_prefill_backup_policy": config.kvcache_prefill_backup_policy,
                "kvcache_seed_max_inflight_decode": (
                    config.kvcache_seed_max_inflight_decode
                ),
                "kvcache_seed_max_decode_transfer_queue_reqs": (
                    config.kvcache_seed_max_decode_transfer_queue_reqs
                ),
                "kvcache_direct_max_decode_transfer_queue_reqs": (
                    config.kvcache_direct_max_decode_transfer_queue_reqs
                ),
                "kvcache_prefill_priority_eviction": (
                    config.kvcache_prefill_priority_eviction
                ),
                "kvcache_prefill_direct_priority": (
                    config.kvcache_prefill_direct_priority
                ),
                "kvcache_prefill_normal_priority": (
                    config.kvcache_prefill_normal_priority
                ),
                "pool_poll_interval_s": config.pool_poll_interval_s,
                "pool_poll_include_sessions": config.pool_poll_include_sessions,
                "enable_backpressure": config.enable_backpressure,
                "backpressure_max_pause_s": config.backpressure_max_pause_s,
                "kvcache_migration_reject_threshold": config.kvcache_migration_reject_threshold,
                "sample_profile": config.sample_profile,
                "min_initial_input_tokens": config.min_initial_input_tokens,
                "max_initial_input_tokens": config.max_initial_input_tokens,
                "max_append_input_tokens": config.max_append_input_tokens,
                "max_output_tokens": config.max_output_tokens,
                "min_overlap_ratio": config.min_overlap_ratio,
                "sample_summary": asdict(sample_summary),
                "topology": {
                    "model_path": config.topology.model_path,
                    "router_url": topology.router_url,
                    "transfer_backend": topology.transfer_backend,
                    "force_rdma": topology.force_rdma,
                    "ib_device": topology.ib_device,
                    "prefill_workers": [
                        worker.worker_id for worker in topology.prefill_workers
                    ],
                    "decode_workers": [
                        worker.worker_id for worker in topology.decode_workers
                    ],
                    "direct_workers": [
                        worker.worker_id for worker in topology.direct_workers
                    ],
                },
            },
            handle,
            indent=2,
            sort_keys=True,
        )

    return BenchmarkArtifacts(
        run_dir=run_dir,
        sampled_trace_path=sampled_trace_path,
        metrics_path=run_dir / "request-metrics.jsonl",
        summary_path=run_dir / "request-metrics.jsonl.summary.json",
        benchmark_config_path=benchmark_config_path,
    )


def _decode_policy_for(policy_name: str) -> str:
    if policy_name == "sticky":
        return "manual"
    if policy_name == "kv-aware":
        return "consistent_hashing"
    return "round_robin"


def _header_mode_for(policy_name: str) -> str:
    if policy_name == "sticky":
        return "routing-key"
    if policy_name == "kv-aware":
        return "target-worker"
    return "none"
