import logging
import os
import shlex
import signal
import subprocess
import tarfile
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import TYPE_CHECKING, Any

from tests.e2e.nightly.multi_node.external_dp.scripts.external_dp_config import (
    ROUTING_DISAGGREGATED_PREFILL,
    ExternalDPConfig,
    RankInfo,
)
from tests.e2e.nightly.multi_node.scripts.benchmark_results import (
    build_task_entry,
    extract_hardware,
    filter_environment,
    get_vllm_version,
    write_results_json,
)

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from tests.e2e.nightly.multi_node.external_dp.scripts.runtime import ServerCommand

SENSITIVE_ENV_TOKENS = ("TOKEN", "SECRET", "PASSWORD", "ACCESS_KEY")


def format_server_cmd(cmd: list[str], env: dict[str, str] | None = None) -> str:
    env_parts: list[str] = []
    for key, value in sorted((env or {}).items()):
        display_value = "***" if any(token in key.upper() for token in SENSITIVE_ENV_TOKENS) else str(value)
        env_parts.append(f"{key}={shlex.quote(display_value)}")
    return " ".join([*env_parts, shlex.join(cmd)])


def start_logged_process(cmd: list[str], env: dict[str, str], log_file: Path) -> subprocess.Popen:
    log_file.parent.mkdir(parents=True, exist_ok=True)
    merged_env = {**os.environ, **env}
    with log_file.open("ab") as f:
        f.write(f"Starting command: {format_server_cmd(cmd, env)}\n".encode())
        f.flush()
        return subprocess.Popen(
            cmd,
            stdout=f,
            stderr=subprocess.STDOUT,
            env=merged_env,
            start_new_session=True,
        )


def terminate_process_tree(pid: int, timeout: int = 30) -> None:
    try:
        import psutil
    except ModuleNotFoundError:
        try:
            os.killpg(pid, signal.SIGTERM)
        except ProcessLookupError:
            return
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                os.kill(pid, 0)
            except ProcessLookupError:
                return
            time.sleep(0.2)
        try:
            os.killpg(pid, signal.SIGKILL)
        except ProcessLookupError:
            return
        return

    try:
        parent = psutil.Process(pid)
    except psutil.NoSuchProcess:
        return

    children = parent.children(recursive=True)
    for process in children:
        process.terminate()
    parent.terminate()

    gone, alive = psutil.wait_procs([parent, *children], timeout=timeout)
    del gone
    for process in alive:
        process.kill()


def is_http_ready(url: str, timeout: float = 5.0) -> bool:
    try:
        with urllib.request.urlopen(url, timeout=timeout) as response:
            return 200 <= response.status < 300
    except (urllib.error.URLError, TimeoutError, OSError):
        return False


def wait_http_ready(url: str, timeout: int, interval: float = 2.0) -> None:
    deadline = time.monotonic() + timeout
    last_error: Exception | None = None
    while time.monotonic() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=5) as response:
                if 200 <= response.status < 300:
                    return
        except (urllib.error.URLError, TimeoutError, OSError) as exc:
            last_error = exc
        time.sleep(interval)
    raise TimeoutError(f"Timed out waiting for HTTP ready: {url}; last_error={last_error}")


def wait_http_unready(url: str, timeout: int, interval: float = 5.0) -> None:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if not is_http_ready(url):
            return
        time.sleep(interval)
    raise TimeoutError(f"Timed out waiting for HTTP unready: {url}")


def collect_logs(src_dir: Path, output_tar: Path) -> None:
    if not src_dir.exists():
        return
    output_tar.parent.mkdir(parents=True, exist_ok=True)
    with tarfile.open(output_tar, "w:gz") as tar:
        tar.add(src_dir, arcname=src_dir.name)


def _common_command_envs(commands: list["ServerCommand"]) -> dict[str, str]:
    if not commands:
        return {}

    common_keys = set(commands[0].env)
    for command in commands[1:]:
        common_keys.intersection_update(command.env)

    common_envs: dict[str, str] = {}
    for key in sorted(common_keys):
        values = {command.env[key] for command in commands}
        if len(values) == 1:
            common_envs[key] = next(iter(values))
    return common_envs


