"""
AutoTag Service Module

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

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

# 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",
    },
}

# Default prompts
DEFAULT_PROMPT = """Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).
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.
Prioritize visual accuracy.
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):
        """Initialize the AutoTag service.

        Args:
            models_dir: Directory for storing models. If None, uses ComfyUI's
                       folder_paths.models_dir / "LLM" path.
        """
        self.logger = get_logger("autotag.service")

        # 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

    @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

    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"]
                    )
            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, HFCacheInfo
        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 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}. Must be 'gguf' or 'hf'"
            )

        try:
            from huggingface_hub import 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)
            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:
            if model_type == "gguf":
                self._tagger = self._load_gguf_tagger(use_gpu)
            else:
                self._tagger = self._load_hf_tagger()

            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_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()

            # Try to clear CUDA cache if available
            try:
                import torch

                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass

            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) -> List[str]:
        """Generate tags for an image.

        Args:
            image_path: Path to the image file
            prompt: Custom prompt for tag generation. Uses default if None.

        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(f"Image not found: {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)
        else:
            model, processor, device, compute_dtype = tagger_obj
            raw_tags = self._generate_hf(
                model, processor, device, compute_dtype, image, use_prompt
            )

        # Parse tags from response
        tags = self._parse_tags(raw_tags)
        return 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:",
        )

        # 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