import json
import os
import threading
from collections.abc import Callable
from contextlib import AbstractContextManager, ExitStack, contextmanager
from pathlib import Path
from typing import Any

import pytest

from tests.dfx.conftest import (
    create_paired_omni_benchmark_pytest_params,
    create_test_parameter_mapping,
    get_benchmark_params_for_server,
    get_runtime_resource_label,
    is_diffusion_perf_config,
    load_benchmark_configs,
    run_benchmark,
)
from tests.helpers.runtime import OmniServer

# Optional JSON field ``mark`` is applied as pytest marks via
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).


os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def _get_config_file_from_argv() -> str | None:
    """Read ``--test-config-file`` from ``sys.argv`` at import time so parametrization can use it."""
    import sys

    for i, arg in enumerate(sys.argv):
        if arg == "--test-config-file" and i + 1 < len(sys.argv):
            return sys.argv[i + 1]
        if arg.startswith("--test-config-file="):
            return arg.split("=", 1)[1]
    return None


_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"

CONFIG_FILE_PATH = _get_config_file_from_argv()
if CONFIG_FILE_PATH is None:
    _all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR)
    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)]
    print(
        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts case(s) from "
        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; "
        f"use -m to filter, e.g. -m tts)"
    )
else:
    BENCHMARK_CONFIGS = load_benchmark_configs(CONFIG_FILE_PATH)

DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS)
paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)

_omni_server_lock = threading.Lock()


class _SingleActiveContext:
    """Reuse one active context while its configuration key is unchanged."""

    def __init__(self) -> None:
        self._key: Any = None
        self._stack: ExitStack | None = None
        self._value: Any = None

    def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any:
        if self._stack is not None and key == self._key:
            return self._value

        self.close()
        stack = ExitStack()
        value = stack.enter_context(factory())
        self._key = key
        self._stack = stack
        self._value = value
        return value

    def close(self) -> None:
        stack = self._stack
        self._key = None
        self._stack = None
        self._value = None
        if stack is not None:
            stack.close()


@contextmanager
def _start_omni_server(server_param):
    test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param

    print(f"Starting OmniServer with test: {test_name}, model: {model}")

    server_args: list[str] = []
    if use_omni:
        server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
    # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
    # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
    if stage_config_path:
        server_args = ["--deploy-config", stage_config_path] + server_args
    if stage_overrides:
        server_args = ["--stage-overrides", stage_overrides] + server_args
    if extra_cli_args:
        server_args = list(extra_cli_args) + server_args
    with OmniServer(model, server_args, use_omni=use_omni) as server:
        server.test_name = test_name
        print("OmniServer started successfully")
        yield server
        print("OmniServer stopping...")

    print("OmniServer stopped")


@pytest.fixture(scope="module")
def omni_server_context():
    """Start vLLM-Omni server as a subprocess with actual model weights.
    Reuse it for adjacent benchmark cases with the same server configuration.
    Multi-stage initialization can take 10-20+ minutes.
    """
    with _omni_server_lock:
        active_context = _SingleActiveContext()
        try:
            yield active_context
        finally:
            active_context.close()


@pytest.fixture
def omni_server(request, omni_server_context):
    return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))


@pytest.fixture
def benchmark_params(request):
    """Benchmark parameters fixture; paired with ``omni_server`` via parametrization."""
    test_name, param_index = request.param

    all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)

    if not all_params:
        raise ValueError(f"No benchmark parameters found for test: {test_name}")

    if param_index >= len(all_params):
        raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")

    current = param_index + 1
    total = len(all_params)
    print(f"\n  Running benchmark {current}/{total} for {test_name}")

    return {
        "test_name": test_name,
        "params": all_params[param_index],
    }


def assert_result(result, num_prompt) -> None:
    assert result["completed"] == num_prompt, "Request failures exist"


@pytest.mark.benchmark
@pytest.mark.parametrize(
    "omni_server,benchmark_params",
    paired_benchmark_params,
    indirect=["omni_server", "benchmark_params"],
)
def test_performance_benchmark(omni_server, benchmark_params):
    test_name = benchmark_params["test_name"]
    params = benchmark_params["params"]
    dataset_name = params.get("dataset_name", "")

    host = omni_server.host
    port = omni_server.port
    model = omni_server.model

    print(f"Running benchmark for model: {model}")
    print(f"Benchmark parameters: {benchmark_params}")

    resource_label = get_runtime_resource_label()

    def to_list(value, default=None):
        if value is None:
            return [] if default is None else [default]
        return [value] if not isinstance(value, (list, tuple)) else list(value)

    qps_list = to_list(params.get("request_rate"))
    num_prompt_list = to_list(params.get("num_prompts"))
    max_concurrency_list = to_list(params.get("max_concurrency"))

    max_len = max(len(qps_list), len(max_concurrency_list))
    if len(num_prompt_list) == 1 and max_len > 1:
        num_prompt_list = num_prompt_list * max_len
    elif max_len == 1 and len(num_prompt_list) > 1:
        if len(qps_list) == 1:
            qps_list = qps_list * len(num_prompt_list)
        if len(max_concurrency_list) == 1:
            max_concurrency_list = max_concurrency_list * len(num_prompt_list)
        max_len = max(len(qps_list), len(max_concurrency_list))
    elif len(num_prompt_list) != max_len and max_len > 0:
        raise ValueError("The number of prompts does not match the QPS or max_concurrency")

    args = ["--host", host, "--port", str(port)]
    exclude_keys = {
        "request_rate",
        "baseline",
        "num_prompts",
        "max_concurrency",
        "task",
        "enabled",
        "eval_phase",
        "trust_remote_code",
    }

    for key, value in params.items():
        if key in exclude_keys or value is None:
            continue

        arg_name = f"--{key.replace('_', '-')}"

        if isinstance(value, bool) and value:
            args.append(arg_name)
        elif isinstance(value, dict):
            json_str = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
            args.extend([arg_name, json_str])
        elif not isinstance(value, bool):
            args.extend([arg_name, str(value)])

    for config in BENCHMARK_CONFIGS:
        if config.get("test_name") != test_name:
            continue
        server_params = config.get("server_params") or {}
        if server_params.get("trust_remote_code") or params.get("trust_remote_code"):
            args.append("--trust-remote-code")
        break

    # QPS / request-rate sweep
    for sweep_index, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
        args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
        result = run_benchmark(
            args=args,
            test_name=test_name,
            flow=qps,
            dataset_name=dataset_name,
            num_prompt=num_prompt,
            baseline_config=params.get("baseline"),
            sweep_index=sweep_index,
            random_input_len=params.get("random_input_len"),
            random_output_len=params.get("random_output_len"),
            resource_label=resource_label,
        )
        assert_result(result, num_prompt)

    # concurrency test
    for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
        args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
        result = run_benchmark(
            args=args,
            test_name=test_name,
            flow=concurrency,
            dataset_name=dataset_name,
            num_prompt=num_prompt,
            baseline_config=params.get("baseline"),
            sweep_index=sweep_index,
            random_input_len=params.get("random_input_len"),
            random_output_len=params.get("random_output_len"),
            resource_label=resource_label,
        )
        assert_result(result, num_prompt)