def _extract_dtype(config: ExternalDPConfig, commands: list["ServerCommand"]) -> str:
    has_w8a8 = "w8a8" in config.model.lower()
    has_quant_ascend = any("--quantization ascend" in command.display_cmd for command in commands)
    return "w8a8" if has_w8a8 and has_quant_ascend else "bf16"


def _extract_features(commands: list["ServerCommand"]) -> list[str]:
    if not commands:
        return []
    features: list[str] = []
    command_args = [command.cmd for command in commands]
    command_displays = [" ".join(shlex.quote(arg) for arg in command.cmd) for command in commands]

    if any("--async-scheduling" in cmd for cmd in command_args):
        features.append("async_scheduling")
    if any("--enable-expert-parallel" in cmd for cmd in command_args):
        features.append("expert_parallel")
    if any("--speculative-config" in cmd for cmd in command_args):
        features.append("speculative")
    if any("cudagraph_mode" in display for display in command_displays):
        features.append("aclgraph")

    feature_envs = {
        "VLLM_ASCEND_ENABLE_FLASHCOMM": "flashcomm",
        "VLLM_ASCEND_ENABLE_FLASHCOMM1": "flashcomm1",
        "VLLM_ASCEND_ENABLE_TOPK_OPTIMIZE": "topk_optimize",
        "VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE": "matmul_allreduce",
        "VLLM_ASCEND_ENABLE_MLAPO": "mlapo",
        "VLLM_ASCEND_ENABLE_FUSED_MC2": "fused_mc2",
    }
    for env_key, feature_name in feature_envs.items():
        values = [str(command.env.get(env_key, "0")) for command in commands]
        if any(value not in ("0", "", "false", "False") for value in values):
            features.append(feature_name)
    return features


def _build_serve_cmd(
    config: ExternalDPConfig,
    ranks: list[RankInfo],
    commands: list["ServerCommand"],
) -> dict[str, Any]:
    entries: dict[str, str] = {}
    for rank, command in zip(ranks, commands):
        prefix = rank.role
        if config.routing.type == ROUTING_DISAGGREGATED_PREFILL:
            prefix = "prefill" if rank.role == "prefiller" else "decode"
        entries[f"{prefix}-node{rank.node_index}-rank{rank.local_rank}"] = command.display_cmd
    key = "external_dp_pd" if config.routing.type == ROUTING_DISAGGREGATED_PREFILL else "external_dp"
    return {key: entries}


def build_benchmark_results(
    *,
    config: ExternalDPConfig,
    ranks: list[RankInfo],
    commands: list["ServerCommand"],
    results: list[Any],
) -> dict[str, Any]:
    valid_items = [(case["case_name"], case) for case in config.benchmark_cases]
    tasks = [build_task_entry(key, case, result) for (key, case), result in zip(valid_items, results)]
    runner = os.environ.get("VLLM_CI_RUNNER", "")
    common_envs = _common_command_envs(commands)

    return {
        "model_name": config.model,
        "hardware": extract_hardware(runner),
        "dtype": _extract_dtype(config, commands),
        "feature": _extract_features(commands),
        "vllm_version": get_vllm_version(),
        "vllm_ascend_version": os.environ.get("VLLM_ASCEND_REF", ""),
        "tasks": tasks,
        "serve_cmd": _build_serve_cmd(config, ranks, commands),
        "environment": filter_environment(common_envs),
    }


def write_benchmark_results_json(
    *,
    config: ExternalDPConfig,
    ranks: list[RankInfo],
    commands: list["ServerCommand"],
    results: list[Any],
    output_dir: Path | None = None,
) -> Path:
    output = build_benchmark_results(config=config, ranks=ranks, commands=commands, results=results)
    job_name = os.environ.get("BENCHMARK_JOB_NAME", "") or config.test_name.replace(" ", "-")
    return write_results_json(output, job_name=job_name, output_dir=output_dir)