# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from typing import Optional

from dynamo.planner.utils.decode_planner import DecodePlanner
from dynamo.planner.utils.dryrun_plot_utils import create_dryrun_plot
from dynamo.planner.utils.planner_config import PlannerConfig
from dynamo.planner.utils.planner_core import (
    PlannerSharedState,
    _apply_component_gpu_budget,
    _apply_global_gpu_budget,
)
from dynamo.planner.utils.prefill_planner import PrefillPlanner
from dynamo.planner.utils.trace_data_extractor import extract_metrics_from_mooncake


def run_sla_planner_dryrun(
    config: PlannerConfig,
    dataset: str,
    start_num_p: int = 1,
    start_num_d: int = 1,
    output_plot: str = "dryrun_plot.png",
) -> None:
    if config.enable_load_scaling:
        raise ValueError(
            "Load-based scaling is not supported in dryrun mode. "
            "Set enable_load_scaling to false in the config."
        )

    if config.prefill_engine_num_gpu is None:
        config.prefill_engine_num_gpu = 1
    if config.decode_engine_num_gpu is None:
        config.decode_engine_num_gpu = 1

    warmup_metrics = None
    if config.load_predictor_warmup_trace is not None:
        warmup_metrics = extract_metrics_from_mooncake(
            config.load_predictor_warmup_trace,
            config.throughput_adjustment_interval,
        )

    metrics = extract_metrics_from_mooncake(
        dataset, config.throughput_adjustment_interval
    )
    if not metrics:
        raise ValueError("Empty metrics dataset: cannot run dryrun")

    mode = config.mode
    prefill_planner: Optional[PrefillPlanner] = None
    decode_planner: Optional[DecodePlanner] = None
    if mode == "disagg":
        shared_state = PlannerSharedState()
        prefill_planner = PrefillPlanner(
            None, config, dryrun=True, shared_state=shared_state
        )
        decode_planner = DecodePlanner(
            None, config, dryrun=True, shared_state=shared_state
        )
    elif mode == "prefill":
        prefill_planner = PrefillPlanner(None, config, dryrun=True)
    elif mode == "decode":
        decode_planner = DecodePlanner(None, config, dryrun=True)
    else:
        raise ValueError(f"Invalid planner mode: {mode}")

    def compute_safe_p_thpt(num_p: int, isl: float, ttft: float):
        """safe throughput is maximum throughput that the engine can handle given the TTFT SLA"""
        assert prefill_planner is not None
        actual_ttft = prefill_planner.prefill_interpolator.interpolate_ttft(isl)
        if actual_ttft > ttft:
            return 0
        return num_p * prefill_planner.prefill_interpolator.interpolate_thpt_per_gpu(
            isl
        )

    def compute_safe_d_thpt(num_d: int, isl: float, osl: float, itl: float):
        """safe throughput is maximum throughput that the engine can handle given the ITL SLA"""
        assert decode_planner is not None
        (
            pred_decode_thpt_per_gpu,
            actual_itl,
            _,
        ) = decode_planner.decode_interpolator.find_best_throughput_per_gpu(
            itl=itl, context_length=isl + osl / 2
        )
        if actual_itl > itl:
            return 0
        return num_d * pred_decode_thpt_per_gpu

    time_series = [0]
    rr = [metrics[0]["request_count"]]
    est_rr = [metrics[0]["request_count"]]
    isl = [metrics[0]["avg_isl"]]
    est_isl = [metrics[0]["avg_isl"]]
    osl = [metrics[0]["avg_osl"]]
    est_osl = [metrics[0]["avg_osl"]]

    interval = config.throughput_adjustment_interval

    if prefill_planner is not None:
        num_p = [start_num_p]
        p_thpt = [rr[0] * isl[0]]
        safe_p_thpt = [compute_safe_p_thpt(start_num_p, isl[0], config.ttft) * interval]
        prefill_planner.dryrun_observe_traffic_stats(rr[0], isl[0], osl[0])
    else:
        num_p = [0]
        p_thpt = [0]
        safe_p_thpt = [0]

    if decode_planner is not None:
        num_d = [start_num_d]
        d_thpt = [rr[0] * osl[0]]
        safe_d_thpt = [
            compute_safe_d_thpt(start_num_d, isl[0], osl[0], config.itl) * interval
        ]
        decode_planner.dryrun_observe_traffic_stats(rr[0], isl[0], osl[0])
    else:
        num_d = [0]
        d_thpt = [0]
        safe_d_thpt = [0]

    predictor_planner = prefill_planner or decode_planner
    assert predictor_planner is not None

    for metric in metrics[1:]:
        time_series.append(time_series[-1] + interval)

        _est_rr, _est_isl, _est_osl = predictor_planner.predict_load()
        est_rr.append(_est_rr)
        est_isl.append(_est_isl)
        est_osl.append(_est_osl)

        _num_p = (
            prefill_planner._compute_replica_requirements(_est_rr, _est_isl, _est_osl)
            if prefill_planner is not None
            else 0
        )
        _num_d = (
            decode_planner._compute_replica_requirements(_est_rr, _est_isl, _est_osl)
            if decode_planner is not None
            else 0
        )

        if prefill_planner is not None and decode_planner is not None:
            _num_p, _num_d = _apply_global_gpu_budget(_num_p, _num_d, config)
        elif prefill_planner is not None:
            _num_p = _apply_component_gpu_budget(
                _num_p, config.prefill_engine_num_gpu, config
            )
        elif decode_planner is not None:
            _num_d = _apply_component_gpu_budget(
                _num_d, config.decode_engine_num_gpu, config
            )

        num_p.append(_num_p)
        num_d.append(_num_d)

        for planner in [prefill_planner, decode_planner]:
            if planner is not None:
                planner.dryrun_observe_traffic_stats(
                    metric["request_count"], metric["avg_isl"], metric["avg_osl"]
                )

        rr.append(metric["request_count"])
        isl.append(metric["avg_isl"])
        osl.append(metric["avg_osl"])

        p_thpt.append(rr[-1] * isl[-1] if prefill_planner is not None else 0)
        d_thpt.append(rr[-1] * osl[-1] if decode_planner is not None else 0)

        safe_p_thpt.append(
            compute_safe_p_thpt(num_p[-1], isl[-1], config.ttft) * interval
            if prefill_planner is not None
            else 0
        )
        safe_d_thpt.append(
            compute_safe_d_thpt(num_d[-1], isl[-1], osl[-1], config.itl) * interval
            if decode_planner is not None
            else 0
        )

    warmup_time = None
    warmup_rr = None
    warmup_isl = None
    warmup_osl = None
    if warmup_metrics:
        n = len(warmup_metrics)
        warmup_time = [-(n - i) * interval for i in range(n)]
        warmup_rr = [m["request_count"] for m in warmup_metrics]
        warmup_isl = [m["avg_isl"] for m in warmup_metrics]
        warmup_osl = [m["avg_osl"] for m in warmup_metrics]

    create_dryrun_plot(
        time=time_series,
        rr=rr,
        est_rr=est_rr,
        isl=isl,
        est_isl=est_isl,
        osl=osl,
        est_osl=est_osl,
        num_p=num_p,
        p_thpt=p_thpt,
        safe_p_thpt=safe_p_thpt,
        num_d=num_d,
        d_thpt=d_thpt,
        safe_d_thpt=safe_d_thpt,
        output_path=output_plot,
        warmup_time=warmup_time,
        warmup_rr=warmup_rr,
        warmup_isl=warmup_isl,
        warmup_osl=warmup_osl,
    )
