#! /usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Transformers 后端实现

基于 HuggingFace Transformers 的 LLM/MLLM 后端。
支持通过配置动态选择模型类和处理器类。
"""

from typing import List, Optional, Type

from .base import BaseLLMBackend, ChatMessage, GenerateConfig, ModelConfig


_awq_patched = False


def _patch_awq_compat():
    """修复 autoawq 与新版 transformers 的兼容性问题

    autoawq 已被官方弃用，但 transformers 加载 AWQ 模型仍依赖它。
    新版 transformers (>=4.50) 将 PytorchGELUTanh 重命名为 GELUTanh，
    导致 awq.quantize.scale 导入失败。

    此函数在运行时动态 patch，无需修改 awq 源文件。
    """
    global _awq_patched
    if _awq_patched:
        return

    try:
        # 先 patch transformers.activations，添加别名
        from transformers import activations
        if not hasattr(activations, "PytorchGELUTanh"):
            if hasattr(activations, "GELUTanh"):
                activations.PytorchGELUTanh = activations.GELUTanh
        _awq_patched = True
    except Exception:
        pass


def _get_model_class(class_name: str) -> Type:
    """动态获取模型类

    Args:
        class_name: 类名，如 "AutoModelForCausalLM", "Qwen3VLForConditionalGeneration"

    Returns:
        模型类
    """
    import transformers

    # 首先尝试从 transformers 直接获取
    if hasattr(transformers, class_name):
        return getattr(transformers, class_name)

    # 尝试从 transformers.models 的子模块获取
    # 例如 HunYuanVLForConditionalGeneration
    for module_name in dir(transformers.models):
        try:
            module = getattr(transformers.models, module_name)
            if hasattr(module, class_name):
                return getattr(module, class_name)
        except Exception:
            continue

    raise ValueError(f"无法找到模型类: {class_name}")


def _get_processor_class(class_name: str) -> Type:
    """动态获取处理器类

    Args:
        class_name: 类名，如 "AutoTokenizer", "AutoProcessor"

    Returns:
        处理器类
    """
    import transformers

    if hasattr(transformers, class_name):
        return getattr(transformers, class_name)

    raise ValueError(f"无法找到处理器类: {class_name}")


class TransformersBackend(BaseLLMBackend):
    """基于 Transformers 的 LLM 后端

    支持:
    - 纯文本 LLM (AutoModelForCausalLM)
    - 多模态 VL 模型 (AutoModelForVision2Seq 或其他)
    - 流式输出 (TextIteratorStreamer)
    - 动态模型类和处理器类配置
    - LoRA 适配器动态加载与切换

    配置示例:
        # 使用 HunyuanOCR
        config = ModelConfig(
            model_id="tencent/HunyuanOCR",
            model_class="HunYuanVLForConditionalGeneration",
            processor_class="AutoProcessor",
            vision_processor="general",
        )

        # 使用 Qwen3 带 thinking
        config = ModelConfig(
            model_id="Qwen/Qwen3-0.6B",
            chat_template_kwargs={"enable_thinking": True},
        )

        # 使用 LoRA 适配器
        config = ModelConfig(
            model_id="Qwen/Qwen2.5-7B-Instruct",
            lora_modules={"lora1": "/path/to/lora1", "lora2": "/path/to/lora2"},
        )
    """

    def __init__(self):
        super().__init__()
        self._tokenizer = None
        self._processor = None  # 多模态用
        self._lora_adapters: dict[str, str] = {}  # name -> path

    # ============== 实现抽象方法 ==============

    def _load_model_impl(self, model_path: str, config: ModelConfig) -> None:
        """加载 Transformers 模型"""
        import torch

        # 修复 autoawq 与新版 transformers 的兼容性
        _patch_awq_compat()

        # 确定 dtype
        if config.torch_dtype:
            torch_dtype = getattr(torch, config.torch_dtype, torch.float16)
        else:
            # 自动选择最佳 dtype
            # CUDA: 优先 bfloat16 > float16
            # NPU: 使用 float16 (兼容性最佳)
            # MPS/CPU: 使用 float32 (更稳定)
            if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
                torch_dtype = torch.bfloat16
            elif torch.cuda.is_available():
                torch_dtype = torch.float16
            elif self._device == "npu":
                torch_dtype = torch.float16
            else:
                # MPS 和 CPU 使用 float32 更稳定
                torch_dtype = torch.float32

        if self._is_multimodal:
            self._load_multimodal_model(model_path, config, torch_dtype)
        else:
            self._load_text_model(model_path, config, torch_dtype)

        # 加载 LoRA 适配器
        if config.lora_modules:
            self._load_lora_adapters(config.lora_modules)

    def _generate_impl(
        self, messages: List[ChatMessage], config: GenerateConfig
    ) -> tuple[str, int, int]:
        """Transformers 生成实现"""
        import torch

        # 构建输入
        if self._is_multimodal:
            inputs = self._build_multimodal_inputs(messages)
        else:
            inputs = self._build_text_inputs(messages)

        prompt_tokens = inputs["input_ids"].shape[1]

        # 生成参数
        gen_kwargs = {
            "max_new_tokens": config.max_tokens,
            "temperature": config.temperature if config.temperature > 0 else 1.0,
            "top_p": config.top_p,
            "do_sample": config.temperature > 0,
            "pad_token_id": self._get_pad_token_id(),
        }

        if config.cache_implementation:
            gen_kwargs["cache_implementation"] = config.cache_implementation

        with torch.no_grad():
            outputs = self._model.generate(**inputs, **gen_kwargs)

        # 解码
        new_tokens = outputs[0][prompt_tokens:]
        completion_tokens = len(new_tokens)

        if self._is_multimodal:
            text = self._processor.decode(new_tokens, skip_special_tokens=True)
        else:
            text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)

        return text, prompt_tokens, completion_tokens

    def generate_batch_sync(
        self,
        batch: list[tuple[List[ChatMessage], GenerateConfig]],
    ) -> list[tuple[str, int, int]]:
        """批量生成（静态批处理）

        将多个请求 left-pad 到相同长度后一次 generate，
        比逐个请求串行调用更高效。

        Args:
            batch: [(messages, config), ...] 请求列表

        Returns:
            [(text, prompt_tokens, completion_tokens), ...] 结果列表
        """
        import torch

        if not batch:
            return []
        if len(batch) == 1:
            msgs, cfg = batch[0]
            processed = self._process_messages(msgs)
            return [self._generate_impl(processed, cfg)]

        # 仅支持纯文本批处理
        if self._is_multimodal:
            results = []
            for msgs, cfg in batch:
                processed = self._process_messages(msgs)
                results.append(self._generate_impl(processed, cfg))
            return results

        tokenizer = self._tokenizer

        # 1. 将所有消息转为文本
        texts = []
        for msgs, cfg in batch:
            processed = self._process_messages(msgs)
            formatted = []
            for msg in processed:
                content = msg.content if isinstance(msg.content, str) else " ".join(
                    p.text for p in msg.content if p.type == "text" and p.text
                )
                formatted.append({"role": msg.role, "content": content})
            chat_kwargs = {"tokenize": False, "add_generation_prompt": True}
            if self._config and self._config.chat_template_kwargs:
                chat_kwargs.update(self._config.chat_template_kwargs)
            texts.append(tokenizer.apply_chat_template(formatted, **chat_kwargs))

        # 2. Left-pad tokenize
        original_padding_side = tokenizer.padding_side
        tokenizer.padding_side = "left"
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.eos_token_id

        inputs = tokenizer(texts, return_tensors="pt", padding=True)
        inputs = inputs.to(self._device)
        tokenizer.padding_side = original_padding_side

        padded_len = inputs["input_ids"].shape[1]

        # 每个序列的实际 prompt 长度
        prompt_lengths = inputs["attention_mask"].sum(dim=1).tolist()

        # 3. 取 batch 中最大的 max_tokens
        max_tokens = max(cfg.max_tokens for _, cfg in batch)

        # 生成参数（使用第一个请求的采样参数）
        _, first_cfg = batch[0]
        gen_kwargs = {
            "max_new_tokens": max_tokens,
            "temperature": first_cfg.temperature if first_cfg.temperature > 0 else 1.0,
            "top_p": first_cfg.top_p,
            "do_sample": first_cfg.temperature > 0,
            "pad_token_id": self._get_pad_token_id(),
        }

        if first_cfg.cache_implementation:
            gen_kwargs["cache_implementation"] = first_cfg.cache_implementation

        # 4. 批量 generate
        with torch.no_grad():
            outputs = self._model.generate(**inputs, **gen_kwargs)

        # 5. 拆分结果
        results = []
        for i, (msgs, cfg) in enumerate(batch):
            new_tokens = outputs[i][padded_len:]
            # 截断到该请求的 max_tokens
            if len(new_tokens) > cfg.max_tokens:
                new_tokens = new_tokens[:cfg.max_tokens]
            completion_tokens = len(new_tokens)
            text = tokenizer.decode(new_tokens, skip_special_tokens=True)
            results.append((text, int(prompt_lengths[i]), completion_tokens))

        return results

    def _generate_stream_impl(
        self, messages: List[ChatMessage], config: GenerateConfig
    ):
        """Transformers 流式生成实现"""
        from threading import Thread
        from transformers import TextIteratorStreamer

        # 构建输入
        if self._is_multimodal:
            inputs = self._build_multimodal_inputs(messages)
            tokenizer = self._processor.tokenizer
        else:
            inputs = self._build_text_inputs(messages)
            tokenizer = self._tokenizer

        # 创建 streamer
        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
        )

        gen_kwargs = {
            **inputs,
            "max_new_tokens": config.max_tokens,
            "temperature": config.temperature if config.temperature > 0 else 1.0,
            "top_p": config.top_p,
            "do_sample": config.temperature > 0,
            "pad_token_id": self._get_pad_token_id(),
            "streamer": streamer,
        }

        if config.cache_implementation:
            gen_kwargs["cache_implementation"] = config.cache_implementation

        # 在线程中运行生成
        thread = Thread(target=self._model.generate, kwargs=gen_kwargs)
        thread.start()

        # 流式输出
        for text in streamer:
            yield text

        thread.join()

    # ============== 内部方法 ==============

    def _load_text_model(self, model_path: str, config: ModelConfig, torch_dtype) -> None:
        """加载纯文本模型"""
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        # 确定处理器类
        processor_class_name = config.processor_class or "AutoTokenizer"
        ProcessorClass = _get_processor_class(processor_class_name)

        self._tokenizer = ProcessorClass.from_pretrained(
            model_path, trust_remote_code=config.trust_remote_code
        )

        # 确定模型类
        model_class_name = config.model_class or "AutoModelForCausalLM"
        ModelClass = _get_model_class(model_class_name)

        model_kwargs = {
            "torch_dtype": torch_dtype,
            "device_map": self._device,
            "trust_remote_code": config.trust_remote_code,
        }
        if config.attn_implementation:
            model_kwargs["attn_implementation"] = config.attn_implementation

        self._model = ModelClass.from_pretrained(model_path, **model_kwargs)
        self._model.eval()

        if config.compile:
            from loguru import logger
            logger.info(f"Compiling model with mode={config.compile_mode}")
            self._model = torch.compile(self._model, mode=config.compile_mode)

    def _load_multimodal_model(self, model_path: str, config: ModelConfig, torch_dtype) -> None:
        """加载多模态模型"""
        import torch
        from transformers import AutoProcessor, AutoModelForVision2Seq

        # 确定处理器类
        processor_class_name = config.processor_class or "AutoProcessor"
        ProcessorClass = _get_processor_class(processor_class_name)

        processor_kwargs = {"trust_remote_code": config.trust_remote_code}
        # HunyuanOCR 需要 use_fast=False
        if "hunyuan" in model_path.lower():
            processor_kwargs["use_fast"] = False

        self._processor = ProcessorClass.from_pretrained(model_path, **processor_kwargs)

        # 确定模型类
        model_class_name = config.model_class or "AutoModelForVision2Seq"
        ModelClass = _get_model_class(model_class_name)

        model_kwargs = {
            "torch_dtype": torch_dtype,
            "device_map": self._device,
            "trust_remote_code": config.trust_remote_code,
        }
        if config.attn_implementation:
            model_kwargs["attn_implementation"] = config.attn_implementation

        self._model = ModelClass.from_pretrained(model_path, **model_kwargs)
        self._model.eval()

        if config.compile:
            from loguru import logger
            logger.info(f"Compiling model with mode={config.compile_mode}")
            self._model = torch.compile(self._model, mode=config.compile_mode)

    def _build_text_inputs(self, messages: List[ChatMessage]):
        """构建纯文本输入"""
        # 转换为标准格式
        formatted = []
        for msg in messages:
            content = msg.content if isinstance(msg.content, str) else " ".join(
                p.text for p in msg.content if p.type == "text" and p.text
            )
            formatted.append({"role": msg.role, "content": content})

        # 获取 chat_template_kwargs
        chat_kwargs = {"tokenize": False, "add_generation_prompt": True}
        if self._config and self._config.chat_template_kwargs:
            chat_kwargs.update(self._config.chat_template_kwargs)

        text = self._tokenizer.apply_chat_template(formatted, **chat_kwargs)
        inputs = self._tokenizer(text, return_tensors="pt")
        return inputs.to(self._device)

    def _build_multimodal_inputs(self, messages: List[ChatMessage]):
        """构建多模态输入

        根据 vision_processor 配置选择不同的处理方式：
        - qwen_vl: 使用 qwen_vl_utils.process_vision_info (Qwen-VL 系列)
        - general: 通用处理方式 (HunyuanOCR, dots.ocr 等)
        """
        if self._vision_processor == "qwen_vl":
            return self._build_qwen_vl_inputs(messages)
        else:
            return self._build_general_vl_inputs(messages)

    def _build_qwen_vl_inputs(self, messages: List[ChatMessage]):
        """构建 Qwen-VL 风格的输入"""
        from qwen_vl_utils import process_vision_info

        # 转换为 Qwen-VL 格式
        qwen_messages = []
        for msg in messages:
            if isinstance(msg.content, str):
                qwen_messages.append({"role": msg.role, "content": msg.content})
            else:
                content_parts = []
                for part in msg.content:
                    if part.type == "text":
                        content_parts.append({"type": "text", "text": part.text})
                    elif part.type == "image_url" and part.image_url:
                        image_url = part.image_url.url
                        if image_url.startswith("data:"):
                            image = self._process_image(image_url)
                            content_parts.append({"type": "image", "image": image})
                        else:
                            content_parts.append({"type": "image", "image": image_url})
                qwen_messages.append({"role": msg.role, "content": content_parts})

        # 获取 chat_template_kwargs
        chat_kwargs = {"tokenize": False, "add_generation_prompt": True}
        if self._config and self._config.chat_template_kwargs:
            chat_kwargs.update(self._config.chat_template_kwargs)

        # 使用 processor 处理
        text = self._processor.apply_chat_template(qwen_messages, **chat_kwargs)
        image_inputs, video_inputs = process_vision_info(qwen_messages)

        inputs = self._processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        return inputs.to(self._device)

    def _build_general_vl_inputs(self, messages: List[ChatMessage]):
        """构建通用多模态输入 (适用于大多数 VL 模型)"""
        # 提取图片和文本
        images = []
        formatted_messages = []

        for msg in messages:
            if isinstance(msg.content, str):
                formatted_messages.append({"role": msg.role, "content": msg.content})
            else:
                content_parts = []
                for part in msg.content:
                    if part.type == "text" and part.text:
                        content_parts.append({"type": "text", "text": part.text})
                    elif part.type == "image_url" and part.image_url:
                        image = self._process_image(part.image_url.url)
                        images.append(image)
                        content_parts.append({"type": "image"})

                formatted_messages.append({"role": msg.role, "content": content_parts})

        # 获取 chat_template_kwargs
        chat_kwargs = {"tokenize": False, "add_generation_prompt": True}
        if self._config and self._config.chat_template_kwargs:
            chat_kwargs.update(self._config.chat_template_kwargs)

        # 应用 chat template
        text = self._processor.apply_chat_template(formatted_messages, **chat_kwargs)

        # 处理输入
        if images:
            inputs = self._processor(
                text=[text],
                images=images,
                padding=True,
                return_tensors="pt",
            )
        else:
            inputs = self._processor(
                text=[text],
                padding=True,
                return_tensors="pt",
            )

        return inputs.to(self._device)

    def _get_pad_token_id(self) -> int:
        """获取 pad token id"""
        if self._is_multimodal:
            tokenizer = self._processor.tokenizer
        else:
            tokenizer = self._tokenizer

        if tokenizer.pad_token_id is not None:
            return tokenizer.pad_token_id
        return tokenizer.eos_token_id

    def _load_lora_adapters(self, lora_modules: dict[str, str]) -> None:
        """加载 LoRA 适配器

        Args:
            lora_modules: {adapter_name: adapter_path}
        """
        from peft import PeftModel
        from loguru import logger

        for name, path in lora_modules.items():
            logger.info(f"Loading LoRA adapter: {name} from {path}")
            if not self._lora_adapters:
                # 第一个适配器：用 PeftModel 包装
                self._model = PeftModel.from_pretrained(
                    self._model, path, adapter_name=name
                )
            else:
                self._model.load_adapter(path, adapter_name=name)
            self._lora_adapters[name] = path
            logger.info(f"LoRA adapter loaded: {name}")

        # 加载完后禁用所有适配器，默认使用基座模型
        self._model.disable_adapter_layers()

    def set_active_lora(self, adapter_name: Optional[str]) -> None:
        """切换 LoRA 适配器

        Args:
            adapter_name: 适配器名称，None 表示使用基座模型
        """
        if not self._lora_adapters:
            return

        if adapter_name is None:
            self._model.disable_adapter_layers()
        else:
            self._model.enable_adapter_layers()
            self._model.set_adapter(adapter_name)


# 默认后端别名
LLMBackend = TransformersBackend
