import atexit
import logging
import requests
import threading
import time
import torch
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict
from packaging import version
from pydantic import ValidationError
from requests import ConnectionError
from torch import nn
from typing import List, Optional, Union
from urllib.parse import urlparse

from swift.infer_engine import AdapterRequest, RequestConfig
from swift.infer_engine.protocol import ChatCompletionResponse, RolloutInferRequest, RolloutOutput
from swift.metrics import Metric
from swift.utils import (is_trl_available, is_vllm_ascend_available, is_vllm_available, is_vllm_metax_available,
                         synchronize)
from .utils import (broadcast_tensor_for_vllm_weight_sync, format_host_for_url, is_valid_ipv6_address,
                    peft_config_to_dict, resolve_hostname)

if is_vllm_available():
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
    from vllm.distributed.utils import StatelessProcessGroup

    if is_vllm_ascend_available():
        from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator  # noqa

    if is_vllm_metax_available():
        import vllm_metax.patch

if is_trl_available():
    import trl
    trl_verison = version.parse(trl.__version__)

logger = logging.getLogger(__name__)


class VLLMInferClient:
    """Inference-only vLLM client. Posts to /infer/ endpoint.
    No weight synchronization. Used for GKD teacher server etc.
    """

    def __init__(self,
                 base_urls: Optional[List[str]] = None,
                 hosts: List[str] = ['0.0.0.0'],
                 server_ports: List[int] = [8000],
                 connection_timeout: float = 240.0):
        if not is_vllm_available():
            raise ImportError('vLLM is not installed. Please install it with `pip install vllm`.')

        if base_urls is not None:
            self.base_urls = []
            self.hosts = []
            for url in base_urls:
                parsed_url = urlparse(url)
                # Use resolve_hostname instead of gethostbyname for IPv6 support
                host = resolve_hostname(parsed_url.hostname)
                scheme = parsed_url.scheme or 'http'
                base_url_i = f'{scheme}://{parsed_url.netloc}{parsed_url.path}'
                self.base_urls.append(base_url_i)
                self.hosts.append(host)
        else:
            if len(hosts) != len(server_ports):
                raise ValueError('host and server_port must have same length when lists are provided')
            # Format IPv6 addresses correctly in URLs (wrap with brackets)
            self.base_urls = [f'http://{format_host_for_url(h)}:{p}' for h, p in zip(hosts, server_ports)]
            self.hosts = hosts

        self.num_servers = len(self.base_urls)
        self.sessions = [requests.Session() for _ in range(self.num_servers)]
        self.check_server(connection_timeout)

    def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0):
        server_status = [False] * self.num_servers

        def check_single_server(i):
            start_time = time.time()
            url = f'{self.base_urls[i]}/health/'
            while True:
                try:
                    response = requests.get(url, timeout=retry_interval)
                    if response.status_code == 200:
                        server_status[i] = True
                        return
                except Exception:
                    pass

                elapsed = time.time() - start_time
                if elapsed >= total_timeout:
                    return

                time.sleep(retry_interval)

        threads = []
        for i in range(self.num_servers):
            t = threading.Thread(target=check_single_server, args=(i, ))
            t.daemon = True
            t.start()
            threads.append(t)

        for t in threads:
            t.join(total_timeout)

        if not all(server_status):
            failed_servers = [self.base_urls[i] for i, status in enumerate(server_status) if not status]
            raise ConnectionError(f'Servers not reachable after {total_timeout}s: {failed_servers}')

    def infer(
        self,
        infer_requests: List[RolloutInferRequest],
        request_config: Optional[RequestConfig] = None,
        metrics: Optional[List[Metric]] = None,
        *,
        use_tqdm: Optional[bool] = None,
        adapter_request: Optional[AdapterRequest] = None,
    ):
        n = len(infer_requests)
        chunk_size = (n + self.num_servers - 1) // self.num_servers
        chunks = [infer_requests[i:i + chunk_size] for i in range(0, n, chunk_size)]
        chunks += [[]] * (self.num_servers - len(chunks))

        results = [None] * self.num_servers
        errors = [None] * self.num_servers
        if isinstance(request_config, RequestConfig):
            request_config = asdict(request_config)

        def process_chunk(i, chunk):
            try:
                if len(chunk) > 0 and isinstance(chunk[0], RolloutInferRequest):
                    chunk = [asdict(req) for req in chunk]

                response = self.sessions[i].post(
                    f'{self.base_urls[i]}/infer/',
                    json={
                        'infer_requests': chunk,
                        'request_config': request_config,
                        'metrics': metrics,
                        'use_tqdm': use_tqdm,
                        'adapter_request': adapter_request,
                    },
                )

                if response.status_code != 200:
                    errors[i] = Exception(f'Server {i} failed: {response.status_code}, {response.text}')
                    return

                resp_data = response.json()
                parsed: List[Union[RolloutOutput, ChatCompletionResponse]] = []

                for item in resp_data:
                    try:
                        parsed.append(RolloutOutput.model_validate(item))
                    except ValidationError:
                        parsed.append(ChatCompletionResponse(**item))
                results[i] = parsed
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(process_chunk, i, chunk) for i, chunk in enumerate(chunks)]
            for future in futures:
                future.result()

        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors: {all_errors}')

        return [res for server_results in results for res in server_results]


