import json
import os
from concurrent.futures import ThreadPoolExecutor
from typing import Any
from uuid import uuid4

import torch
from peft import PeftModel

from grasp.configs import ModelConfig
from grasp.model import Reasoning, Response, ToolCall


class VllmEngine:
    def __init__(
        self,
        model: str,
        tp_size: int = 1,
        tool_call_parser: str = "qwen3_coder",
        load_format: str = "auto",
        use_ray: bool | None = None,
        **engine_kwargs: Any,
    ) -> None:
        from transformers import AutoTokenizer
        from vllm import LLM

        # from vllm.config import WeightTransferConfig
        from vllm.tool_parsers import ToolParserManager

        # Auto-detect: use Ray if more than 1 GPU available
        if use_ray is None:
            use_ray = torch.cuda.device_count() > 1

        self.use_ray = use_ray
        self.tp_size = tp_size

        llm_kwargs: dict[str, Any] = dict(
            model=model,
            tensor_parallel_size=tp_size,
            load_format=load_format,
            # weight_transfer_config=WeightTransferConfig(backend="nccl"),
            enforce_eager=True,
            **engine_kwargs,
        )

        if use_ray:
            self.llm = _create_ray_llm(tp_size, llm_kwargs)
        else:
            self.llm = LLM(**llm_kwargs)

        # Tool parser runs locally (just text parsing, no GPU needed)
        tokenizer = AutoTokenizer.from_pretrained(model)
        parser_cls = ToolParserManager.get_tool_parser(tool_call_parser)
        self.tool_parser = parser_cls(tokenizer)

        # Weight transfer state (set by init_weight_transfer)
        self.model_update_group = None

    def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        config: ModelConfig,
    ) -> Response:
        from vllm import SamplingParams

        sampling_params = SamplingParams(
            temperature=config.temperature or 1.0,
            top_p=config.top_p or 1.0,
            max_tokens=config.max_completion_tokens,
            seed=config.seed,
            logprobs=0,
        )

        if config.model_kwargs:
            if "top_k" in config.model_kwargs:
                sampling_params.top_k = config.model_kwargs["top_k"]
            if "min_p" in config.model_kwargs:
                sampling_params.min_p = config.model_kwargs["min_p"]

        if self.use_ray:
            import ray

            outputs = ray.get(
                self.llm.chat.remote(
                    messages=messages,
                    tools=tools or None,
                    sampling_params=sampling_params,
                )
            )
        else:
            outputs = self.llm.chat(
                messages=messages,
                tools=tools or None,
                sampling_params=sampling_params,
            )

        if not outputs:
            return Response(id=uuid4().hex, message=None, tool_calls=[])

        return self._convert_output(outputs[0], tools)

    def init_weight_transfer(self, train_model: torch.nn.Module | PeftModel) -> None:
        """Initialize NCCL weight transfer between trainer and vLLM engine.

        Sets up the NCCL process group on both sides concurrently:
        trainer (rank 0) and vLLM workers (rank 1+).
        """
        from vllm.distributed.weight_transfer.nccl_engine import (
            NCCLWeightTransferEngine,
        )
        from vllm.utils.network_utils import get_ip, get_open_port

        master_address = get_ip()
        master_port = get_open_port()
        world_size = self.tp_size + 1  # vLLM workers + trainer

        init_request = dict(
            init_info=dict(
                master_address=master_address,
                master_port=master_port,
                rank_offset=1,
                world_size=world_size,
            )
        )
        trainer_info = dict(
            master_address=master_address,
            master_port=master_port,
            world_size=world_size,
        )

        # Both sides must init concurrently (NCCL collective)
        if self.use_ray:
            import ray

            inference_handle = self.llm.init_weight_transfer_engine.remote(init_request)
            self.model_update_group = NCCLWeightTransferEngine.trainer_init(
                trainer_info
            )
            ray.get(inference_handle)
        else:
            with ThreadPoolExecutor(max_workers=2) as pool:
                inference_future = pool.submit(
                    self.llm.init_weight_transfer_engine, init_request
                )
                trainer_future = pool.submit(
                    NCCLWeightTransferEngine.trainer_init, trainer_info
                )
                self.model_update_group = trainer_future.result()
                inference_future.result()

    def sync_weights(self, train_model: torch.nn.Module | PeftModel) -> None:
        """Broadcast weights from training model to vLLM engine via NCCL.

        Engine should be in sleep state. Both sides run concurrently.
        Handles LoRA models by temporarily merging adapters.
        """
        from vllm.distributed.weight_transfer.nccl_engine import (
            NCCLTrainerSendWeightsArgs,
            NCCLWeightTransferEngine,
        )

        assert self.model_update_group is not None, "Call init_weight_transfer() first"

        # For LoRA, temporarily merge adapters to get full weights
        merged = False
        if isinstance(train_model, PeftModel):
            train_model.merge_adapter()
            merged = True

        try:
            model_for_sync = (
                train_model.base_model.model
                if isinstance(train_model, PeftModel)
                else train_model
            )

            names, dtype_names, shapes = _get_weight_metadata(model_for_sync)
            update_request = dict(
                update_info=dict(
                    names=names,
                    dtype_names=dtype_names,
                    shapes=shapes,
                    packed=True,
                )
            )
            trainer_args = NCCLTrainerSendWeightsArgs(
                group=self.model_update_group,
                packed=True,
            )

            # Both sides must run concurrently (NCCL collective)
            if self.use_ray:
                import ray

                inference_handle = self.llm.update_weights.remote(update_request)
                NCCLWeightTransferEngine.trainer_send_weights(
                    iterator=model_for_sync.named_parameters(),
                    trainer_args=trainer_args,
                )
                ray.get(inference_handle)
            else:
                with ThreadPoolExecutor(max_workers=2) as pool:
                    inference_future = pool.submit(
                        self.llm.update_weights, update_request
                    )
                    trainer_future = pool.submit(
                        NCCLWeightTransferEngine.trainer_send_weights,
                        iterator=model_for_sync.named_parameters(),
                        trainer_args=trainer_args,
                    )
                    trainer_future.result()
                    inference_future.result()
        finally:
            if merged:
                train_model.unmerge_adapter()

    def sleep(self, level: int = 0) -> None:
        if self.use_ray:
            import ray

            ray.get(self.llm.sleep.remote(level=level))
        else:
            self.llm.sleep(level=level)

    def wake_up(self) -> None:
        if self.use_ray:
            import ray

            ray.get(self.llm.wake_up.remote(tags=["scheduling"]))
        else:
            self.llm.wake_up(tags=["scheduling"])

    # ── Output conversion ────────────────────────────────────────────

    def _convert_output(
        self, request_output: Any, tools: list[dict[str, Any]]
    ) -> Response:
        output = request_output.outputs[0]
        resp_id = uuid4().hex

        text = output.text.strip() if output.text else None
        if not text:
            text = None

        # Extract token IDs and logprobs
        token_ids = list(output.token_ids) if output.token_ids else None
        prompt_token_ids = (
            list(request_output.prompt_token_ids)
            if hasattr(request_output, "prompt_token_ids")
            and request_output.prompt_token_ids
            else None
        )

        token_logprobs = None
        token_texts = None
        if output.logprobs:
            token_logprobs = []
            token_texts = []
            for logprob_dict in output.logprobs:
                if logprob_dict:
                    sampled = next(iter(logprob_dict.values()))
                    token_logprobs.append(sampled.logprob)
                    token_texts.append(sampled.decoded_token)
                else:
                    token_logprobs.append(0.0)
                    token_texts.append("")

        # Extract reasoning content if present
        reasoning = None
        if hasattr(output, "reasoning_content") and output.reasoning_content:
            reasoning = Reasoning(
                id=uuid4().hex,
                content=output.reasoning_content.strip() or None,
            )

        # Parse tool calls using vLLM's official parser
        tool_calls = []
        message = text
        if text:
            extracted = self.tool_parser.extract_tool_calls(text, _MockRequest(tools))
            if extracted.tools_called and extracted.tool_calls:
                for tc in extracted.tool_calls:
                    tc_id = getattr(tc, "id", None) or f"call_{uuid4().hex[:8]}"
                    args_str = tc.function.arguments
                    try:
                        args = (
                            json.loads(args_str)
                            if isinstance(args_str, str)
                            else args_str
                        )
                    except json.JSONDecodeError:
                        args = {}
                    tool_calls.append(
                        ToolCall(id=tc_id, name=tc.function.name, args=args)
                    )
                message = extracted.content
                if message is not None:
                    message = message.strip() or None

        return Response(
            id=resp_id,
            message=message,
            reasoning=reasoning,
            tool_calls=tool_calls,
            prompt_token_ids=prompt_token_ids,
            token_ids=token_ids,
            token_logprobs=token_logprobs,
            token_texts=token_texts,
        )


