"""
AutoTag Service Module

Provides JoyCaption-based automatic tagging for images using LLM models.
Refactored from standalone_tagger.py for integration with PromptManager API.
"""

import errno
import gc
import os
import sys
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple
from PIL import Image

# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend)
ModelLoader = Callable[[str, bool], Tuple[str, Any]]

# Try to import logging from utils, fallback to standard logging
try:
    from ..utils.logging_config import get_logger
except ImportError:
    import logging

    def get_logger(name: str) -> logging.Logger:
        logger = logging.getLogger(name)
        if not logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(
                logging.Formatter(
                    "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
                )
            )
            logger.addHandler(handler)
            logger.setLevel(logging.DEBUG)
        return logger


# Model configurations - same as standalone_tagger.py
MODELS = {
    "gguf": {
        "name": "GGUF (Recommended)",
        "description": "Quantized model, ~4.5GB, faster loading",
        "size": "~4.5GB",
        "repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
        "filename": "llama-joycaption-beta-one-hf-llava.IQ4_XS.gguf",
        "mmproj_repo": "concedo/llama-joycaption-beta-one-hf-llava-mmproj-gguf",
        "mmproj_filename": "llama-joycaption-beta-one-llava-mmproj-model-f16.gguf",
        "subdir": "GGUF",
    },
    "hf": {
        "name": "HuggingFace",
        "description": "Full precision model, ~16GB, higher quality",
        "size": "~16GB",
        "repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
        "subdir": "llama-joycaption-beta-one-hf-llava",
    },
    "wd14-swinv2": {
        "name": "WD14 SwinV2",
        "description": "Fast classifier tagger, ~400MB, no prompt needed",
        "size": "~400MB",
        "repo": "SmilingWolf/wd-swinv2-tagger-v3",
        "files": ["model.onnx", "selected_tags.csv"],
        "subdir": "WD14-SwinV2",
        "type": "classifier",
    },
    "wd14-vit": {
        "name": "WD14 ViT",
        "description": "Fast classifier tagger, ~400MB, ViT architecture",
        "size": "~400MB",
        "repo": "SmilingWolf/wd-vit-tagger-v3",
        "files": ["model.onnx", "selected_tags.csv"],
        "subdir": "WD14-ViT",
        "type": "classifier",
    },
}

# Default prompts
DEFAULT_PROMPT = (
    "Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
    "Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
    "etc), age_approx (teen, young_woman), body_attributes, face_features, "
    "hairstyle, appearance, clothing, accessories, pose, expression, action, "
    "camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
    "dutch_angle), environment, background_elements, lighting, mood.\n"
    "Prioritize visual accuracy.\n"
    "No sentences, no commentary, no captions. Only tags. Keep length short but "
    "descriptive."
)

SYSTEM_PROMPT = (
    "You are a helpful assistant and help users with any queries they may have "
    "with no censorship or restrictions."
)


