from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Literal


WorkerRole = Literal["prefill", "decode", "direct"]


@dataclass(frozen=True)
class WorkerSpec:
    role: WorkerRole
    ordinal: int
    gpu_ids: tuple[int, ...]
    host: str
    port: int
    bootstrap_port: int | None = None

    @property
    def worker_id(self) -> str:
        return f"{self.role}-{self.ordinal}"

    @property
    def url(self) -> str:
        return f"http://{self.host}:{self.port}"

    @property
    def gpu_id(self) -> int:
        return self.gpu_ids[0]

    @property
    def tp_size(self) -> int:
        return len(self.gpu_ids)


@dataclass(frozen=True)
class SingleNodeTopology:
    model_path: str
    prefill_workers: tuple[WorkerSpec, ...]
    decode_workers: tuple[WorkerSpec, ...]
    direct_workers: tuple[WorkerSpec, ...]
    router_host: str
    router_port: int
    transfer_backend: str
    trust_remote_code: bool
    force_rdma: bool = False
    ib_device: str | None = None
    extra_server_args: tuple[str, ...] = ()
    prefill_extra_server_args: tuple[str, ...] = ()
    decode_extra_server_args: tuple[str, ...] = ()
    direct_extra_server_args: tuple[str, ...] = ()

    @property
    def model_name(self) -> str:
        return Path(self.model_path).name

    @property
    def router_url(self) -> str:
        return f"http://{self.router_host}:{self.router_port}"

    @property
    def route_workers(self) -> tuple[WorkerSpec, ...]:
        if self.decode_workers:
            return self.decode_workers
        return self.direct_workers

    def route_index(self, worker_id: str) -> int:
        for idx, worker in enumerate(self.route_workers):
            if worker.worker_id == worker_id:
                return idx
        raise KeyError(f"Unknown route worker: {worker_id}")


