# transformers_api.py
"""Utilities for loading & running local HuggingFace Transformer models in ComfyUI LLM-Toolkit.

The helper focuses on *offline* (local) inference and therefore does not require any Internet
connectivity once the model weights have been downloaded.  It provides a thin wrapper that mirrors
our other provider_* modules (openai_api, gemini_api, ollama_api …) so that *send_request* can call
it with a familiar signature.

Key design goals
----------------
• *Precision* – Load models in half-precision (float16 / bfloat16) by default to minimise VRAM
  footprint.  Users may override via the *precision* kwarg.
• *AWQ* – If the model repo or directory name contains the substring "AWQ" we will attempt to load
  it with the *autoawq* loader (if available).  Otherwise we fallback to standard HF loading.
• *Multi-modal* – Basic support for Alibaba Qwen2.5-Omni (text-image-audio) and Qwen-VL/VL-like
  models is included.  The caller can pass *base64_images* (List[str]) and *base64_audio* (List[str])
  which will be converted to PIL images / raw audio arrays via *qwen_omni_utils* when available.
• *Model cache* – Keep models & tokenisers in a global dict so subsequent calls are instant.
• *Thread safety* – The cache is guarded by a simple asyncio lock.

The public coroutine *send_transformers_request* follows the same signature as other providers so
that integration inside *send_request.py* is trivial.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple
import asyncio
import base64
import logging
import os
from pathlib import Path
import textwrap

import torch

# Lazy loading for transformers - don't import at module level
TRANSFORMERS_AVAILABLE = None
AutoModelForCausalLM = None
AutoTokenizer = None
TextIteratorStreamer = None
Qwen2_5OmniForConditionalGeneration = None
Qwen2_5OmniProcessor = None

def _ensure_transformers():
    """Lazy-load transformers only when actually needed."""
    global TRANSFORMERS_AVAILABLE, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
    global Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniProcessor
    
    if TRANSFORMERS_AVAILABLE is not None:
        return TRANSFORMERS_AVAILABLE
    
    try:
        from transformers import (
            AutoModelForCausalLM as _AutoModelForCausalLM,
            AutoTokenizer as _AutoTokenizer,
            TextIteratorStreamer as _TextIteratorStreamer,
        )
        AutoModelForCausalLM = _AutoModelForCausalLM
        AutoTokenizer = _AutoTokenizer
        TextIteratorStreamer = _TextIteratorStreamer
        
        # Optional – might not exist in every transformers release
        try:
            from transformers import (
                Qwen2_5OmniForConditionalGeneration as _Qwen2_5OmniForConditionalGeneration,
                Qwen2_5OmniProcessor as _Qwen2_5OmniProcessor,
            )
            Qwen2_5OmniForConditionalGeneration = _Qwen2_5OmniForConditionalGeneration
            Qwen2_5OmniProcessor = _Qwen2_5OmniProcessor
        except (ImportError, AttributeError):
            pass
        
        TRANSFORMERS_AVAILABLE = True
        logger.info("Transformers library loaded successfully")
    except Exception as e:
        TRANSFORMERS_AVAILABLE = False
        logger.warning(f"Transformers library not available: {e}")
    
    return TRANSFORMERS_AVAILABLE

logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
#  Model cache (global across requests)
# -----------------------------------------------------------------------------
_MODEL_CACHE: Dict[str, Dict[str, Any]] = {}
_CACHE_LOCK = asyncio.Lock()


# -----------------------------------------------------------------------------
#  Helper – determine torch dtype from precision string
# -----------------------------------------------------------------------------

_PRECISION_MAP = {
    "fp16": torch.float16,
    "float16": torch.float16,
    "bf16": torch.bfloat16,
    "bfloat16": torch.bfloat16,
    "fp32": torch.float32,
    "float32": torch.float32,
}

def _to_dtype(precision: str | None) -> torch.dtype:
    if not precision:
        return torch.float16
    precision = precision.lower()
    return _PRECISION_MAP.get(precision, torch.float16)


# -----------------------------------------------------------------------------
#  Helper – decide which device to place the model on
# -----------------------------------------------------------------------------

def _best_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():  # macOS Metal
        return "mps"
    return "cpu"


# -----------------------------------------------------------------------------
#  Loader – main entry used by *send_transformers_request*
# -----------------------------------------------------------------------------

def _is_awq_model(model_name: str) -> bool:
    return "awq" in model_name.lower()

async def _ensure_model_loaded(
    model_name_or_path: str,
    *,
    precision: str = "fp16",
    device: str | None = None,
    trust_remote_code: bool = True,
) -> Dict[str, Any]:
    """Load *model_name_or_path* into memory once and cache it.

    Returns a dict with keys *model*, *tokenizer*, and optionally *processor* (for Qwen-Omni).
    Calls are thread-safe (asyncio).
    """
    if not _ensure_transformers():
        raise RuntimeError(
            "transformers_api – HuggingFace transformers not installed.  Please `pip install transformers`"
        )

    async with _CACHE_LOCK:
        if model_name_or_path in _MODEL_CACHE:
            return _MODEL_CACHE[model_name_or_path]

        device = device or _best_device()
        dtype = _to_dtype(precision)

        logger.info(
            "[Transformers-API] Loading model %s | device=%s | dtype=%s",
            model_name_or_path,
            device,
            dtype,
        )

        # ------------------------------------------------------------------
        #  Special-case Qwen2.5-Omni multimodal
        # ------------------------------------------------------------------
        if (
            "qwen2.5-omni" in model_name_or_path.lower()
            and Qwen2_5OmniForConditionalGeneration is not None
            and Qwen2_5OmniProcessor is not None
        ):
            try:
                model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
                    model_name_or_path,
                    torch_dtype=dtype,
                    device_map="auto" if device == "cuda" else None,  # CPU / MPS handled by HF
                    trust_remote_code=trust_remote_code,
                )
                processor = Qwen2_5OmniProcessor.from_pretrained(model_name_or_path)
                tokenizer = None  # Processor handles tokenisation
                _MODEL_CACHE[model_name_or_path] = {
                    "model": model,
                    "tokenizer": tokenizer,
                    "processor": processor,
                    "modal": "omni",
                }
                return _MODEL_CACHE[model_name_or_path]
            except Exception as exc:
                logger.error("Failed to load Qwen2.5-Omni model %s – %s", model_name_or_path, exc)
                # Fall through to generic loader below instead of raising

        # If the user requested an Omni model but the specialised classes are unavailable, log a hint.
        if "qwen2.5-omni" in model_name_or_path.lower() and Qwen2_5OmniForConditionalGeneration is None:
            logger.warning(
                "Qwen2.5-Omni support not available in your transformers build; "
                "loading as a plain text model.  Install a newer transformers pre-release "
                "to enable multimodal features."
            )

        # ------------------------------------------------------------------
        #  AWQ (AutoAWQ) support – try specialised loader first
        # ------------------------------------------------------------------
        if _is_awq_model(model_name_or_path):
            try:
                from awq import AutoAWQForCausalLM  # type: ignore

                # NB: Using `device_map="auto"` sometimes places the token embedding
                # layer on CPU while the rest of the model (and the input IDs) reside
                # on GPU, which then crashes inside `torch.embedding`.  We therefore
                # load without device mapping first and *explicitly* move the entire
                # model to the requested device afterwards.

                # Prefer the dedicated *from_quantized* helper when available because it sets up
                # quantised layers on the correct device automatically.  Fall back to the plain
                # *from_pretrained* call with `device_map='auto'` otherwise.

                try:
                    model = AutoAWQForCausalLM.from_quantized(  # type: ignore[attr-defined]
                        model_name_or_path,
                        torch_dtype=dtype,
                        device_map="auto" if device == "cuda" else None,
                        trust_remote_code=trust_remote_code,
                    )
                except AttributeError:
                    # Older awq versions – fall back to from_pretrained
                    model = AutoAWQForCausalLM.from_pretrained(
                        model_name_or_path,
                        torch_dtype=dtype,
                        device_map="auto" if device == "cuda" else None,
                        trust_remote_code=trust_remote_code,
                    )

                # Sanity-check: ensure the token embedding sits on the same device as the rest
                try:
                    emb = model.get_input_embeddings()
                    if emb is not None and emb.weight.device != model.device:
                        emb.to(model.device)
                except Exception:
                    pass  # defensive – skip if anything unexpected
                tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
                _MODEL_CACHE[model_name_or_path] = {
                    "model": model,
                    "tokenizer": tokenizer,
                    "modal": "text",
                }
                return _MODEL_CACHE[model_name_or_path]
            except Exception as exc:
                logger.warning(
                    "AWQ specific loader failed for %s (continuing with normal HF loader): %s",
                    model_name_or_path,
                    exc,
                )
                # Fall through to normal loader

        # ------------------------------------------------------------------
        #  Generic text models (AutoModel)
        # ------------------------------------------------------------------
        try:
            model = AutoModelForCausalLM.from_pretrained(
                model_name_or_path,
                torch_dtype=dtype,
                device_map="auto" if device == "cuda" else None,
                trust_remote_code=trust_remote_code,
            )
            tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=trust_remote_code)
            _MODEL_CACHE[model_name_or_path] = {
                "model": model,
                "tokenizer": tokenizer,
                "modal": "text",
            }
            return _MODEL_CACHE[model_name_or_path]
        except Exception as exc:
            logger.error("Could not load model %s via transformers – %s", model_name_or_path, exc)
            raise


# -----------------------------------------------------------------------------
#  Utility – multimodal helpers (images / audio) for Qwen-Omni
# -----------------------------------------------------------------------------

try:
    from qwen_omni_utils import process_mm_info  # provided by qwen-omni-utils package
except Exception:  # pragma: no cover – optional dependency
    def process_mm_info(*_args, **_kwargs):  # type: ignore
        raise NotImplementedError(
            "qwen_omni_utils not installed – required for image/audio inputs with Qwen-Omni"
        )


# -----------------------------------------------------------------------------
#  Public coroutine – mirror *send_openai_request* API
# -----------------------------------------------------------------------------
async def send_transformers_request(
    *,
    base64_images: List[str],
    base64_audio: List[str],
    model: str,
    system_message: str,
    user_message: str,
    messages: List[Dict[str, Any]] | None,
    temperature: float = 0.7,
    max_tokens: int = 1024,
    top_p: float = 0.9,
    repeat_penalty: float = 1.1,
    precision: str = "fp16",
    **kwargs,
) -> Dict[str, Any]:
    """Run local inference with a HF transformer checkpoint.

    Returns an *OpenAI-style* dict so the rest of the toolkit can remain provider-agnostic.
    """
    if not _ensure_transformers():
        return {
            "choices": [
                {"message": {"content": "Transformers not installed.  Please `pip install transformers`."}}]
        }

    try:
        artefacts = await _ensure_model_loaded(model, precision=precision)
        mdl = artefacts["model"]
        tokenizer = artefacts.get("tokenizer")
        processor = artefacts.get("processor")
        modal_type = artefacts.get("modal", "text")

        # Determine the primary device of the model by inspecting the first parameter
        try:
            first_param = next(mdl.parameters())
            device = first_param.device
        except StopIteration:
            device = torch.device("cpu")
        mdl.eval()

        # Build the conversation list for chat-style models
        if messages is None:
            messages = []
        if system_message:
            messages = [{"role": "system", "content": system_message}] + messages
        if user_message:
            messages = messages + [{"role": "user", "content": user_message}]

        # ------------------------------------------------------------------
        #  Qwen2.5-Omni – multimodal path
        # ------------------------------------------------------------------
        if modal_type == "omni":
            if processor is None:
                raise RuntimeError("Qwen-Omni processor not initialised")

            # Convert base64 images / audio to the expected format via helper
            audios, images, videos = process_mm_info(
                messages,  # uses the conversation to pull URLs but we feed empty lists next
                use_audio_in_video=False,
            )
            # Replace the image/audio placeholders explicitly – images list expects raw bytes or PIL
            # For simplicity we treat base64_images as url-style – processor can handle raw bytes too.
            for b64 in base64_images or []:
                images.append({"type": "image", "image": f"data:image/png;base64,{b64}"})
            for b64 in base64_audio or []:
                audios.append({"type": "audio", "audio": base64.b64decode(b64)})

            text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
            inputs = processor(text=text_prompt, audio=audios, images=images, videos=videos, return_tensors="pt", padding=True)
            inputs = inputs.to(device)

            out_ids, out_audio = mdl.generate(
                **inputs,
                do_sample=True,
                temperature=temperature,
                max_new_tokens=max_tokens,
                top_p=top_p,
                repetition_penalty=repeat_penalty,
                return_audio=False,  # We'll stick to text for now
            )
            output_texts = processor.batch_decode(out_ids, skip_special_tokens=True)
            content = output_texts[0]
        else:
            # ------------------------------------------------------------------
            #  Generic text-only generation
            # ------------------------------------------------------------------
            if tokenizer is None:
                raise RuntimeError("Tokenizer missing for text model")

            # Prefer the model-specific chat template when the tokenizer provides one.
            if hasattr(tokenizer, "apply_chat_template"):
                model_name_lower = model.lower()
                # Kimi-VL and Qwen3 models need `enable_thinking=True` for system prompts
                if "qwen3" in model_name_lower or ("kimi-vl" in model_name_lower and "thinking" in model_name_lower):
                    try:
                        prompt_text = tokenizer.apply_chat_template(
                            messages,
                            tokenize=False,
                            add_generation_prompt=True,
                            enable_thinking=True,
                        )
                    except TypeError:
                        # Fallback for older tokenizer versions without the kwarg
                        prompt_text = tokenizer.apply_chat_template(
                            messages,
                            tokenize=False,
                            add_generation_prompt=True,
                        )
                else:
                    prompt_text = tokenizer.apply_chat_template(
                        messages,
                        tokenize=False,
                        add_generation_prompt=True,
                    )
            else:
                # Fallback to simple concatenation for generic models
                prompt_parts: List[str] = []
                for msg in messages:
                    if msg["role"] == "system":
                        prompt_parts.append(f"[INST] {msg['content']} [/INST]")
                    elif msg["role"] == "user":
                        prompt_parts.append(f"<|user|>: {msg['content']}")
                    elif msg["role"] == "assistant":
                        prompt_parts.append(f"<|assistant|>: {msg['content']}")
                prompt_text = "\n".join(prompt_parts) + "\n<|assistant|>:"

            input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids.to(device)
            with torch.no_grad():
                gen_ids = mdl.generate(
                    input_ids,
                    do_sample=True,
                    temperature=temperature,
                    max_new_tokens=max_tokens,
                    top_p=top_p,
                    repetition_penalty=repeat_penalty,
                )
            
            output_ids = gen_ids[0][input_ids.shape[-1] :]
            model_name_lower = model.lower()

            if "qwen3" in model_name_lower:
                try:
                    eot_token_id = 151668  # Token ID for </think> in Qwen3
                    output_ids_list = output_ids.tolist()
                    index = len(output_ids_list) - output_ids_list[::-1].index(eot_token_id)
                    content_ids = output_ids_list[index:]
                    content = tokenizer.decode(content_ids, skip_special_tokens=True).strip()
                except ValueError:
                    content = tokenizer.decode(output_ids, skip_special_tokens=True).strip()
            
            elif "kimi-vl" in model_name_lower and "thinking" in model_name_lower:
                bot_token = "◁think▷"
                eot_token = "◁/think▷"
                decoded_text = tokenizer.decode(output_ids, skip_special_tokens=True)
                
                if bot_token in decoded_text and eot_token in decoded_text:
                    _, summary = decoded_text.split(eot_token, 1)
                    content = summary.strip()
                else:
                    content = decoded_text.strip()
            else:
                content = tokenizer.decode(output_ids, skip_special_tokens=True)

        return {
            "choices": [
                {
                    "message": {
                        "content": content,
                    }
                }
            ]
        }
    except Exception as exc:
        logger.error("transformers_api – inference error: %s", exc, exc_info=True)
        return {"choices": [{"message": {"content": f"Error: {exc}"}}]}


# -----------------------------------------------------------------------------
#  Public coroutine – streaming version (yields chunks)
# -----------------------------------------------------------------------------
async def send_transformers_request_stream(
    *,
    base64_images: List[str],
    base64_audio: List[str],
    model: str,
    system_message: str,
    user_message: str,
    messages: List[Dict[str, Any]] | None,
    temperature: float = 0.7,
    max_tokens: int = 1024,
    top_p: float = 0.9,
    repeat_penalty: float = 1.1,
    precision: str = "fp16",
    **kwargs,
):
    """Async generator that yields text chunks as they are generated."""
    if not _ensure_transformers():
        yield "[Transformers not installed]"
        return

    try:
        artefacts = await _ensure_model_loaded(model, precision=precision)
        mdl = artefacts["model"]
        tokenizer = artefacts.get("tokenizer")
        modal_type = artefacts.get("modal", "text")

        if modal_type != "text":
            # For now fall back to full generation for multimodal models.
            full_resp = await send_transformers_request(
                base64_images=base64_images,
                base64_audio=base64_audio,
                model=model,
                system_message=system_message,
                user_message=user_message,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                repeat_penalty=repeat_penalty,
                precision=precision,
            )
            content = full_resp.get("choices", [{}])[0].get("message", {}).get("content", "")
            if content:
                for _chunk in textwrap.wrap(content, 120):
                    yield _chunk
            return

        if tokenizer is None:
            raise RuntimeError("Tokenizer missing for text model (streaming)")

        try:
            first_param = next(mdl.parameters())
            device = first_param.device
        except StopIteration:
            device = torch.device("cpu")
        mdl.eval()

        # Build prompt
        if messages is None:
            messages = []
        if system_message:
            messages = [{"role": "system", "content": system_message}] + messages
        if user_message:
            messages = messages + [{"role": "user", "content": user_message}]

        # Use model-specific chat template if available to prevent token leakage
        if hasattr(tokenizer, "apply_chat_template"):
            model_name_lower = model.lower()
            # Kimi-VL and Qwen3 models need `enable_thinking=True` for system prompts
            if "qwen3" in model_name_lower or ("kimi-vl" in model_name_lower and "thinking" in model_name_lower):
                try:
                    prompt = tokenizer.apply_chat_template(
                        messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
                    )
                except (TypeError, ValueError):
                    prompt = tokenizer.apply_chat_template(
                        messages, tokenize=False, add_generation_prompt=True
                    )
            else:
                prompt = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
        else:
            # Legacy fallback for models without a chat template
            prompt_parts: List[str] = []
            for msg in messages:
                if msg["role"] == "system":
                    prompt_parts.append(f"[INST] {msg['content']} [/INST]")
                elif msg["role"] == "user":
                    prompt_parts.append(f"<|user|>: {msg['content']}")
                elif msg["role"] == "assistant":
                    prompt_parts.append(f"<|assistant|>: {msg['content']}")
            prompt = "\n".join(prompt_parts) + "\n<|assistant|>:"

        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

        from threading import Thread
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

        generation_kwargs = dict(
            input_ids=input_ids,
            streamer=streamer,
            do_sample=True,
            temperature=temperature,
            max_new_tokens=max_tokens,
            top_p=top_p,
            repetition_penalty=repeat_penalty,
        )

        # Run generate in background thread so we can iterate streamer concurrently
        gen_thread = Thread(target=mdl.generate, kwargs=generation_kwargs, daemon=True)
        gen_thread.start()

        loop = asyncio.get_running_loop()

        stop = False
        while not stop:
            def _next_chunk():
                try:
                    return next(streamer)
                except StopIteration:
                    return None

            chunk = await loop.run_in_executor(None, _next_chunk)
            if chunk is None:
                stop = True
            elif chunk:
                yield chunk

        gen_thread.join(timeout=0)
    except Exception as exc:
        logger.warning(
            "transformers_api – streaming failed (%s). Falling back to blocking generation.",
            exc,
        )
        try:
            full_resp = await send_transformers_request(
                base64_images=base64_images,
                base64_audio=base64_audio,
                model=model,
                system_message=system_message,
                user_message=user_message,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                top_p=top_p,
                repeat_penalty=repeat_penalty,
                precision=precision,
            )
            content = full_resp.get("choices", [{}])[0].get("message", {}).get("content", "")
            if content:
                for _chunk in textwrap.wrap(content, 120):
                    yield _chunk
        except Exception as e2:
            logger.error("transformers_api – fallback blocking generation also failed: %s", e2, exc_info=True)
            yield f"[Error: {e2}]"