import torch
from transformers import AutoProcessor, AutoModelForCausalLM, AutoTokenizer, AutoConfig
from PIL import Image
import os
import comfy.model_management as mm
import hashlib
import gc

class UnifiedVisionPromptGenerator:
    # Class-level cache for models
    _model_cache = {}
    _processor_cache = {}
    _current_model_key = None  # Track which model is currently loaded
    
    # Model type detection
    MODEL_TYPES = {
        'qwen2vl': [
            'Qwen/Qwen2-VL-2B-Instruct',
            'Qwen/Qwen2-VL-7B-Instruct',
            'Qwen/Qwen2.5-VL-3B-Instruct',
            'Qwen/Qwen2.5-VL-7B-Instruct',
            'Qwen/Qwen2.5-VL-72B-Instruct',
        ],
        'moondream': [
            'vikhyatk/moondream2',
            'vikhyatk/moondream1',
        ],
        'smolvlm': [
            'HuggingFaceTB/SmolVLM-Instruct',
            'HuggingFaceTB/SmolVLM-256M-Instruct',
            'HuggingFaceTB/SmolVLM-500M-Instruct',
        ],
    }
    
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        # Flatten all models into one list
        all_models = []
        for model_list in cls.MODEL_TYPES.values():
            all_models.extend(model_list)
        
        return {
            "required": {
                "question": ("STRING", {"default": "Describe this image in detail.", "multiline": True}),
                "max_new_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.1}),
                "seed": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
                "model": (all_models, {"default": "Qwen/Qwen2.5-VL-7B-Instruct"}),
                "keep_model_loaded": ("BOOLEAN", {"default": False}),  # Option to cache or not
            },
            "optional": {
                "image": ("IMAGE",),
            }
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)
    FUNCTION = "generate_prompt"
    CATEGORY = "GraftingRayman/Image Processing"

    def hash_seed(self, seed):
        """Hash the seed for reproducibility."""
        seed_bytes = str(seed).encode('utf-8')
        hash_object = hashlib.sha256(seed_bytes)
        return int(hash_object.hexdigest(), 16) % (2**32)

    def get_model_type(self, model_path):
        """Detect which model type we're using."""
        for model_type, model_list in self.MODEL_TYPES.items():
            if model_path in model_list:
                return model_type
        # Default to qwen2vl for unknown models
        return 'qwen2vl'

    @classmethod
    def clear_cache(cls):
        """Clear all cached models from memory."""
        print("[UnifiedVisionPromptGenerator] Clearing model cache...")
        
        # Move models to CPU and delete
        for key, model in cls._model_cache.items():
            try:
                if hasattr(model, 'to'):
                    model.to('cpu')
                del model
            except Exception as e:
                print(f"Error clearing model {key}: {e}")
        
        cls._model_cache.clear()
        cls._processor_cache.clear()
        cls._current_model_key = None
        
        # Force garbage collection and clear CUDA cache
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
        
        print("[UnifiedVisionPromptGenerator] Cache cleared.")

    @classmethod
    def unload_if_different(cls, new_cache_key):
        """Unload current model if switching to a different one."""
        if cls._current_model_key and cls._current_model_key != new_cache_key:
            print(f"[UnifiedVisionPromptGenerator] Switching models, unloading {cls._current_model_key}")
            
            # Unload the old model
            if cls._current_model_key in cls._model_cache:
                old_model = cls._model_cache[cls._current_model_key]
                try:
                    if hasattr(old_model, 'to'):
                        old_model.to('cpu')
                    del cls._model_cache[cls._current_model_key]
                    del cls._processor_cache[cls._current_model_key]
                except Exception as e:
                    print(f"Error unloading old model: {e}")
            
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    def load_model(self, model_path, device, keep_loaded=False):
        """Load the processor and model with optional caching."""
        try:
            model_type = self.get_model_type(model_path)
            cache_key = f"{model_path}_{model_type}"
            
            # If not keeping models loaded, clear cache before loading new model
            if not keep_loaded:
                self.clear_cache()
            else:
                # If keeping loaded but switching models, unload the old one
                self.unload_if_different(cache_key)
            
            # Check cache first
            if cache_key in self._model_cache and cache_key in self._processor_cache:
                print(f"[UnifiedVisionPromptGenerator] Using cached model: {model_path}")
                model = self._model_cache[cache_key]
                processor = self._processor_cache[cache_key]
                
                # Ensure model is on correct device
                if hasattr(model, 'parameters') and next(model.parameters()).device != device:
                    print(f"[UnifiedVisionPromptGenerator] Moving cached model to {device}")
                    model = model.to(device)
                    self._model_cache[cache_key] = model
                
                self._current_model_key = cache_key
                return processor, model, model_type
            
            print(f"[UnifiedVisionPromptGenerator] Loading model: {model_path}")
            
            # Load based on model type
            if model_type == 'qwen2vl':
                try:
                    from transformers import Qwen2VLForConditionalGeneration
                    from qwen_vl_utils import process_vision_info
                except ImportError:
                    print("Error: Qwen2VL requires 'qwen-vl-utils'. Install with: pip install qwen-vl-utils")
                    raise
                    
                processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
                model = Qwen2VLForConditionalGeneration.from_pretrained(
                    model_path,
                    trust_remote_code=True,
                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
                ).to(device)
                
            elif model_type == 'moondream':
                tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
                model = AutoModelForCausalLM.from_pretrained(
                    model_path,
                    trust_remote_code=True,
                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
                ).to(device)
                processor = tokenizer
                
            else:  # smolvlm
                processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
                try:
                    model = AutoModelForCausalLM.from_pretrained(
                        model_path,
                        trust_remote_code=True,
                        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                        attn_implementation="flash_attention_2" if torch.cuda.is_available() else "eager"
                    ).to(device)
                except:
                    model = AutoModelForCausalLM.from_pretrained(
                        model_path,
                        trust_remote_code=True,
                        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
                    ).to(device)
            
            # Only cache if keep_loaded is True
            if keep_loaded:
                self._model_cache[cache_key] = model
                self._processor_cache[cache_key] = processor
                self._current_model_key = cache_key
                print(f"[UnifiedVisionPromptGenerator] Model cached in VRAM")
            else:
                print(f"[UnifiedVisionPromptGenerator] Model loaded (not cached)")
            
            return processor, model, model_type
            
        except Exception as e:
            print(f"Error loading model {model_path}: {str(e)}")
            raise

    def convert_image(self, image):
        """Safely convert image tensor to PIL Image."""
        try:
            if isinstance(image, torch.Tensor):
                if image.dim() == 4:
                    image = image[0]
                elif image.dim() != 3:
                    raise ValueError(f"Unexpected image dimensions: {image.dim()}")
                
                if image.is_cuda:
                    image = image.cpu()
                
                image_np = image.mul(255).clamp(0, 255).byte().numpy()
                
                if image_np.shape[0] in [1, 3, 4]:
                    image_np = image_np.transpose(1, 2, 0)
                
                if image_np.shape[2] == 1:
                    image_np = image_np.repeat(3, axis=2)
                elif image_np.shape[2] == 4:
                    image_np = image_np[:, :, :3]
                
                return Image.fromarray(image_np)
            elif isinstance(image, Image.Image):
                return image
            else:
                raise ValueError(f"Unsupported image type: {type(image)}")
                
        except Exception as e:
            print(f"Error converting image: {str(e)}")
            raise

    def generate_qwen2vl(self, processor, model, device, question, image, max_new_tokens, temperature):
        """Generate using Qwen2-VL models."""
        from qwen_vl_utils import process_vision_info
        
        if image is not None:
            pil_image = self.convert_image(image)
            temp_path = "/tmp/temp_qwen_image.png"
            pil_image.save(temp_path)
            
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": f"file://{temp_path}"},
                        {"type": "text", "text": question},
                    ],
                }
            ]
        else:
            messages = [{"role": "user", "content": [{"type": "text", "text": question}]}]

        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        
        inputs = processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(device)

        with torch.no_grad():
            generated_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
            )
        
        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        
        output_text = processor.batch_decode(
            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
        )[0]
        
        return output_text

    def generate_moondream(self, processor, model, device, question, image, max_new_tokens):
        """Generate using Moondream models."""
        if image is None:
            return "Moondream requires an image input."
        
        pil_image = self.convert_image(image)
        enc_image = model.encode_image(pil_image)
        
        with torch.no_grad():
            try:
                answer = model.answer_question(enc_image, question, processor, max_tokens=max_new_tokens)
            except:
                answer = model.answer_question(enc_image, question, processor)
        
        # Clean up the answer - remove question if it's echoed
        answer = answer.strip()
        
        patterns_to_remove = [
            question,
            f"{question}:",
            f"{question} ",
            f"Question: {question}",
            f"Q: {question}",
        ]
        
        for pattern in patterns_to_remove:
            if answer.lower().startswith(pattern.lower()):
                answer = answer[len(pattern):].strip()
                answer = answer.lstrip(':-').strip()
                break
        
        if answer.lower().startswith("answer:"):
            answer = answer[7:].strip()
        elif answer.lower().startswith("a:"):
            answer = answer[2:].strip()
        
        return answer

    def generate_smolvlm(self, processor, model, device, question, image, max_new_tokens, temperature):
        """Generate using SmolVLM models."""
        if image is not None:
            pil_image = self.convert_image(image)
            
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image"},
                        {"type": "text", "text": question}
                    ]
                }
            ]
            
            prompt_text = processor.apply_chat_template(messages, add_generation_prompt=True)
            inputs = processor(text=prompt_text, images=[pil_image], return_tensors="pt").to(device)
        else:
            messages = [{"role": "user", "content": [{"type": "text", "text": question}]}]
            prompt_text = processor.apply_chat_template(messages, add_generation_prompt=True)
            inputs = processor(text=prompt_text, return_tensors="pt").to(device)

        with torch.no_grad():
            generated_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
            )

        generated_ids_trimmed = [
            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
        ]
        
        output_text = processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
        )[0]
        
        return output_text.strip()

    def generate_prompt(self, question, max_new_tokens, temperature, seed, model, keep_model_loaded=False, image=None):
        """Generate a prompt using the selected vision model."""
        generated_ids = None
        inputs = None
        pil_image = None
        model_obj = None
        should_unload = not keep_model_loaded
        
        try:
            device = mm.get_torch_device()
            offload_device = mm.unet_offload_device()

            processor, model_obj, model_type = self.load_model(model, device, keep_model_loaded)

            if seed:
                torch.manual_seed(self.hash_seed(seed))
                if torch.cuda.is_available():
                    torch.cuda.manual_seed(self.hash_seed(seed))

            # Generate based on model type
            if model_type == 'qwen2vl':
                generated_prompt = self.generate_qwen2vl(processor, model_obj, device, question, image, max_new_tokens, temperature)
                
            elif model_type == 'moondream':
                generated_prompt = self.generate_moondream(processor, model_obj, device, question, image, max_new_tokens)
                
            elif model_type == 'smolvlm':
                generated_prompt = self.generate_smolvlm(processor, model_obj, device, question, image, max_new_tokens, temperature)
            
            else:
                generated_prompt = "Unknown model type."

            print(f"Model Type: {model_type}")
            print(f"Generated prompt: {generated_prompt}")
            
            return (generated_prompt,)
            
        except Exception as e:
            print(f"Error in generate_prompt: {str(e)}")
            import traceback
            traceback.print_exc()
            return ("",)
            
        finally:
            # Clean up intermediate tensors
            if generated_ids is not None:
                del generated_ids
            if inputs is not None:
                del inputs
            if pil_image is not None:
                del pil_image
            
            # If not keeping model loaded, unload it
            if should_unload and model_obj is not None:
                try:
                    print("[UnifiedVisionPromptGenerator] Unloading model after inference")
                    if hasattr(model_obj, 'to'):
                        model_obj.to('cpu')
                    self.clear_cache()
                except Exception as e:
                    print(f"Error during model cleanup: {e}")
            
            # Always clean up ComfyUI memory
            mm.soft_empty_cache()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()

NODE_CLASS_MAPPINGS = {
    "UnifiedVisionPromptGenerator": UnifiedVisionPromptGenerator
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "UnifiedVisionPromptGenerator": "Unified Vision Prompt Generator"
}