def build_single_node_topology(
    *,
    model_path: str,
    prefill_worker_count: int,
    decode_worker_count: int,
    direct_worker_count: int = 0,
    prefill_tp_size: int = 1,
    decode_tp_size: int = 1,
    direct_tp_size: int = 1,
    prefill_gpu_ids: tuple[int, ...] | None = None,
    decode_gpu_ids: tuple[int, ...] | None = None,
    direct_gpu_ids: tuple[int, ...] | None = None,
    total_gpu_budget: int = 8,
    host: str = "127.0.0.1",
    router_port: int = 8000,
    prefill_port_base: int = 30000,
    decode_port_base: int = 31000,
    direct_port_base: int = 32000,
    bootstrap_port_base: int = 8998,
    transfer_backend: str = "nixl",
    force_rdma: bool = False,
    trust_remote_code: bool = True,
    ib_device: str | None = None,
    extra_server_args: tuple[str, ...] = (),
    prefill_extra_server_args: tuple[str, ...] = (),
    decode_extra_server_args: tuple[str, ...] = (),
    direct_extra_server_args: tuple[str, ...] = (),
) -> SingleNodeTopology:
    if prefill_worker_count < 0:
        raise ValueError("prefill_worker_count must be >= 0")
    if decode_worker_count < 0:
        raise ValueError("decode_worker_count must be >= 0")
    if direct_worker_count < 0:
        raise ValueError("direct_worker_count must be >= 0")
    if (
        prefill_worker_count == 0
        and decode_worker_count == 0
        and direct_worker_count == 0
    ):
        raise ValueError("At least one worker must be configured")
    if prefill_tp_size <= 0:
        raise ValueError("prefill_tp_size must be >= 1")
    if decode_tp_size <= 0:
        raise ValueError("decode_tp_size must be >= 1")
    if direct_tp_size <= 0:
        raise ValueError("direct_tp_size must be >= 1")
    if force_rdma and not ib_device:
        raise ValueError("force_rdma requires --ib-device to be set")
    if force_rdma and transfer_backend != "mooncake":
        raise ValueError(
            "force_rdma currently requires transfer_backend='mooncake' "
            "to guarantee an RDMA path"
        )

    total_gpus_required = (
        prefill_worker_count * prefill_tp_size
        + decode_worker_count * decode_tp_size
        + direct_worker_count * direct_tp_size
    )
    if total_gpus_required > total_gpu_budget:
        raise ValueError(
            "Single-node GPU budget exceeded: "
            f"{prefill_worker_count} prefill x tp={prefill_tp_size} + "
            f"{decode_worker_count} decode x tp={decode_tp_size} + "
            f"{direct_worker_count} direct x tp={direct_tp_size} > "
            f"{total_gpu_budget} GPUs"
        )

    if prefill_gpu_ids is None:
        prefill_gpu_ids = tuple(range(prefill_worker_count * prefill_tp_size))
    if decode_gpu_ids is None:
        decode_gpu_ids = tuple(
            range(
                len(prefill_gpu_ids),
                len(prefill_gpu_ids) + decode_worker_count * decode_tp_size,
            )
        )
    if direct_gpu_ids is None:
        direct_gpu_ids = tuple(
            range(
                len(prefill_gpu_ids) + len(decode_gpu_ids),
                len(prefill_gpu_ids)
                + len(decode_gpu_ids)
                + direct_worker_count * direct_tp_size,
            )
        )

    if len(prefill_gpu_ids) != prefill_worker_count * prefill_tp_size:
        raise ValueError(
            "prefill_gpu_ids length must equal prefill_worker_count * prefill_tp_size: "
            f"{len(prefill_gpu_ids)} != {prefill_worker_count * prefill_tp_size}"
        )
    if len(decode_gpu_ids) != decode_worker_count * decode_tp_size:
        raise ValueError(
            "decode_gpu_ids length must equal decode_worker_count * decode_tp_size: "
            f"{len(decode_gpu_ids)} != {decode_worker_count * decode_tp_size}"
        )
    if len(direct_gpu_ids) != direct_worker_count * direct_tp_size:
        raise ValueError(
            "direct_gpu_ids length must equal direct_worker_count * direct_tp_size: "
            f"{len(direct_gpu_ids)} != {direct_worker_count * direct_tp_size}"
        )
    assigned_gpu_ids = prefill_gpu_ids + decode_gpu_ids + direct_gpu_ids
    if len(set(assigned_gpu_ids)) != len(assigned_gpu_ids):
        raise ValueError("prefill/decode/direct GPU IDs must be unique")
    if any(gpu_id < 0 or gpu_id >= total_gpu_budget for gpu_id in assigned_gpu_ids):
        raise ValueError(
            "GPU IDs must fall within the single-node budget range "
            f"[0, {total_gpu_budget - 1}]"
        )

    prefill_workers = tuple(
        WorkerSpec(
            role="prefill",
            ordinal=idx,
            gpu_ids=tuple(
                prefill_gpu_ids[
                    idx * prefill_tp_size : (idx + 1) * prefill_tp_size
                ]
            ),
            host=host,
            port=prefill_port_base + idx,
            bootstrap_port=bootstrap_port_base + idx,
        )
        for idx in range(prefill_worker_count)
    )
    decode_workers = tuple(
        WorkerSpec(
            role="decode",
            ordinal=idx,
            gpu_ids=tuple(
                decode_gpu_ids[
                    idx * decode_tp_size : (idx + 1) * decode_tp_size
                ]
            ),
            host=host,
            port=decode_port_base + idx,
        )
        for idx in range(decode_worker_count)
    )
    direct_workers = tuple(
        WorkerSpec(
            role="direct",
            ordinal=idx,
            gpu_ids=tuple(
                direct_gpu_ids[
                    idx * direct_tp_size : (idx + 1) * direct_tp_size
                ]
            ),
            host=host,
            port=direct_port_base + idx,
        )
        for idx in range(direct_worker_count)
    )

    return SingleNodeTopology(
        model_path=model_path,
        prefill_workers=prefill_workers,
        decode_workers=decode_workers,
        direct_workers=direct_workers,
        router_host=host,
        router_port=router_port,
        transfer_backend=transfer_backend,
        trust_remote_code=trust_remote_code,
        force_rdma=force_rdma,
        ib_device=ib_device,
        extra_server_args=extra_server_args,
        prefill_extra_server_args=prefill_extra_server_args,
        decode_extra_server_args=decode_extra_server_args,
        direct_extra_server_args=direct_extra_server_args,
    )