class _MockRequest:
    def __init__(self, tools: list[dict[str, Any]]) -> None:
        self.tools = tools
        self.tool_choice = "auto"


def _create_ray_llm(tp_size: int, llm_kwargs: dict[str, Any]) -> Any:
    import ray
    from ray.util.placement_group import placement_group
    from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
    from vllm import LLM

    if not ray.is_initialized():
        ray.init()

    pg = placement_group([{"GPU": 1, "CPU": 0}] * tp_size)
    ray.get(pg.ready())
    scheduling = PlacementGroupSchedulingStrategy(
        placement_group=pg,
        placement_group_capture_child_tasks=True,
        placement_group_bundle_index=0,
    )

    bundle_indices = ",".join(str(i) for i in range(tp_size))

    class _RayLLM(LLM):
        def __init__(self, *args, **kwargs):
            os.environ["VLLM_RAY_BUNDLE_INDICES"] = bundle_indices
            super().__init__(*args, **kwargs)

    if tp_size > 1:
        llm_kwargs["distributed_executor_backend"] = "ray"

    return ray.remote(
        num_cpus=0,
        num_gpus=0,
        scheduling_strategy=scheduling,
    )(_RayLLM).remote(**llm_kwargs)


def _get_weight_metadata(
    model: torch.nn.Module,
) -> tuple[list[str], list[str], list[list[int]]]:
    """Extract weight names, dtype names, and shapes from model parameters."""
    names = []
    dtype_names = []
    shapes = []
    for name, p in model.named_parameters():
        names.append(name)
        dtype_names.append(str(p.dtype).split(".")[-1])
        shapes.append(list(p.shape))
    return names, dtype_names, shapes