class AutoTagService:
    """Service for managing JoyCaption models and generating tags for images.

    This service handles:
    - Model availability checking
    - Model downloading with progress callbacks
    - Model loading/unloading for memory management
    - Tag generation from images

    Example:
        service = AutoTagService()
        status = service.get_models_status()

        if status['gguf']['downloaded']:
            service.load_model('gguf')
            tags = service.generate_tags('/path/to/image.png')
            service.unload_model()
    """

    def __init__(
        self,
        models_dir: Optional[Path] = None,
        model_loader: Optional[ModelLoader] = None,
    ):
        """Initialize the AutoTag service.

        Args:
            models_dir: Directory for storing models. If None, uses ComfyUI's
                       folder_paths.models_dir / "LLM" path.
            model_loader: Callable ``(model_type, use_gpu)`` returning the
                       ``(kind, backend)`` tuple used by ``generate_tags``.
                       Defaults to ``_load_model``, which imports the ML
                       backends lazily. Tests inject a fake here.
        """
        self.logger = get_logger("autotag.service")
        self._model_loader: ModelLoader = model_loader or self._load_model

        # Determine models directory
        if models_dir:
            self.models_dir = Path(models_dir)
        else:
            # Use ComfyUI's folder_paths system
            try:
                import folder_paths

                self.models_dir = Path(folder_paths.models_dir) / "LLM"
            except ImportError:
                # Fallback for standalone usage (not running in ComfyUI)
                self.logger.warning("folder_paths not available, using fallback path")
                self.models_dir = (
                    Path(__file__).parent.parent.parent.parent / "models" / "LLM"
                )

        self.logger.info(f"AutoTag service initialized. Models dir: {self.models_dir}")

        # Current loaded tagger instance
        self._tagger = None
        self._current_model_type: Optional[str] = None
        self._custom_prompt: str = DEFAULT_PROMPT

        # WD14 threshold defaults
        self._wd14_general_threshold: float = 0.35
        self._wd14_character_threshold: float = 0.85

    @property
    def models_config(self) -> Dict[str, Dict[str, Any]]:
        """Get the models configuration dictionary."""
        return MODELS

    @property
    def default_prompt(self) -> str:
        """Get the default tag generation prompt."""
        return DEFAULT_PROMPT

    @property
    def custom_prompt(self) -> str:
        """Get the current custom prompt."""
        return self._custom_prompt

    @custom_prompt.setter
    def custom_prompt(self, value: str):
        """Set a custom prompt for tag generation."""
        self._custom_prompt = value

    @property
    def wd14_general_threshold(self) -> float:
        return self._wd14_general_threshold

    @wd14_general_threshold.setter
    def wd14_general_threshold(self, value: float):
        self._wd14_general_threshold = max(0.0, min(1.0, value))

    @property
    def wd14_character_threshold(self) -> float:
        return self._wd14_character_threshold

    @wd14_character_threshold.setter
    def wd14_character_threshold(self, value: float):
        self._wd14_character_threshold = max(0.0, min(1.0, value))

    def get_models_status(self) -> Dict[str, Dict[str, Any]]:
        """Get availability status for all model types.

        Returns:
            Dictionary with status for each model type:
            {
                'gguf': {
                    'name': 'GGUF (Recommended)',
                    'description': '...',
                    'size': '~4.5GB',
                    'downloaded': True,
                    'model_exists': True,
                    'mmproj_exists': True,  # GGUF only
                    'model_path': '/path/to/model'
                },
                'hf': {...}
            }
        """
        status = {}

        for model_type, config in MODELS.items():
            model_status = {
                "name": config["name"],
                "description": config["description"],
                "size": config["size"],
                "downloaded": False,
                "model_path": None,
            }

            if model_type == "gguf":
                model_exists, mmproj_exists = self._check_gguf_models()
                model_status["model_exists"] = model_exists
                model_status["mmproj_exists"] = mmproj_exists
                model_status["downloaded"] = model_exists and mmproj_exists
                if model_status["downloaded"]:
                    model_status["model_path"] = str(
                        self.models_dir / config["subdir"] / config["filename"]
                    )
            elif model_type.startswith("wd14"):
                model_status["downloaded"] = self._check_wd14_model(model_type)
                if model_status["downloaded"]:
                    model_status["model_path"] = str(
                        self.models_dir / config["subdir"] / "model.onnx"
                    )
            else:  # hf
                model_status["downloaded"] = self._check_hf_model()
                if model_status["downloaded"]:
                    # Get the actual path (local or cache)
                    model_status["model_path"] = str(self._get_hf_model_path())

            status[model_type] = model_status

        return status

    def _check_gguf_models(self) -> Tuple[bool, bool]:
        """Check if GGUF model and mmproj files exist.

        Returns:
            Tuple of (model_exists, mmproj_exists)
        """
        config = MODELS["gguf"]
        gguf_dir = self.models_dir / config["subdir"]
        model_path = gguf_dir / config["filename"]
        mmproj_path = gguf_dir / config["mmproj_filename"]
        return model_path.exists(), mmproj_path.exists()

    def _check_hf_model(self) -> bool:
        """Check if HuggingFace model exists.

        Checks both the local models directory and the HuggingFace cache.

        Returns:
            True if model directory contains config.json
        """
        config = MODELS["hf"]

        # Check local directory first
        model_dir = self.models_dir / config["subdir"]
        self.logger.debug(f"Checking local HF model path: {model_dir}")
        if (model_dir / "config.json").exists():
            self.logger.debug("Found HF model in local directory")
            return True

        # Check HuggingFace cache as fallback
        hf_cache_path = self._get_hf_cache_path(config["repo"])
        if hf_cache_path:
            self.logger.debug(f"Found HF model in cache: {hf_cache_path}")
            return True

        self.logger.debug("HF model not found in local dir or cache")
        return False

    def _get_hf_cache_path(self, repo_id: str) -> Optional[Path]:
        """Get the path to a model in the HuggingFace cache.

        Args:
            repo_id: The HuggingFace repo ID
                (e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava')

        Returns:
            Path to the cached model directory, or None if not found
        """
        try:
            from huggingface_hub import scan_cache_dir
        except ImportError:
            self.logger.debug("huggingface_hub not available for cache check")
            return None

        try:
            cache_info = scan_cache_dir()
            for repo in cache_info.repos:
                if repo.repo_id == repo_id and repo.repo_type == "model":
                    # Get the latest revision's snapshot path
                    for revision in repo.revisions:
                        snapshot_path = revision.snapshot_path
                        if (Path(snapshot_path) / "config.json").exists():
                            return Path(snapshot_path)
        except Exception as e:
            self.logger.debug(f"Error scanning HF cache: {e}")

        return None

    def _get_hf_model_path(self) -> Optional[Path]:
        """Get the actual path to the HuggingFace model.

        Checks local directory first, then HuggingFace cache.

        Returns:
            Path to the model directory, or None if not found
        """
        config = MODELS["hf"]

        # Check local directory first
        model_dir = self.models_dir / config["subdir"]
        if (model_dir / "config.json").exists():
            return model_dir

        # Check HuggingFace cache as fallback
        cache_path = self._get_hf_cache_path(config["repo"])
        if cache_path:
            return cache_path

        return None

    def _check_wd14_model(self, model_type: str) -> bool:
        """Check if WD14 model files exist."""
        config = MODELS[model_type]
        model_dir = self.models_dir / config["subdir"]
        return all((model_dir / f).exists() for f in config["files"])

    def _download_wd14_model(
        self,
        model_type: str,
        progress_callback: Optional[Callable[[str, float], None]] = None,
    ) -> bool:
        """Download WD14 model files from HuggingFace."""
        from huggingface_hub import hf_hub_download

        config = MODELS[model_type]
        model_dir = self.models_dir / config["subdir"]
        model_dir.mkdir(parents=True, exist_ok=True)

        files = config["files"]
        for i, filename in enumerate(files):
            file_path = model_dir / filename
            if not file_path.exists():
                progress = int((i / len(files)) * 80) + 10
                if progress_callback:
                    progress_callback(f"Downloading {filename}...", progress)

                self.logger.info(f"Downloading WD14 file: {filename}")
                hf_hub_download(
                    repo_id=config["repo"],
                    filename=filename,
                    local_dir=str(model_dir),
                    local_dir_use_symlinks=False,
                )

        if progress_callback:
            progress_callback("Download complete", 100)
        return True

    def _load_wd14_tagger(self, model_type: str, use_gpu: bool = True):
        """Load WD14 ONNX model and tag list."""
        try:
            import onnxruntime as ort
        except ImportError:
            raise RuntimeError(
                "onnxruntime is required for WD14 models. "
                "Install with: pip install onnxruntime (or onnxruntime-gpu for CUDA)"
            )
        import csv

        config = MODELS[model_type]
        model_dir = self.models_dir / config["subdir"]
        model_path = model_dir / "model.onnx"
        tags_path = model_dir / "selected_tags.csv"

        self.logger.info(f"Loading WD14 model from {model_dir}...")

        # Set up ONNX session
        providers = []
        if use_gpu:
            providers.append("CUDAExecutionProvider")
        providers.append("CPUExecutionProvider")

        session = ort.InferenceSession(str(model_path), providers=providers)

        # Read tag list from CSV
        tags_list = []
        with open(tags_path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                tags_list.append(row)

        self.logger.info(
            f"WD14 model loaded ({len(tags_list)} tags, "
            f"provider: {session.get_providers()[0]})"
        )
        return ("wd14", (session, tags_list))

    def _generate_wd14(
        self,
        session,
        tags_list: List[Dict[str, str]],
        image: Image.Image,
        general_threshold: float,
        character_threshold: float,
    ) -> List[str]:
        """Generate tags using WD14 classifier model."""
        import numpy as np

        # Get model input dimensions
        input_shape = session.get_inputs()[0].shape
        target_size = input_shape[2]  # e.g. 448

        # Pad to square with white background
        max_dim = max(image.size)
        padded = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
        padded.paste(
            image, ((max_dim - image.width) // 2, (max_dim - image.height) // 2)
        )

        # Resize to model input size
        padded = padded.resize((target_size, target_size), Image.Resampling.BILINEAR)

        # Convert to numpy: RGB -> BGR, float32, add batch dim
        img_array = np.array(padded, dtype=np.float32)
        img_array = img_array[:, :, ::-1]  # RGB to BGR
        img_array = np.expand_dims(img_array, axis=0)

        # Run inference
        input_name = session.get_inputs()[0].name
        output = session.run(None, {input_name: img_array})
        probabilities = output[0][0]  # First output, first batch

        # Filter tags by category and threshold
        result_tags = []
        for i, tag_info in enumerate(tags_list):
            if i >= len(probabilities):
                break
            prob = float(probabilities[i])
            category = tag_info.get("category", "0")
            tag_name = tag_info.get("name", "").strip()

            if not tag_name:
                continue

            # Category 0 = general, 4 = character, 9 = rating
            if category == "0" and prob >= general_threshold:
                result_tags.append((tag_name, prob))
            elif category == "4" and prob >= character_threshold:
                result_tags.append((tag_name, prob))

        # Sort by confidence descending
        result_tags.sort(key=lambda x: x[1], reverse=True)
        return [tag for tag, _ in result_tags]

    def download_model(
        self,
        model_type: str,
        progress_callback: Optional[Callable[[str, float], None]] = None,
    ) -> bool:
        """Download a model with optional progress updates.

        Args:
            model_type: Either 'gguf' or 'hf'
            progress_callback: Optional callback(status_message, progress_percent)

        Returns:
            True if download successful, False otherwise

        Raises:
            ValueError: If model_type is not valid
        """
        if model_type not in MODELS:
            raise ValueError(
                f"Invalid model type: {model_type}. "
                f"Must be one of: {', '.join(MODELS.keys())}"
            )

        try:
            from huggingface_hub import (  # noqa: F401 - availability probe
                hf_hub_download,
                snapshot_download,
            )
        except ImportError:
            self.logger.error("huggingface_hub not installed")
            if progress_callback:
                progress_callback("Error: huggingface_hub not installed", 0)
            return False

        try:
            if model_type == "gguf":
                return self._download_gguf_models(progress_callback)
            elif model_type.startswith("wd14"):
                return self._download_wd14_model(model_type, progress_callback)
            else:
                return self._download_hf_model(progress_callback)
        except Exception as e:
            self.logger.error(f"Download failed: {e}")
            if progress_callback:
                progress_callback(f"Error: {str(e)}", 0)
            return False

    def _download_gguf_models(
        self, progress_callback: Optional[Callable[[str, float], None]] = None
    ) -> bool:
        """Download GGUF model and mmproj files."""
        from huggingface_hub import hf_hub_download

        config = MODELS["gguf"]
        gguf_dir = self.models_dir / config["subdir"]
        gguf_dir.mkdir(parents=True, exist_ok=True)

        model_path = gguf_dir / config["filename"]
        mmproj_path = gguf_dir / config["mmproj_filename"]

        # Download main model
        if not model_path.exists():
            if progress_callback:
                progress_callback(f"Downloading {config['filename']}...", 10)

            self.logger.info(f"Downloading GGUF model: {config['filename']}")
            hf_hub_download(
                repo_id=config["repo"],
                filename=config["filename"],
                local_dir=str(gguf_dir),
                local_dir_use_symlinks=False,
            )
            self.logger.info("GGUF model downloaded")

        if progress_callback:
            progress_callback("Main model ready", 50)

        # Download mmproj
        if not mmproj_path.exists():
            if progress_callback:
                progress_callback(f"Downloading {config['mmproj_filename']}...", 60)

            self.logger.info(f"Downloading mmproj: {config['mmproj_filename']}")
            hf_hub_download(
                repo_id=config["mmproj_repo"],
                filename=config["mmproj_filename"],
                local_dir=str(gguf_dir),
                local_dir_use_symlinks=False,
            )
            self.logger.info("mmproj downloaded")

        if progress_callback:
            progress_callback("Download complete", 100)

        return True

    def _download_hf_model(
        self, progress_callback: Optional[Callable[[str, float], None]] = None
    ) -> bool:
        """Download HuggingFace model."""
        from huggingface_hub import snapshot_download

        config = MODELS["hf"]
        model_dir = self.models_dir / config["subdir"]

        if not self._check_hf_model():
            if progress_callback:
                progress_callback(f"Downloading {config['repo']}...", 10)

            self.logger.info(f"Downloading HF model: {config['repo']}")
            snapshot_download(
                repo_id=config["repo"],
                local_dir=str(model_dir),
                local_dir_use_symlinks=False,
            )
            self.logger.info("HF model downloaded")

        if progress_callback:
            progress_callback("Download complete", 100)

        return True

    def load_model(self, model_type: str, use_gpu: bool = True) -> bool:
        """Load a model into memory for tag generation.

        Args:
            model_type: Either 'gguf' or 'hf'
            use_gpu: Whether to use GPU acceleration (default True)

        Returns:
            True if model loaded successfully

        Raises:
            ValueError: If model_type is invalid
            RuntimeError: If model not downloaded or loading fails
        """
        if model_type not in MODELS:
            raise ValueError(f"Invalid model type: {model_type}")

        # Unload existing model first
        if self._tagger is not None:
            self.unload_model()

        status = self.get_models_status()
        if not status[model_type]["downloaded"]:
            raise RuntimeError(f"Model {model_type} not downloaded")

        try:
            self._tagger = self._model_loader(model_type, use_gpu)
            self._current_model_type = model_type
            self.logger.info(f"Model {model_type} loaded successfully")
            return True

        except Exception as e:
            self.logger.error(f"Failed to load model {model_type}: {e}")
            self._tagger = None
            self._current_model_type = None
            raise RuntimeError(f"Failed to load model: {e}")

    def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]:
        """Default loader: dispatch to a backend, importing it only now."""
        if model_type == "gguf":
            return self._load_gguf_tagger(use_gpu)
        if model_type.startswith("wd14"):
            return self._load_wd14_tagger(model_type, use_gpu)
        return self._load_hf_tagger()

    def _load_gguf_tagger(self, use_gpu: bool = True):
        """Load GGUF-based tagger."""
        from llama_cpp import Llama
        from llama_cpp.llama_chat_format import Llava15ChatHandler

        config = MODELS["gguf"]
        gguf_dir = self.models_dir / config["subdir"]
        model_path = gguf_dir / config["filename"]
        mmproj_path = gguf_dir / config["mmproj_filename"]

        self.logger.info("Loading GGUF model...")
        n_gpu_layers = -1 if use_gpu else 0

        tagger = Llama(
            model_path=str(model_path),
            n_ctx=4096,
            n_batch=2048,
            n_threads=4,
            n_gpu_layers=n_gpu_layers,
            verbose=False,
            chat_handler=Llava15ChatHandler(clip_model_path=str(mmproj_path)),
            offload_kqv=True,
        )

        self.logger.info("GGUF model loaded")
        return ("gguf", tagger)

    def _load_hf_tagger(self, quantization: str = "8bit"):
        """Load HuggingFace-based tagger."""
        import torch
        from transformers import (
            AutoProcessor,
            LlavaForConditionalGeneration,
            BitsAndBytesConfig,
        )

        # Get the actual model path (local or cache)
        model_path = self._get_hf_model_path()
        if model_path is None:
            raise RuntimeError(
                "HuggingFace model not found in local directory or cache"
            )

        self.logger.info(f"Loading HuggingFace model from {model_path}...")
        device = "cuda" if torch.cuda.is_available() else "cpu"

        processor = AutoProcessor.from_pretrained(str(model_path))
        model_kwargs = {"device_map": "cuda" if device == "cuda" else "cpu"}

        if quantization == "8bit":
            qnt_config = BitsAndBytesConfig(
                load_in_8bit=True,
                bnb_8bit_compute_dtype=torch.float16,
                bnb_8bit_use_double_quant=True,
                llm_int8_skip_modules=["vision_tower", "multi_modal_projector"],
            )
            model = LlavaForConditionalGeneration.from_pretrained(
                str(model_path),
                torch_dtype=torch.float16,
                quantization_config=qnt_config,
                **model_kwargs,
            )
        else:
            model = LlavaForConditionalGeneration.from_pretrained(
                str(model_path), torch_dtype=torch.bfloat16, **model_kwargs
            )

        model.eval()
        self.logger.info("HuggingFace model loaded")
        # Track the compute dtype for pixel_values conversion
        compute_dtype = torch.float16 if quantization == "8bit" else torch.bfloat16
        return ("hf", (model, processor, device, compute_dtype))

    def unload_model(self):
        """Unload the current model and free memory."""
        if self._tagger is not None:
            self.logger.info(f"Unloading model: {self._current_model_type}")
            del self._tagger
            self._tagger = None
            self._current_model_type = None
            gc.collect()

            # Clear the CUDA cache only if torch is already loaded; importing
            # it here would pull in a heavy dependency the backend never used.
            torch = sys.modules.get("torch")
            if torch is not None and torch.cuda.is_available():
                torch.cuda.empty_cache()

            self.logger.info("Model unloaded, memory freed")

    def is_model_loaded(self) -> bool:
        """Check if a model is currently loaded."""
        return self._tagger is not None

    def get_loaded_model_type(self) -> Optional[str]:
        """Get the type of currently loaded model."""
        return self._current_model_type

    def generate_tags(
        self,
        image_path: str,
        prompt: Optional[str] = None,
        general_threshold: Optional[float] = None,
        character_threshold: Optional[float] = None,
    ) -> List[str]:
        """Generate tags for an image.

        Args:
            image_path: Path to the image file
            prompt: Custom prompt for tag generation (LLM models only).
                Uses default if None.
            general_threshold: Confidence threshold for general tags (WD14 only).
            character_threshold: Confidence threshold for character tags (WD14 only).

        Returns:
            List of generated tags

        Raises:
            RuntimeError: If no model is loaded
            FileNotFoundError: If image doesn't exist
        """
        if self._tagger is None:
            raise RuntimeError("No model loaded. Call load_model() first.")

        if not os.path.exists(image_path):
            raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path))

        use_prompt = prompt or self._custom_prompt

        # Load image
        image = Image.open(image_path)
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Generate based on model type
        model_type, tagger_obj = self._tagger

        if model_type == "gguf":
            raw_tags = self._generate_gguf(tagger_obj, image, use_prompt)
            return self._parse_tags(raw_tags)
        elif model_type == "wd14":
            session, tags_list = tagger_obj
            gen_thresh = (
                general_threshold
                if general_threshold is not None
                else self._wd14_general_threshold
            )
            char_thresh = (
                character_threshold
                if character_threshold is not None
                else self._wd14_character_threshold
            )
            return self._generate_wd14(
                session, tags_list, image, gen_thresh, char_thresh
            )
        else:
            model, processor, device, compute_dtype = tagger_obj
            raw_tags = self._generate_hf(
                model, processor, device, compute_dtype, image, use_prompt
            )
            return self._parse_tags(raw_tags)

    def _generate_gguf(self, model, image: Image.Image, prompt: str) -> str:
        """Generate tags using GGUF model."""
        import base64
        import io

        # Resize image
        image = image.resize((336, 336), Image.Resampling.BILINEAR)

        # Encode to base64
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        buffer.seek(0)
        img_base64 = base64.b64encode(buffer.read()).decode("utf-8")
        data_uri = f"data:image/png;base64,{img_base64}"

        # Create message
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": data_uri}},
                ],
            },
        ]

        # Generate
        response = model.create_chat_completion(
            messages=messages,
            max_tokens=512,
            temperature=0.6,
            top_p=0.9,
            stop=["</s>", "User:", "Assistant:"],
            stream=False,
        )

        return response["choices"][0]["message"]["content"].strip()

    def _generate_hf(
        self,
        model,
        processor,
        device: str,
        compute_dtype,
        image: Image.Image,
        prompt: str,
    ) -> str:
        """Generate tags using HuggingFace model."""
        import torch

        # Resize image
        image = image.resize((336, 336), Image.Resampling.LANCZOS)

        convo = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]
        convo_string = processor.apply_chat_template(
            convo, tokenize=False, add_generation_prompt=True
        )

        inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(
            device
        )

        # Convert pixel_values to the model's compute dtype
        # (float16 for 8-bit quantized, bfloat16 for non-quantized)
        if "pixel_values" in inputs and inputs["pixel_values"] is not None:
            inputs["pixel_values"] = inputs["pixel_values"].to(compute_dtype)

        with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True):
            generate_ids = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=True,
                temperature=0.6,
                top_p=0.9,
                use_cache=True,
            )[0]

        generate_ids = generate_ids[inputs["input_ids"].shape[1] :]
        return processor.tokenizer.decode(
            generate_ids, skip_special_tokens=True
        ).strip()

    def _parse_tags(self, raw_output: str) -> List[str]:
        """Parse raw model output into a list of clean tags.

        Args:
            raw_output: Raw text output from the model

        Returns:
            List of cleaned, deduplicated tags
        """
        # Patterns to exclude
        exclude_prefixes = (
            "copyright:",
            "meta:",
            "photo_",
            "photo:",
            "prepend:",
            "append:",
        )

        # Split by common delimiters
        tags = []

        # Handle comma-separated tags
        for part in raw_output.split(","):
            tag = part.strip().lower()
            # Remove any quotes or extra characters
            tag = tag.strip("\"'")
            # Replace spaces with underscores (Danbooru style)
            tag = tag.replace(" ", "_")
            # Remove empty tags
            if tag and len(tag) > 1:
                # Filter out unwanted tag patterns
                if not tag.startswith(exclude_prefixes):
                    tags.append(tag)

        # Deduplicate while preserving order
        seen = set()
        unique_tags = []
        for tag in tags:
            if tag not in seen:
                seen.add(tag)
                unique_tags.append(tag)

        return unique_tags


# Singleton instance for API use
_service_instance: Optional[AutoTagService] = None


def get_autotag_service() -> AutoTagService:
    """Get or create the singleton AutoTagService instance."""
    global _service_instance
    if _service_instance is None:
        _service_instance = AutoTagService()
    return _service_instance