import os
import torch
import pickle

import torch.distributed as dist

from typing import Callable
from abc import ABC, abstractmethod
from multiprocessing.synchronize import Event
from multiprocessing.shared_memory import SharedMemory

from diffulex.config import Config
from diffulex.distributed.parallel_state import fetch_parallel_state, init_parallel_state, init_process_group, reset_parallel_state
from diffulex.sampler.base import merge_sample_outputs
from diffulex.sampler import AutoSampler
from diffulex.engine.request import DllmReq
from diffulex.attention.metadata import set_warming_up, reset_warming_up
from diffulex.model import AutoModelForDiffusionLM
from diffulex.engine.strategy_registry import DiffulexStrategyRegistry
from diffulex.logger import get_logger
from diffulex.profiling import TorchProfileSession, record_function
from diffulex.vllm_compat import reset_vllm_compat_state, vllm_current_config


logger = get_logger(__name__)


class ModelRunnerBase(
    ABC,
):
    """Base class for model runners supporting different model types."""

    def __init__(self, config: Config, rank: int, event: Event | list[Event]):
        self.config = config
        hf_config = config.hf_config
        self.block_size = config.block_size
        self.page_size = config.kv_cache_page_size
        self.enforce_eager = config.enforce_eager
        config.enforce_eager = self.enforce_eager
        self.world_size = config.tensor_parallel_size
        self.rank = rank
        self.event = event

        if config.device_ids:
            device_id = config.device_ids[rank]
        else:
            device_id = config.device_start + rank
        assert 0 <= device_id < torch.cuda.device_count(), f"Invalid device_id {device_id}."

        # Initialize model, sampler, and kv cache
        init_method = f"tcp://{config.master_addr}:{config.master_port}"
        init_process_group(
            tp_size=config.tensor_parallel_size,
            ep_size=config.expert_parallel_size,
            dp_size=config.data_parallel_size,
            rank=rank,
            init_method=init_method,
            device_id=device_id,
            backend="nccl",
            timeout_seconds=config.distributed_timeout_seconds,
        )
        parallel_state = init_parallel_state(
            tp_size=config.tensor_parallel_size,
            ep_size=config.expert_parallel_size,
            dp_size=config.data_parallel_size,
        )
        a2a_requires_eager = (
            config.moe_dispatcher_backend == "naive"
            or (
                config.moe_dispatcher_backend == "deepep"
                and getattr(config, "deepep_mode", "auto") == "normal"
            )
        )
        if a2a_requires_eager and not self.enforce_eager:
            logger.warning(
                "Forcing enforce_eager=True for this expert-parallel topology/backend "
                "(tp_size=%s, dp_size=%s, ep_size=%s, moe_dispatcher_backend=%s, deepep_mode=%s).",
                parallel_state.tp_size,
                parallel_state.dp_size,
                parallel_state.ep_size,
                config.moe_dispatcher_backend,
                getattr(config, "deepep_mode", "auto"),
            )
            self.enforce_eager = True
            config.enforce_eager = True
        self.world_size = parallel_state.world_size
        self.rank = parallel_state.global_rank
        self.profile_session = TorchProfileSession("model_runner", rank=self.rank)
        self.dp_rank = parallel_state.dp_rank
        self.dp_world_size = parallel_state.dp_size
        self.cross_dp_ep = parallel_state.is_cross_dp_ep
        self.model_parallel_rank = parallel_state.model_parallel_rank
        self.is_model_parallel_root = self.model_parallel_rank == 0
        # Choose CUDA device for this TP rank.
        # config.device_ids is already a list of logical CUDA device indices (respecting CUDA_VISIBLE_DEVICES).
        # Do NOT add rank again, otherwise rank 1 with device_ids=[0,1] becomes device 2.
        torch.cuda.set_device(device_id)

        self.default_dtype = torch.get_default_dtype()
        self.default_dtype = (
            hf_config.torch_dtype if hasattr(hf_config, "torch_dtype") and hf_config.torch_dtype else torch.bfloat16
        )
        torch.set_default_dtype(self.default_dtype)
        torch.set_default_device(f"cuda:{device_id}")

        with vllm_current_config(config):
            self.model = self.load_model(config)
        self.sampler = self.load_sampler(config)
        self.allocate_kv_cache()
        self.warmup_model()

        if not self.enforce_eager:
            self.capture_cudagraph()

        self.start_worker_loop()

    def exit(self):
        if hasattr(self, "profile_session"):
            self.profile_session.stop()
        if not getattr(self, "_runner_exited", False):
            self._runner_exited = True
        else:
            return

        if not self.enforce_eager:
            for name in ("graphs", "graph_vars", "prefill_graphs", "graph_pool", "graph_capture_stream"):
                if hasattr(self, name):
                    try:
                        delattr(self, name)
                    except Exception:
                        logger.debug("Failed to delete CUDA graph attribute %s.", name, exc_info=True)

        if hasattr(self, "shm"):
            try:
                self.shm.close()
            except Exception:
                logger.debug("Failed to close shared memory on rank %s.", self.rank, exc_info=True)
            if self.rank == 0:
                try:
                    self.shm.unlink()
                except FileNotFoundError:
                    pass
                except Exception:
                    logger.debug("Failed to unlink shared memory on rank 0.", exc_info=True)

        try:
            torch.cuda.synchronize()
        except Exception:
            logger.debug("CUDA synchronize failed during runner exit on rank %s.", self.rank, exc_info=True)
        try:
            if dist.is_available() and dist.is_initialized():
                dist.destroy_process_group()
        except Exception:
            logger.debug("Failed to destroy process group on rank %s.", self.rank, exc_info=True)
        reset_vllm_compat_state()
        reset_parallel_state()

    def start_worker_loop(self):
        # Allocate shared memory for inter-process communication
        torch.set_default_device("cpu")
        torch.set_default_dtype(self.default_dtype)
        if self.world_size > 1:
            if self.rank == 0:
                try:
                    shm = SharedMemory(name=self.config.shm_name)
                    shm.close()
                    shm.unlink()
                except FileNotFoundError:
                    pass
                shm_size = 2**22
                self.shm = SharedMemory(name=self.config.shm_name, create=True, size=shm_size)
                dist.barrier()
            else:
                dist.barrier()
                self.shm = SharedMemory(name=self.config.shm_name)
                self.loop()

    def loop(self):
        try:
            while True:
                method_name, args = self.read_shm()
                self.call(method_name, *args)
                if method_name == "exit":
                    break
        except KeyboardInterrupt:
            self.exit()
            raise
        except BaseException:
            self.exit()
            raise

    def read_shm(self):
        assert self.world_size > 1 and self.rank
        self.event.wait()
        n = int.from_bytes(self.shm.buf[0:4], "little")
        method_name, *args = pickle.loads(self.shm.buf[4 : n + 4])
        self.event.clear()
        return method_name, args

    def write_shm(self, method_name, *args):
        assert self.world_size > 1 and not self.rank
        data = pickle.dumps([method_name, *args])
        n = len(data)

        if n + 4 > len(self.shm.buf):
            raise ValueError(
                f"Serialized data size ({n} bytes) exceeds shared memory buffer size ({len(self.shm.buf)} bytes). "
                f"Consider increasing shared memory size or reducing batch size."
            )

        self.shm.buf[0:4] = n.to_bytes(4, "little")
        self.shm.buf[4 : n + 4] = data
        for event in self.event:
            event.set()

    def call(self, method_name, *args):
        if self.world_size > 1 and self.rank == 0:
            self.write_shm(method_name, *args)
        method = getattr(self, method_name, None)
        if method_name == "run":
            self.profile_session.start()
            with record_function(f"diffulex.model_runner.rank{self.rank}.run"):
                result = method(*args)
            self.profile_session.step()
            return result
        with record_function(f"diffulex.model_runner.rank{self.rank}.{method_name}"):
            return method(*args)

    def load_model(self, config: Config):
        """Instantiate the underlying model; override to customize."""
        return AutoModelForDiffusionLM.from_config(config)

    def load_sampler(self, config: Config):
        """Instantiate the sampler implementation; override to customize."""
        return AutoSampler.from_config(config)

    def evict_sampler_state(self, req_ids: list[int] | list[str]) -> None:
        evict_fn = getattr(self.sampler, "evict_req_states", None)
        if evict_fn is not None:
            evict_fn(req_ids)

    def filter_local_reqs(self, reqs: list[DllmReq]) -> list[DllmReq]:
        if self.dp_world_size == 1:
            return reqs
        return [req for req in reqs if getattr(req, "dp_rank", 0) == self.dp_rank]

    def gather_dp_sample_output(self, sample_output):
        if self.dp_world_size == 1:
            return sample_output if self.is_model_parallel_root else None

        if not self.is_model_parallel_root:
            return None

        parallel_state = fetch_parallel_state()
        dp_group = parallel_state.get_dp_group()
        if dp_group is None:
            return sample_output

        gathered_outputs = [None] * self.dp_world_size if self.dp_rank == 0 else None
        dist.gather_object(sample_output, gathered_outputs, dst=0, group=dp_group)
        if self.dp_rank != 0:
            return None
        return merge_sample_outputs(gathered_outputs)

    @abstractmethod
    def _prefill_warmup(self):
        """Run template-specific prefill warmup."""
        pass

    def warmup_model(self):
        # TODO: attention metadata needs optimize for strategy awareness in warm-up
        if os.getenv("DIFFULEX_SKIP_WARMUP", "0") == "1":
            logger.warning("Skipping model warmup because DIFFULEX_SKIP_WARMUP=1.")
            return
        logger.info("Warming up model...")
        set_warming_up(True)
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        self._prefill_warmup()
        reset_warming_up()

    def allocate_kv_cache(self):
        config = self.config
        hf_config = config.hf_config
        free, total = torch.cuda.mem_get_info()
        used = total - free
        peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
        current = torch.cuda.memory_stats()["allocated_bytes.all.current"]
        parallel_state = fetch_parallel_state()
        num_kv_heads = (
            getattr(
                hf_config,
                "num_key_value_heads",
                getattr(hf_config, "n_kv_heads", None),
            )
            // parallel_state.get_tp_world_size()
        )

        if hasattr(hf_config, "head_dim"):
            head_dim = hf_config.head_dim
        elif hasattr(hf_config, "hidden_size") and hasattr(hf_config, "num_attention_heads"):
            head_dim = hf_config.hidden_size // hf_config.num_attention_heads
        else:
            raise AttributeError(f"Cannot determine head_dim from config: {type(hf_config)}")

        storage_dtype = torch.bfloat16
        itemsize = torch.empty(1, dtype=storage_dtype).element_size()
        page_bytes = 2 * hf_config.num_hidden_layers * self.page_size * num_kv_heads * head_dim * itemsize
        get_num_pages = lambda gpu_memory_utilization: (
            int(total * gpu_memory_utilization - used - peak + current) // page_bytes
        )
        try:
            num_pages = get_num_pages(config.gpu_memory_utilization)
            assert num_pages > 0
        except Exception:
            gpu_memory_utilization = config.gpu_memory_utilization
            while num_pages <= 200:
                logger.warning(
                    f"GPU memory utilization {gpu_memory_utilization} is too low to allocate kv cache. "
                    "Automatically adding 0.05."
                )
                gpu_memory_utilization += 0.05
                num_pages = get_num_pages(gpu_memory_utilization)
            logger.info(f"Set gpu_memory_utilization to {gpu_memory_utilization:.2f} to allocate kv cache.")
            config.gpu_memory_utilization = gpu_memory_utilization

        config.num_pages = num_pages
        logger.info(f"Allocated {config.num_pages} pages of size {self.page_size} for kv cache on rank {self.rank}.")

        # Cache the list of Attention-like modules once, to keep binding logic consistent
        # across cache layout branches (and avoid duplicated traversal).
        attn_modules = [m for m in self.model.modules() if hasattr(m, "k_cache") and hasattr(m, "v_cache")]

        if config.kv_cache_layout == "distinct":
            x = config.k_cache_hdim_split_factor_x
            self.k_cache = torch.zeros(
                hf_config.num_hidden_layers,
                config.num_pages,
                num_kv_heads,
                head_dim // x,
                self.page_size,
                x,
                dtype=storage_dtype,
            )
            self.v_cache = torch.zeros(
                hf_config.num_hidden_layers,
                config.num_pages,
                num_kv_heads,
                head_dim,
                self.page_size,
                dtype=storage_dtype,
            )
            for layer_id, module in enumerate(attn_modules):
                module.k_cache = self.k_cache[layer_id]
                module.v_cache = self.v_cache[layer_id]

        elif config.kv_cache_layout == "unified":
            self.kv_cache = torch.zeros(
                2,
                hf_config.num_hidden_layers,
                config.num_pages,
                self.page_size,
                num_kv_heads,
                head_dim,
                dtype=storage_dtype,
            )
            for layer_id, module in enumerate(attn_modules):
                module.k_cache = self.kv_cache[0, layer_id]
                module.v_cache = self.kv_cache[1, layer_id]

        else:
            raise ValueError(
                "Unsupported kv_cache_layout: {layout}. Supported values are 'distinct' and 'unified'.".format(
                    layout=config.kv_cache_layout
                )
            )

    def prepare_page_tables(self, reqs: list[DllmReq]):
        if not reqs:
            return torch.empty((0, 1), dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        max_len = max(len(req.page_table) for req in reqs)
        page_tables = [req.page_table + [-1] * (max_len - len(req.page_table)) for req in reqs]
        page_tables = torch.tensor(page_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True)
        return page_tables

    def init_attn_metadata_fn(self, set_fn: Callable, reset_fn: Callable, fetch_fn: Callable):
        self.set_attn_metadata = set_fn
        self.reset_attn_metadata = reset_fn
        self.fetch_attn_metadata = fetch_fn

    @abstractmethod
    def prepare_prefill(self, reqs: list[DllmReq]):
        """Model-specific prefill preparation."""
        pass

    @abstractmethod
    def prepare_decode(self, reqs: list[DllmReq]):
        """Model-specific decode preparation."""
        pass

    def prepare_sample(self, reqs: list[DllmReq]):
        temperatures = []
        for req in reqs:
            temperatures.append(req.temperature)
        temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True)
        return temperatures

    @abstractmethod
    @torch.inference_mode()
    def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor):
        """Model-specific forward pass."""
        pass

    @abstractmethod
    def run(self, reqs: list[DllmReq]) -> list[int]:
        """Main inference pipeline."""
        pass

    @abstractmethod
    @torch.inference_mode()
    def capture_cudagraph(self):
        """Model-specific CUDA graph capture."""
        pass


RunnerFactory = Callable[[Config, int, Event | list[Event]], "ModelRunnerBase"]


class AutoModelRunner(DiffulexStrategyRegistry):
    @classmethod
    def from_config(cls, config: Config, rank: int, event: Event | list[Event]):
        # Ensure project root is in sys.path for spawn mode subprocesses
        import sys
        import os

        if not any("diffulex_kernel" in p for p in sys.path):
            # Try to find project root by locating diffulex package
            diffulex_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
            if os.path.basename(diffulex_path) == "diffulex":
                project_root = os.path.dirname(diffulex_path)
                if project_root not in sys.path:
                    sys.path.insert(0, project_root)

        cls._ensure_strategies_loaded()
        cls._MODULE_MAPPING: dict[str, RunnerFactory]
        candidates: list[str] = []
        if config.decoding_strategy:
            candidates.append(config.decoding_strategy)
        candidates.append(cls._DEFAULT_KEY)

        for key in candidates:
            factory = cls._MODULE_MAPPING.get(key)
            if factory is not None:
                return factory(config, rank, event)

        available = ", ".join(cls.available_modules()) or "<none>"
        raise ValueError(
            "No model runner registered for decoding_strategy="
            f"'{config.decoding_strategy}'. Available runners: {available}."
        )