class VLLMClient(VLLMInferClient):

    def __init__(self,
                 base_urls: Optional[List[str]] = None,
                 hosts: List[str] = ['0.0.0.0'],
                 server_ports: List[int] = [8000],
                 group_ports: Optional[Union[int, List[int]]] = None,
                 connection_timeout: float = 240.0):
        super().__init__(
            base_urls=base_urls, hosts=hosts, server_ports=server_ports, connection_timeout=connection_timeout)

        if group_ports is None:
            group_ports = [51216 + i for i in range(self.num_servers)]

        if isinstance(group_ports, int):
            self.group_ports = [group_ports + i for i in range(self.num_servers)]
        elif isinstance(group_ports, list) and len(group_ports) == self.num_servers:
            self.group_ports = group_ports
        else:
            raise ValueError('group_port must be int or list of length num_servers')

        self.pynccl_comms = []

    def init_communicator(self, device: Union[int, str] = 0):
        self.pynccl_comms = []
        for i in range(self.num_servers):
            response = self.sessions[i].get(f'{self.base_urls[i]}/get_world_size/')
            if response.status_code != 200:
                raise Exception(f'Server {i} failed: {response.text}')
            vllm_world_size = response.json()['world_size']

            world_size = vllm_world_size + 1
            rank = vllm_world_size
            # Use '::' for IPv6 hosts, '0.0.0.0' for IPv4 hosts
            bind_host = '::' if is_valid_ipv6_address(self.hosts[i]) else '0.0.0.0'
            response = self.sessions[i].post(
                f'{self.base_urls[i]}/init_communicator/',
                json={
                    'host': bind_host,
                    'port': self.group_ports[i],
                    'world_size': world_size,
                })
            if response.status_code != 200:
                raise Exception(f'Server {i} init failed: {response.text}')

            time.sleep(0.1)

            pg = StatelessProcessGroup.create(
                host=self.hosts[i], port=self.group_ports[i], rank=rank, world_size=world_size)
            if is_vllm_ascend_available():
                import torch_npu
                torch_npu.npu.set_device(device)
            comm = PyNcclCommunicator(pg, device=device)
            self.pynccl_comms.append(comm)

        atexit.register(self.close_communicator)

    def update_named_param(self, name: str, weights: torch.Tensor):
        dtype = str(weights.dtype)
        shape = tuple(weights.shape)

        errors = [None] * self.num_servers

        def _update_single_server(i):
            try:
                response = self.sessions[i].post(
                    f'{self.base_urls[i]}/update_named_param/',
                    json={
                        'name': name,
                        'dtype': dtype,
                        'shape': shape
                    },
                )
                if response.status_code != 200:
                    raise Exception(f'Server {i} update failed: {response.text}')

                synchronize()
                broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], weights, src=self.pynccl_comms[i].rank)
                synchronize()
                self.pynccl_comms[i].group.barrier()
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()

        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors: {all_errors}')

    def update_adapter_flattened_param(self, peft_config, metadatas, flattened_tensor):
        """
        Adds a LoRA adapter to the model on all servers using flattened tensor.

        Args:
            peft_config: PEFT configuration for LoRA adapter.
            metadatas: List of FlattenedTensorMetadata objects.
            flattened_tensor: The flattened tensor containing all adapter parameters.
        """
        errors = [None] * self.num_servers
        peft_config = peft_config_to_dict(peft_config)
        metadatas = [m.model_dump() if hasattr(m, 'model_dump') else m.dict() for m in metadatas]

        def _update_single_server(i):
            try:
                data = {
                    'peft_config': {
                        **peft_config
                    },
                    'metadatas': metadatas,
                }

                response = self.sessions[i].post(
                    f'{self.base_urls[i]}/update_adapter_flattened_param/',
                    json=data,
                )
                if response.status_code != 200:
                    raise Exception(f'Server {i} update adapter failed: {response.text}')

                synchronize()
                broadcast_tensor_for_vllm_weight_sync(
                    self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank)
                synchronize()
                self.pynccl_comms[i].group.barrier()
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()

        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors: {all_errors}')

    def update_adapter_param(self, peft_config, lora_params):
        """
        Adds a LoRA adapter to the model on all servers without flattening.
        Sends each tensor individually.

        Args:
            peft_config: PEFT configuration for LoRA adapter.
            lora_params: OrderedDict of (name, tensor) pairs for LoRA parameters.
        """
        errors = [None] * self.num_servers
        peft_config = peft_config_to_dict(peft_config)

        # Build metadata for each tensor
        lora_tensors_metadata = []
        for name, param in lora_params.items():
            metadata = {
                'name': name,
                'dtype': str(param.dtype),
                'shape': tuple(param.shape),
                'start_idx': 0,  # Not used in non-flattened mode
                'end_idx': param.numel(),  # Not used in non-flattened mode
                'numel': param.numel(),
            }
            lora_tensors_metadata.append(metadata)

        def _update_single_server(i):
            try:
                data = {
                    'peft_config': {
                        **peft_config
                    },
                    'lora_tensors_metadata': lora_tensors_metadata,
                }

                response = self.sessions[i].post(
                    f'{self.base_urls[i]}/update_adapter_param/',
                    json=data,
                )
                if response.status_code != 200:
                    raise Exception(f'Server {i} update adapter failed: {response.text}')

                # Broadcast each tensor individually
                synchronize()
                for name, param in lora_params.items():
                    broadcast_tensor_for_vllm_weight_sync(self.pynccl_comms[i], param, src=self.pynccl_comms[i].rank)
                synchronize()
                self.pynccl_comms[i].group.barrier()
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()

        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors: {all_errors}')

    def update_flattened_params(self, metadatas, flattened_tensor):
        """
        Updates model parameters using flattened tensor data.

        Args:
            metadatas: List of FlattenedTensorMetadata objects
            flattened_tensor: The flattened tensor containing all parameters
        """
        errors = [None] * self.num_servers
        metadatas = [m.model_dump() if hasattr(m, 'model_dump') else m.dict() for m in metadatas]

        def _update_single_server(i):
            try:
                data = {
                    'metadatas': metadatas,
                }

                response = self.sessions[i].post(
                    f'{self.base_urls[i]}/update_flattened_params/',
                    json=data,
                )
                if response.status_code != 200:
                    raise Exception(f'Server {i} update flattened params failed: {response.text}')

                synchronize()
                broadcast_tensor_for_vllm_weight_sync(
                    self.pynccl_comms[i], flattened_tensor, src=self.pynccl_comms[i].rank)
                synchronize()
                self.pynccl_comms[i].group.barrier()
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_update_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()

        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors: {all_errors}')

    def process_weights_after_loading(self):
        """Trigger process_weights_after_loading on all vLLM workers.

        Must be called **once** after ALL weight buckets have been
        sent via ``update_flattened_params``.  This mirrors the
        pattern used by verl and ROLL.
        """
        errors = [None] * self.num_servers

        def _process_single_server(i):
            try:
                response = self.sessions[i].post(f'{self.base_urls[i]}/process_weights_after_loading/', )
                if response.status_code != 200:
                    raise Exception(f'Server {i} process_weights_after_loading failed: {response.text}')
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_process_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()
        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors on process_weights_after_loading: {all_errors}')

    def update_model_params(self, model: nn.Module):
        for name, param in model.named_parameters():
            self.update_named_param(name, param.data)

    def reset_prefix_cache(self):
        errors = [None] * self.num_servers

        def _reset_single_server(i):
            try:
                response = self.sessions[i].post(f'{self.base_urls[i]}/reset_prefix_cache/')
                if response.status_code != 200:
                    raise Exception(f'Server {i} reset failed: {response.text}')
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()
        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors on reset_prefix_cache: {all_errors}')

    def reset_encoder_cache(self):
        errors = [None] * self.num_servers

        def _reset_single_server(i):
            try:
                response = self.sessions[i].post(f'{self.base_urls[i]}/reset_encoder_cache/')
                if response.status_code != 200:
                    raise Exception(f'Server {i} reset failed: {response.text}')
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()
        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors on reset_encoder_cache: {all_errors}')

    def reset_mm_cache(self):
        errors = [None] * self.num_servers

        def _reset_single_server(i):
            try:
                response = self.sessions[i].post(f'{self.base_urls[i]}/reset_mm_cache/')
                if response.status_code != 200:
                    raise Exception(f'Server {i} reset failed: {response.text}')
            except Exception as e:
                errors[i] = e

        with ThreadPoolExecutor(max_workers=self.num_servers) as executor:
            futures = [executor.submit(_reset_single_server, i) for i in range(self.num_servers)]
            for future in futures:
                future.result()
        all_errors = [e for e in errors if e is not None]
        if all_errors:
            raise RuntimeError(f'Multiple errors on reset_mm_cache: {all_errors}')

    def get_engine_type(self):
        # assume that all server has same engine type
        response = self.sessions[0].post(f'{self.base_urls[0]}/get_engine_type/')
        if response.status_code != 200:
            raise Exception(f'Engine type request failed: {response.text}')

        result = response.json()
        self.use_async_engine = result['engine_type'] == 'AsyncLLMEngine'
        self.enable_multi_turn = result.get('enable_multi_turn', False)
        self.use_gym_env = result.get('use_gym_env', False)
        self.enable_lora = result.get('enable_lora', False)
        return result

    def get_model_state_keys(self):
        """Fetch runtime vLLM model parameter names from server."""
        response = self.sessions[0].get(f'{self.base_urls[0]}/get_model_state_keys/')
        if response.status_code != 200:
            raise Exception(f'Get model state keys failed: {response.text}')
        data = response.json()
        return data.get('keys', [])

    def close_communicator(self):
        for i in range(self.num_servers):
            try:
                response = self.sessions[i].post(f'{self.base_urls[i]}/close_communicator/')
                if response.status_code != 200:
                    logger.warning(f'Server {i} close failed: {response.text}')
            except Exception as e:
                logger.warning(f'Error closing server {i} communicator: {str(e)}')
