import os
import io
import re
import yaml
import json
import torch
import torchvision
import cv2
import base64
import logging
import datetime
import requests
import time
import numpy as np
from io import BytesIO
from aiohttp import web
from dotenv import load_dotenv
from PIL import Image, ImageOps, ImageSequence
from typing import Tuple, Optional, Dict, Union, List, Any
import node_helpers
from torchvision.transforms import functional as TF
import folder_paths


from typing import Union, List, Tuple

logger = logging.getLogger(__name__)

def process_frames(self, images_tensor, frame_sample_count, max_pixels=512*512):
    """
    Process image frames for the model by sampling and resizing.

    Args:
        images_tensor: Input tensor in shape [B,H,W,C]
        frame_sample_count: Number of frames to sample
        max_pixels: Max pixels for each frame

    Returns:
        List of processed PIL images
    """
    try:
        # Get the batch size (number of frames)
        batch_size = images_tensor.shape[0]

        # Sample the frames evenly from the batch
        if batch_size <= frame_sample_count:
            # Use all frames if we have fewer than requested
            sampled_indices = list(range(batch_size))
        else:
            # Sample evenly across the frames
            sampled_indices = [
                int(i * (batch_size - 1) / (frame_sample_count - 1))
                for i in range(frame_sample_count)
            ]

        # Extract the sampled frames from the tensor
        sampled_frames = [images_tensor[i] for i in sampled_indices]

        # Convert to PIL images with proper preprocessing
        pil_images = []
        for frame in sampled_frames:
            # Convert tensor to numpy
            frame_np = frame.cpu().numpy()

            # Scale from [0,1] to [0,255] if needed
            if frame_np.max() <= 1.0:
                frame_np = (frame_np * 255).astype(np.uint8)
            else:
                frame_np = frame_np.astype(np.uint8)

            # Convert to PIL
            pil_image = Image.fromarray(frame_np)

            # Ensure RGB mode
            if pil_image.mode != "RGB":
                pil_image = pil_image.convert("RGB")

            # Calculate resize dimensions if needed
            if max_pixels > 0:
                width, height = pil_image.size
                if width * height > max_pixels:
                    # Calculate new dimensions while maintaining aspect ratio
                    ratio = math.sqrt(max_pixels / (width * height))
                    new_width = int(width * ratio)
                    new_height = int(height * ratio)

                    # Ensure dimensions are multiples of 32 for better compatibility
                    new_width = (new_width // 32) * 32
                    new_height = (new_height // 32) * 32

                    # Resize using LANCZOS for better quality
                    pil_image = pil_image.resize((new_width, new_height), Image.LANCZOS)

            # Add final checks
            if pil_image.size[0] > 2048 or pil_image.size[1] > 2048:
                # Limit maximum dimension to 2048
                pil_image.thumbnail((2048, 2048), Image.LANCZOS)

            pil_images.append(pil_image)

        return pil_images

    except Exception as e:
        logger.error(f"Error processing frames: {e}")
        # Return a single black frame as fallback
        fallback_size = (512, 512)
        return [Image.new("RGB", fallback_size, (0, 0, 0))]

def resize_image_max_side(img, max_size):
    """Resize image so its longest side is max_size while maintaining aspect ratio"""
    ratio = max_size / max(img.size)
    if ratio < 1:  # Only resize if image is larger than max_size
        new_size = tuple(int(dim * ratio) for dim in img.size)
        return img.resize(new_size, Image.LANCZOS)
    return img

def prepare_batch_images(images):
    """
    Convert images to list of batches.
    Handles tensor, list, and single image inputs while preserving dimensions.
    
    Args:
        images: torch.Tensor or list of tensors
        
    Returns:
        List of image tensors
    """
    try:
        if images is None:
            return []
            
        if isinstance(images, torch.Tensor):
            # Handle 4D tensor [B,H,W,C] - split into list of [H,W,C]
            if images.dim() == 4:
                return [images[i] for i in range(images.shape[0])]
            # Handle 3D tensor [H,W,C] - wrap in list
            elif images.dim() == 3:
                return [images]
            else:
                raise ValueError(f"Invalid tensor dimensions: {images.dim()}")
                
        # Handle list input - validate each element
        if isinstance(images, list):
            for i, img in enumerate(images):
                if not isinstance(img, torch.Tensor):
                    raise ValueError(f"Image {i} is not a tensor")
            return images
            
        # Handle single image
        return [images]
        
    except Exception as e:
        logger.error(f"Error in prepare_batch_images: {str(e)}")
        return []

def process_auto_mode_images(images, mask=None, batch_size=4):
    """
    Process images and masks for auto mode with proper mask dimensionality handling.
    
    Args:
        images: Input images tensor [B,H,W,C] or list of tensors
        mask: Mask tensor [B,H,W] or [B,1,H,W] or list of tensors
        batch_size: Maximum size of each batch (default 4)
        
    Returns:
        Tuple of (image_batches, mask_batches) where each is a list of tensors
    """
    try:
        # Convert images to list format
        if images is None or (isinstance(images, (list, tuple)) and len(images) == 0):
            # Return a tuple of empty lists
            return ([], [])

        if isinstance(images, torch.Tensor):
            if images.dim() == 4:  # [B,H,W,C]
                images = [images[i] for i in range(images.shape[0])]
            elif images.dim() == 3:  # [H,W,C]
                images = [images]
            else:
                raise ValueError(f"Invalid image tensor dimensions: {images.dim()}")
        
        # Split images into batches
        image_batches = []
        current_batch = []
        
        for img in images:
            if len(current_batch) == batch_size:
                image_batches.append(torch.stack(current_batch))
                current_batch = []
            current_batch.append(img)
            
        if current_batch:  # Don't forget the last batch
            image_batches.append(torch.stack(current_batch))

        # Process masks
        mask_batches = []
        
        if mask is not None:
            # Standardize mask format
            if isinstance(mask, torch.Tensor):
                # Handle different mask dimensions
                if mask.dim() == 2:  # [H,W]
                    mask = mask.unsqueeze(0)  # -> [1,H,W]
                elif mask.dim() == 3:  # [B,H,W] or [1,H,W]
                    if mask.shape[0] != len(images):
                        # Broadcast mask to match batch size
                        mask = mask.repeat(len(images), 1, 1)
                elif mask.dim() == 4:  # [B,1,H,W] or similar
                    mask = mask.squeeze(1)  # Remove channel dim -> [B,H,W]
                
                # Split mask into batches matching image batches
                start_idx = 0
                for img_batch in image_batches:
                    batch_size = img_batch.size(0)
                    mask_batch = mask[start_idx:start_idx + batch_size]
                    
                    mask_batches.append(mask_batch)
                    start_idx += batch_size
            else:
                # Handle list of masks
                mask_list = mask if isinstance(mask, list) else [mask] * len(images)
                start_idx = 0
                for img_batch in image_batches:
                    batch_size = img_batch.size(0)
                    mask_slice = mask_list[start_idx:start_idx + batch_size]
                    
                    # Convert and stack masks
                    mask_tensors = []
                    for m in mask_slice:
                        if isinstance(m, torch.Tensor):
                            if m.dim() == 2:
                                m = m.unsqueeze(0)  # Add batch dim
                            m = m.unsqueeze(-1)  # Add channel dim at end
                        else:
                            # Convert non-tensor masks
                            m = torch.tensor(m, dtype=torch.float32)
                            if m.dim() == 2:
                                m = m.unsqueeze(0).unsqueeze(-1)
                            elif m.dim() == 3:
                                m = m.unsqueeze(-1)
                        mask_tensors.append(m)
                    
                    mask_batch = torch.stack(mask_tensors)
                    mask_batches.append(mask_batch)
                    start_idx += batch_size
        else:
            # Create default masks matching image batches
            for img_batch in image_batches:
                mask_batch = torch.ones((img_batch.size(0), img_batch.size(1), 
                                       img_batch.size(2)),  # Removed extra dimension
                                  dtype=torch.float32,
                                  device=img_batch.device)
                mask_batches.append(mask_batch)

        return image_batches, mask_batches

    except Exception as e:
        logger.error(f"Error in process_auto_mode_images: {str(e)}")
        raise

def convert_images_for_api(images, target_format='tensor'):
    """
    Convert images to the specified format for API consumption.
    Supports conversion to: tensor, base64, pil
    """
    if images is None:
        return None

    # Handle single tensor input with ComfyUI compatibility
    if isinstance(images, torch.Tensor):
        if images.dim() == 3:  # Single image
            images = images.unsqueeze(0)
        # Permute tensor to ComfyUI format (B, H, W, C) -> (B, C, H, W)
        images = images.permute(0, 3, 1, 2)

        if target_format == 'tensor':
            return images
        elif target_format == 'base64':
            return [tensor_to_base64(img) for img in images]
        elif target_format == 'pil':
            return [TF.to_pil_image(img) for img in images]
        else:
            raise ValueError(f"Unsupported target format for tensor: {target_format}")

    # Handle list of tensors input
    elif isinstance(images, list) and all(isinstance(x, torch.Tensor) for x in images):
        # Filter out tensors with unsupported channel counts
        supported_images = []
        for idx, img in enumerate(images):
            if img.shape[0] in [1, 3]:
                supported_images.append(img)
            elif img.shape[0] > 3:
                logger.warning(f"Skipping tensor at index {idx} with {img.shape[0]} channels.")
            else:
                logger.warning(f"Skipping tensor at index {idx} with unsupported number of channels: {img.shape[0]}")
        if not supported_images:
            raise ValueError("No supported image tensors found in the input list.")

        if target_format == 'tensor':
            return torch.stack(supported_images).permute(0, 3, 1, 2)  # Ensure correct format
        elif target_format == 'base64':
            return [tensor_to_base64(img) for img in supported_images]
        elif target_format == 'pil':
            return [TF.to_pil_image(img) for img in supported_images]
        else:
            raise ValueError(f"Unsupported target format for list of tensors: {target_format}")

    # Handle base64 input
    elif isinstance(images, str) or (isinstance(images, list) and all(isinstance(x, str) for x in images)):
        base64_list = [images] if isinstance(images, str) else images
        if target_format == 'base64':
            return base64_list

        # Convert base64 to PIL first
        pil_images = [base64_to_pil(b64) for b64 in base64_list]
        if target_format == 'pil':
            return pil_images
        elif target_format == 'tensor':
            tensors = [pil_to_tensor(img) for img in pil_images]
            return torch.stack(tensors).permute(0, 2, 3, 1)  # Convert to ComfyUI format (B,H,W,C)
        else:
            raise ValueError(f"Unsupported target format for base64 input: {target_format}")

    # Handle list of PIL images input
    elif isinstance(images, (list, tuple)) and all(isinstance(x, Image.Image) for x in images):
        if target_format == 'pil':
            return images
        elif target_format == 'base64':
            return [pil_image_to_base64(img) for img in images]
        elif target_format == 'tensor':
            tensors = [pil_to_tensor(img) for img in images]
            return torch.stack(tensors).permute(0, 2, 3, 1)  # Maintain ComfyUI format
        else:
            raise ValueError(f"Unsupported target format for PIL input: {target_format}")

    # If none of the above conditions are met, attempt to convert using the default method
    # Ensure that images can be saved (i.e., are PIL Images)
    else:
        try:
            encoded_images = []
            for img in images:
                if not isinstance(img, Image.Image):
                    raise ValueError(f"Expected PIL.Image, got {type(img)}")
                buffered = BytesIO()
                img.save(buffered, format="PNG")  # Adjust format if needed
                img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
                encoded_images.append(img_str)
            return encoded_images
        except Exception as e:
            raise ValueError(f"Unsupported image format or target format: {target_format}. Error: {str(e)}") from e

def convert_single_image(image, target_format):
    """Helper function to convert a single image"""
    if isinstance(image, str) and image.startswith('data:image'):
        # Convert base64 to PIL
        base64_data = image.split('base64,')[1]
        image_data = base64.b64decode(base64_data)
        image = Image.open(BytesIO(image_data))
    
    if target_format == 'pil':
        return image
    elif target_format == 'tensor':
        return pil_to_tensor(image)
    elif target_format == 'base64':
        return pil_image_to_base64(image)

def load_placeholder_image(placeholder_image_path):
        
        # Ensure the placeholder image exists
        if not os.path.exists(placeholder_image_path):
            # Create a proper RGB placeholder image
            placeholder = Image.new('RGB', (512, 512), color=(73, 109, 137))
            os.makedirs(os.path.dirname(placeholder_image_path), exist_ok=True)
            placeholder.save(placeholder_image_path)
        
        img = node_helpers.pillow(Image.open, placeholder_image_path)
        
        output_images = []
        output_masks = []
        w, h = None, None

        excluded_formats = ['MPO']
        
        for i in ImageSequence.Iterator(img):
            i = node_helpers.pillow(ImageOps.exif_transpose, i)

            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")

            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]
            
            if image.size[0] != w or image.size[1] != h:
                continue
            
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)

def process_images_for_comfy(images, placeholder_image_path=None, response_key='data', field_name='b64_json', field2_name=""):
    """Process images for ComfyUI, ensuring consistent sizes."""
    def _process_single_image(image):
        try:
            if image is None:
                return load_placeholder_image(placeholder_image_path)

            # Handle JSON/API response
            if isinstance(image, dict):
                try:
                    # Only attempt to extract from response if response_key is provided
                    if response_key and response_key in image:
                        items = image[response_key]
                        if isinstance(items, list):
                            for item in items:
                                # Only attempt to get field_name if it's provided
                                if field2_name and field_name:
                                    image_data = item.get(field2_name, {}).get(field_name)
                                elif field_name:
                                    image_data = item.get(field_name)
                                else:
                                    continue
                                
                                if image_data:
                                    # Convert the first valid image found
                                    if isinstance(image_data, str):
                                        if image_data.startswith(('data:image', 'http:', 'https:')):
                                            image = image_data  # Will be handled by URL processing below
                                        else:
                                            # Handle base64 directly
                                            image_data = base64.b64decode(image_data)
                                            image = Image.open(BytesIO(image_data))
                                            break
                    
                    if isinstance(image, dict):
                        logger.warning(f"No valid image found in response under key '{response_key}'")
                        return load_placeholder_image(placeholder_image_path)
                except Exception as e:
                    logger.error(f"Error processing API response: {str(e)}")
                    return load_placeholder_image(placeholder_image_path)

            # Convert various input types to PIL Image
            if isinstance(image, torch.Tensor):
                # Ensure tensor is in correct format [B,H,W,C] or [H,W,C]
                if image.dim() == 4:
                    if image.shape[-1] != 3:  # Wrong channel dimension
                        image = image.squeeze(1)  # Remove channel dim if [B,1,H,W]
                        if image.shape[-1] != 3:  # Still wrong shape
                            image = image.permute(0, 2, 3, 1)  # [B,C,H,W] -> [B,H,W,C]
                    image = image.squeeze(0)  # Remove batch dim
                elif image.dim() == 3 and image.shape[0] == 3:
                    image = image.permute(1, 2, 0)  # [C,H,W] -> [H,W,C]
                
                # Convert to numpy and scale to 0-255 range
                image = (image.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
                image = Image.fromarray(image)
                
            elif isinstance(image, np.ndarray):
                # Handle numpy arrays
                if image.dtype != np.uint8:
                    image = (image * 255).clip(0, 255).astype(np.uint8)
                if image.shape[-1] != 3 and image.shape[0] == 3:
                    image = np.transpose(image, (1, 2, 0))
                image = Image.fromarray(image)
                
            elif isinstance(image, str):
                if image.startswith('data:image'):
                    base64_data = image.split('base64,')[1]
                    image_data = base64.b64decode(base64_data)
                    image = Image.open(BytesIO(image_data)).convert('RGB')
                elif image.startswith(('http:', 'https:')):
                    response = requests.get(image)
                    image = Image.open(BytesIO(response.content)).convert('RGB')
                else:
                    image = Image.open(image).convert('RGB')

            # Ensure we have a PIL Image at this point
            if not isinstance(image, Image.Image):
                raise ValueError(f"Failed to convert to PIL Image: {type(image)}")

            # Convert PIL to tensor in ComfyUI format
            img_array = np.array(image).astype(np.float32) / 255.0
            img_tensor = torch.from_numpy(img_array)
            
            # Ensure NHWC format
            if img_tensor.dim() == 3:  # [H,W,C]
                img_tensor = img_tensor.unsqueeze(0)  # Add batch dim: [1,H,W,C]
            
            # Create mask
            mask_tensor = torch.ones((1, img_tensor.shape[1], img_tensor.shape[2]), 
                                   dtype=torch.float32)

            return img_tensor, mask_tensor

        except Exception as e:
            logger.error(f"Error processing single image: {str(e)}")
            return load_placeholder_image(placeholder_image_path)

    try:
        # Handle API responses
        if isinstance(images, dict) and response_key in images:
            # Process each item in API response
            all_tensors = []
            all_masks = []
            
            items = images[response_key]
            if isinstance(items, list):
                for item in items:
                    try:
                        img_tensor, mask_tensor = _process_single_image({response_key: [item]})
                        all_tensors.append(img_tensor)
                        all_masks.append(mask_tensor)
                    except Exception as e:
                        logger.error(f"Error processing response item: {str(e)}")
                        continue
                
                if all_tensors:
                    return torch.cat(all_tensors, dim=0), torch.cat(all_masks, dim=0)
            
            # If no valid images processed, return placeholder
            return load_placeholder_image(placeholder_image_path)

        # Handle list/batch of images
        if isinstance(images, (list, tuple)):
            all_tensors = []
            all_masks = []
            
            for img in images:
                try:
                    img_tensor, mask_tensor = _process_single_image(img)
                    all_tensors.append(img_tensor)
                    all_masks.append(mask_tensor)
                except Exception as e:
                    logger.error(f"Error processing batch image: {str(e)}")
                    continue
            
            if all_tensors:
                return torch.cat(all_tensors, dim=0), torch.cat(all_masks, dim=0)
            
            return load_placeholder_image(placeholder_image_path)

        # Handle single image
        return _process_single_image(images)

    except Exception as e:
        logger.error(f"Error in process_images_for_comfy: {str(e)}")
        return _process_single_image(None)

def process_mask(retrieved_mask, image_tensor):
    """
    Process the retrieved_mask to ensure it's in the correct format.
    The mask should be a tensor of shape (B, H, W), matching image_tensor's batch size and dimensions.
    """
    try:
        # Handle torch.Tensor
        if isinstance(retrieved_mask, torch.Tensor):
            # Normalize dimensions
            if retrieved_mask.dim() == 2:  # (H, W)
                retrieved_mask = retrieved_mask.unsqueeze(0)  # Add batch dimension
            elif retrieved_mask.dim() == 3:
                if retrieved_mask.shape[0] != image_tensor.shape[0]:
                    # Adjust batch size
                    retrieved_mask = retrieved_mask.repeat(image_tensor.shape[0], 1, 1)
            elif retrieved_mask.dim() == 4:
                # If mask has a channel dimension, reduce it
                retrieved_mask = retrieved_mask.squeeze(1)
            else:
                raise ValueError(f"Invalid mask tensor dimensions: {retrieved_mask.shape}")

            # Ensure proper format
            retrieved_mask = retrieved_mask.float()
            if retrieved_mask.max() > 1.0:
                retrieved_mask = retrieved_mask / 255.0

            # Ensure mask dimensions match image dimensions
            if retrieved_mask.shape[1:] != image_tensor.shape[2:]:
                # Resize mask to match image dimensions
                retrieved_mask = torch.nn.functional.interpolate(
                    retrieved_mask.unsqueeze(1),
                    size=(image_tensor.shape[2], image_tensor.shape[3]),
                    mode='nearest'
                ).squeeze(1)

            return retrieved_mask

        # Handle PIL Image
        elif isinstance(retrieved_mask, Image.Image):
            mask_array = np.array(retrieved_mask.convert('L')).astype(np.float32) / 255.0
            mask_tensor = torch.from_numpy(mask_array)
            mask_tensor = mask_tensor.unsqueeze(0)  # Add batch dimension

            # Adjust batch size
            if mask_tensor.shape[0] != image_tensor.shape[0]:
                mask_tensor = mask_tensor.repeat(image_tensor.shape[0], 1, 1)

            # Resize if needed
            if mask_tensor.shape[1:] != image_tensor.shape[2:]:
                mask_tensor = torch.nn.functional.interpolate(
                    mask_tensor.unsqueeze(1),
                    size=(image_tensor.shape[2], image_tensor.shape[3]),
                    mode='nearest'
                ).squeeze(1)

            return mask_tensor

        # Handle numpy array
        elif isinstance(retrieved_mask, np.ndarray):
            mask_array = retrieved_mask.astype(np.float32)
            if mask_array.max() > 1.0:
                mask_array = mask_array / 255.0
            if mask_array.ndim == 2:
                pass  # (H, W)
            elif mask_array.ndim == 3:
                mask_array = np.mean(mask_array, axis=2)  # Convert to grayscale
            else:
                raise ValueError(f"Invalid mask array dimensions: {mask_array.shape}")

            mask_tensor = torch.from_numpy(mask_array)
            mask_tensor = mask_tensor.unsqueeze(0)  # Add batch dimension

            # Adjust batch size
            if mask_tensor.shape[0] != image_tensor.shape[0]:
                mask_tensor = mask_tensor.repeat(image_tensor.shape[0], 1, 1)

            # Resize if needed
            if mask_tensor.shape[1:] != image_tensor.shape[2:]:
                mask_tensor = torch.nn.functional.interpolate(
                    mask_tensor.unsqueeze(1),
                    size=(image_tensor.shape[2], image_tensor.shape[3]),
                    mode='nearest'
                ).squeeze(1)

            return mask_tensor

        # Handle other types (e.g., file paths, base64 strings)
        elif isinstance(retrieved_mask, str):
            # Attempt to process as file path or base64 string
            if os.path.exists(retrieved_mask):
                pil_image = Image.open(retrieved_mask).convert('L')
            elif retrieved_mask.startswith('data:image'):
                base64_data = retrieved_mask.split('base64,')[1]
                image_data = base64.b64decode(base64_data)
                pil_image = Image.open(BytesIO(image_data)).convert('L')
            else:
                raise ValueError(f"Invalid mask string: {retrieved_mask}")
            return process_mask(pil_image, image_tensor)

        else:
            raise ValueError(f"Unsupported mask type: {type(retrieved_mask)}")

    except Exception as e:
        logger.error(f"Error processing mask: {str(e)}")
        # Return a default mask matching the image dimensions
        return torch.ones((image_tensor.shape[0], image_tensor.shape[2], image_tensor.shape[3]), dtype=torch.float32)

def convert_mask_to_grayscale_alpha(mask_input):
    """
    Convert mask to grayscale alpha channel.
    Handles tensors, PIL images and numpy arrays.
    Returns tensor in shape [B,1,H,W].
    """
    if isinstance(mask_input, torch.Tensor):
        # Handle tensor input
        if mask_input.dim() == 2:  # [H,W]
            return mask_input.unsqueeze(0).unsqueeze(0)  # Add batch and channel dims
        elif mask_input.dim() == 3:  # [C,H,W] or [B,H,W]
            if mask_input.shape[0] in [1,3,4]:  # Assume channel-first
                if mask_input.shape[0] == 4:  # Use alpha channel
                    return mask_input[3:4].unsqueeze(0)
                else:  # Convert to grayscale
                    weights = torch.tensor([0.299, 0.587, 0.114]).to(mask_input.device)
                    return (mask_input * weights.view(-1,1,1)).sum(0).unsqueeze(0).unsqueeze(0)
        elif mask_input.dim() == 4:  # [B,C,H,W]
            if mask_input.shape[1] == 4:  # Use alpha channel
                return mask_input[:,3:4]
            else:  # Convert to grayscale
                weights = torch.tensor([0.299, 0.587, 0.114]).to(mask_input.device)
                return (mask_input * weights.view(1,-1,1,1)).sum(1).unsqueeze(1)
                
    elif isinstance(mask_input, Image.Image):
        # Convert PIL image to grayscale
        mask = mask_input.convert('L')
        tensor = torch.from_numpy(np.array(mask)).float() / 255.0
        return tensor.unsqueeze(0).unsqueeze(0)  # Add batch and channel dims
        
    elif isinstance(mask_input, np.ndarray):
        # Handle numpy array
        if mask_input.ndim == 2:  # [H,W]
            tensor = torch.from_numpy(mask_input).float()
            return tensor.unsqueeze(0).unsqueeze(0)
        elif mask_input.ndim == 3:  # [H,W,C]
            if mask_input.shape[2] == 4:  # Use alpha channel
                tensor = torch.from_numpy(mask_input[:,:,3]).float()
            else:  # Convert to grayscale
                tensor = torch.from_numpy(np.dot(mask_input[...,:3], [0.299, 0.587, 0.114])).float()
            return tensor.unsqueeze(0).unsqueeze(0)
            
    raise ValueError(f"Unsupported mask input type: {type(mask_input)}")

def tensor_to_base64(tensor: torch.Tensor) -> str:
    """Convert a tensor to a base64-encoded PNG image string."""
    try:
        # Ensure the tensor is in [0, 1] range
        tensor = torch.clamp(tensor, 0, 1)

        # Handle different tensor dimensions
        if tensor.dim() == 3:
            # [C, H, W]
            if tensor.shape[0] == 1:
                # Grayscale image, convert to RGB by repeating channels
                image = tensor.squeeze(0).permute(1, 2, 0).cpu().numpy()  # [H, W, C]
                image = np.repeat(image, 3, axis=2)
            elif tensor.shape[0] == 3:
                # RGB image
                image = tensor.permute(1, 2, 0).cpu().numpy()
            else:
                # Handle tensors with more than 3 channels: select the first 3 channels
                logger.warning(f"Unsupported number of channels: {tensor.shape[0]}. Selecting first 3 channels.")
                if tensor.shape[0] >= 3:
                    image = tensor[:3, :, :].permute(1, 2, 0).cpu().numpy()
                else:
                    raise ValueError(f"Unsupported number of channels: {tensor.shape[0]}")
        elif tensor.dim() == 2:
            # [H, W] Grayscale image
            image = tensor.unsqueeze(-1).cpu().numpy()
            image = np.repeat(image, 3, axis=2)
        else:
            raise ValueError(f"Unsupported tensor shape for conversion: {tensor.shape}")

        # Convert to uint8
        image = (image * 255).astype(np.uint8)

        # Create PIL Image
        pil_image = Image.fromarray(image)

        # Save image to buffer
        buffered = BytesIO()
        pil_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return img_str
    except Exception as e:
        logger.error(f"Error converting tensor to base64: {str(e)}", exc_info=True)
        raise

def tensor_to_pil(tensor):
    """
    Convert a tensor to a PIL image with better error handling and format detection.
    
    Args:
        tensor: A PyTorch tensor representing an image
        
    Returns:
        PIL.Image: The converted PIL image
    """
    try:
        # Ensure tensor is on CPU
        tensor = tensor.cpu()
        
        # Handle different tensor shapes
        if tensor.dim() == 4 and tensor.shape[0] == 1:  # [1, C, H, W] or [1, H, W, C]
            tensor = tensor.squeeze(0)  # Remove batch dimension
            
        # Determine if we have a channels-first or channels-last format
        if tensor.dim() == 3:
            # Handle both [C, H, W] and [H, W, C] formats
            if tensor.shape[0] in [1, 3, 4]:  # Channels-first format [C, H, W]
                tensor = tensor.permute(1, 2, 0)  # Convert to [H, W, C]
                
        # Special case for grayscale
        if tensor.dim() == 2:
            # Add a channel dimension for grayscale [H, W] -> [H, W, 1]
            tensor = tensor.unsqueeze(-1)
            
        # Convert to numpy array
        tensor_np = tensor.numpy()
        
        # Scale to 0-255 range for uint8
        tensor_np = np.clip(tensor_np * 255, 0, 255).astype(np.uint8)
        
        # Create PIL image
        pil_image = Image.fromarray(tensor_np)
        return pil_image
        
    except Exception as e:
        logger.error(f"Error in tensor_to_pil: {e}")
        raise ValueError(f"Failed to convert tensor to PIL image: {e}")

def pil_to_tensor(pil_image):
    # Convert PIL image to tensor
    tensor = torch.from_numpy(np.array(pil_image)).float() / 255.0
    return tensor.permute(2, 0, 1) if tensor.dim() == 3 else tensor.unsqueeze(0)

def base64_to_pil(base64_str):
    """Convert base64 string to PIL Image"""
    if base64_str.startswith('data:image'):
        base64_str = base64_str.split('base64,')[1]
    image_data = base64.b64decode(base64_str)
    return Image.open(BytesIO(image_data))

def pil_image_to_base64(pil_image: Image.Image) -> str:
    """Converts a PIL Image to a data URL."""
    try:
        buffered = io.BytesIO()
        pil_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return f"data:image/png;base64,{img_str}"
    except Exception as e:
        logger.error(f"Error converting image to data URL: {str(e)}", exc_info=True)
        raise

def clean_text(generated_text, remove_weights=True, remove_author=True):
    """Clean text while preserving intentional line breaks."""
    # Split into lines first to preserve breaks
    lines = generated_text.split('\n')
    cleaned_lines = []
    
    for line in lines:
        if line.strip():  # Only process non-empty lines
            # Remove author attribution if requested
            if remove_author:
                line = re.sub(r"\bby:.*", "", line)

            # Remove weights if requested
            if remove_weights:
                line = re.sub(r"\(([^)]*):[\d\.]*\)", r"\1", line)
                line = re.sub(r"(\w+):[\d\.]*(?=[ ,]|$)", r"\1", line)

            # Remove markup tags
            line = re.sub(r"<[^>]*>", "", line)

            # Remove lonely symbols and formatting
            line = re.sub(r"(?<=\s):(?=\s)", "", line)
            line = re.sub(r"(?<=\s);(?=\s)", "", line)
            line = re.sub(r"(?<=\s),(?=\s)", "", line)
            line = re.sub(r"(?<=\s)#(?=\s)", "", line)

            # Clean up extra spaces while preserving line structure
            line = re.sub(r"\s{2,}", " ", line)
            line = re.sub(r"\.,", ",", line)
            line = re.sub(r",,", ",", line)

            # Remove audio tags from the line
            if "<audio" in line:
                print(f"iF_prompt_MKR: Audio has been generated.")
                line = re.sub(r"<audio.*?>.*?</audio>", "", line)

            cleaned_lines.append(line.strip())

    # Join with newlines to preserve line structure
    return "\n".join(cleaned_lines)

def get_api_key(api_key_name, engine):
    """
    Retrieve API key from environment variables or .env file.
    
    Args:
        api_key_name (str): Name of the API key environment variable
        engine (str): Name of the engine being used
        
    Returns:
        str: API key if found and valid
        
    Raises:
        ValueError: If API key is missing or invalid
    """
    local_engines = ["ollama", "llamacpp", "kobold", "lmstudio", "textgen", "sentence_transformers", "transformers"]
    # Try to get the key from .env first
    load_dotenv()
    api_key = os.getenv(api_key_name)
    if engine.lower() in local_engines:
        print(f"You are using {engine} as the engine, no API key is required.")
        return "1234"
    
    # Special handling for HuggingFace
    if engine.lower() == "huggingface":
        # Try both conventional and HF-specific env var names
        api_key = os.getenv("HUGGINGFACE_API_KEY") or os.getenv("HF_AUTH_TOKEN")
        if api_key:
            if validate_huggingface_token(api_key):
                return api_key
            else:
                raise ValueError("Invalid HuggingFace API key")
        raise ValueError("No HuggingFace API key found in environment variables")
    
    elif api_key:
        print(f"API key for {api_key_name} found in .env file or environment variables")
        return api_key
    
    print(f"API key for {api_key_name} not found in .env file or environment variables")
    raise ValueError(f"{api_key_name} not found. Please set it in your .env file or as an environment variable.")

def get_models(engine, base_ip, port, api_key):

    if engine == "ollama":
        api_url = f"http://{base_ip}:{port}/api/tags"
        try:
            response = requests.get(api_url)
            response.raise_for_status()
            models = [model["name"] for model in response.json().get("models", [])]
            return models
        except Exception as e:
            print(f"Failed to fetch models from Ollama: {e}")
            return []
    
    
    elif engine == "huggingface":
        fallback_models = [
            # Vision Language Models (VLM)
            "meta-llama/Llama-3.2-11B-Vision-Instruct",
            "Qwen/Qwen2-VL-7B-Chat",
            "Qwen/Qwen2-VL-7B",
            "Qwen/Qwen2-VL-2B-Chat",
            "Qwen/Qwen2-VL-2B",
            "Qwen/Qwen2-VL-7B-Instruct",
            "Qwen/Qwen2-VL-2B-Instruct",
            "microsoft/phi-2",
            "HuggingFaceH4/zephyr-7b-beta",
            
            # Text to Image Models
            "stabilityai/sdxl-turbo",
            "stabilityai/stable-diffusion-xl-base-1.0",
            "stabilityai/stable-diffusion-2-1",
            "runwayml/stable-diffusion-v1-5",
            "CompVis/stable-diffusion-v1-4",
            "stabilityai/stable-diffusion-3-base",
            "stabilityai/stable-diffusion-3-medium",
            "stabilityai/stable-diffusion-3-small",
            "black-forest-labs/FLUX.1-dev",
            "playgroundai/playground-v2-256px",
            "playgroundai/playground-v2-1024px",
            
            # Image to Image Models
            "timbrooks/instruct-pix2pix",
            "lambdalabs/sd-image-variations-diffusers",
            "diffusers/controlnet-canny-sdxl-1.0",
            
            # Specialized Models
            "kandinsky-community/kandinsky-3",
            "stabilityai/stable-cascade",
            "dataautogpt3/OpenDalle3",
            "ByteDance/SDXL-Lightning",
            
            # ControlNet Models
            "lllyasviel/control_v11p_sd15_canny",
            "lllyasviel/control_v11p_sd15_openpose",
            "lllyasviel/control_v11p_sd15_depth",
            
            # Text Feature Extraction
            "sentence-transformers/all-MiniLM-L6-v2",
            "sentence-transformers/all-mpnet-base-v2",
            
            # Image Feature Extraction  
            "openai/clip-vit-base-patch32",
            "openai/clip-vit-large-patch14",
            
            # Text Classification
            "distilbert-base-uncased-finetuned-sst-2-english",
            "roberta-base-openai-detector",
            
            # Text Generation
            "gpt2",
            "facebook/opt-350m",
            
            # Translation 
            "Helsinki-NLP/opus-mt-en-fr",
            "Helsinki-NLP/opus-mt-fr-en",
            
            # Question Answering
            "deepset/roberta-base-squad2",
            "distilbert-base-cased-distilled-squad"
        ]

        try:
            # Verify API key
            if not api_key or api_key == "1234":
                print("No valid HuggingFace API key provided. Using fallback models.")
                return fallback_models

            headers = {
                "Authorization": f"Bearer {api_key}",
                "Accept": "application/json"
            }

            # Check inference API endpoint directly
            inference_url = "https://api-inference.huggingface.co/status"
            response = requests.get(inference_url, headers=headers)
            
            if response.status_code != 200:
                print("Failed to verify HuggingFace Inference API access. Using fallback models.")
                return fallback_models

            # Get models available for inference API
            models_url = "https://api-inference.huggingface.co/framework/all"
            response = requests.get(models_url, headers=headers)

            if response.status_code == 200:
                api_models = []
                data = response.json()
                
                # Extract models supporting inference API
                for framework in data:
                    for model in framework.get("models", []):
                        model_id = model.get("model_id")
                        if model_id:
                            api_models.append(model_id)
                
                # Combine with fallback models and remove duplicates
                combined_models = list(dict.fromkeys(api_models + fallback_models))
                return combined_models
            else:
                print(f"Failed to fetch inference models. Status code: {response.status_code}")
                return fallback_models

        except Exception as e:
            print(f"Error fetching HuggingFace models: {str(e)}")
            return fallback_models

    elif engine == "deepseek":
        fallback_models = [
            "deepseek-reasoner",
            "deepseek-chat",
            "deepseek-coder"
        ]

        #api_key = get_api_key("DEEPSEEK_API_KEY", engine)
        if not api_key or api_key == "1234":
            print("Warning: Invalid DeepSeek API key. Using fallback model list.")
            return fallback_models

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        api_url = "https://api.deepseek.com/v1/models"  # Adjust URL if needed
        try:
            response = requests.get(api_url, headers=headers)
            response.raise_for_status()
            api_models = [model["id"] for model in response.json()["data"]]
            print(f"Successfully fetched {len(api_models)} models from DeepSeek API")
            
            # Combine API models with fallback models, prioritizing API models
            combined_models = list(set(api_models + fallback_models))
            return combined_models
        except Exception as e:
            print(f"Failed to fetch models from DeepSeek: {e}")
            print(f"Returning fallback list of {len(fallback_models)} DeepSeek models")
            return fallback_models

    elif engine == "lmstudio":
        api_url = f"http://{base_ip}:{port}/v1/models"
        try:
            print(f"Attempting to connect to {api_url}")
            response = requests.get(api_url, timeout=10)
            print(f"Response status code: {response.status_code}")
            print(f"Response content: {response.text}")
            if response.status_code == 200:
                data = response.json()
                models = [model["id"] for model in data["data"]]
                return models
            else:
                print(f"Failed to fetch models from LM Studio. Status code: {response.status_code}")
                return []
        except requests.exceptions.RequestException as e:
            print(f"Error connecting to LM Studio server: {e}")
            return []

    elif engine == "textgen":
        api_url = f"http://{base_ip}:{port}/v1/internal/model/list"
        try:
            response = requests.get(api_url)
            response.raise_for_status()
            models = response.json()["model_names"]
            return models
        except Exception as e:
            print(f"Failed to fetch models from text-generation-webui: {e}")
            return []

    elif engine == "kobold":
        api_url = f"http://{base_ip}:{port}/api/v1/model"
        try:
            response = requests.get(api_url)
            response.raise_for_status()
            model = response.json()["result"]
            return [model]
        except Exception as e:
            print(f"Failed to fetch models from Kobold: {e}")
            return []

    elif engine == "llamacpp":
        api_url = f"http://{base_ip}:{port}/v1/models"
        try:
            response = requests.get(api_url)
            response.raise_for_status()
            models = [model["id"] for model in response.json()["data"]]
            return models
        except Exception as e:
            print(f"Failed to fetch models from llama.cpp: {e}")
            return []

    elif engine == "vllm":
        api_url = f"http://{base_ip}:{port}/v1/models"
        try:
            response = requests.get(api_url)
            response.raise_for_status()
            # Adapt this based on vLLM"s actual API response structure
            models = [model["id"] for model in response.json()["data"]] 
            return models
        except Exception as e:
            print(f"Failed to fetch models from vLLM: {e}")
            return []

    elif engine == "openai":
        fallback_models = [
            # GPT-4o Models
            "gpt-4o",
            "gpt-4o-2024-05-13",
            "gpt-4o-2024-08-06",
            "gpt-4o-2024-11-20",
            "gpt-4o-audio-preview",
            "gpt-4o-audio-preview-2024-10-01",
            "gpt-4o-audio-preview-2024-12-17",
            "gpt-4o-mini",
            "gpt-4o-mini-2024-07-18",
            "gpt-4o-mini-audio-preview",
            "gpt-4o-mini-audio-preview-2024-12-17",
            "gpt-4o-mini-realtime-preview",
            "gpt-4o-mini-realtime-preview-2024-12-17",
            "gpt-4o-realtime-preview",
            "gpt-4o-realtime-preview-2024-10-01",
            "gpt-4o-realtime-preview-2024-12-17",

            # GPT-4 Models
            "gpt-4",
            "gpt-4-0125-preview",
            "gpt-4-0613",
            "gpt-4-1106-preview",
            "gpt-4-turbo",
            "gpt-4-turbo-2024-04-09",
            "gpt-4-turbo-preview",

            # GPT-3.5 Models
            "gpt-3.5-turbo",
            "gpt-3.5-turbo-0125",
            "gpt-3.5-turbo-1106",
            "gpt-3.5-turbo-16k",
            "gpt-3.5-turbo-instruct",
            "gpt-3.5-turbo-instruct-0914",

            # DALL-E Models
            "dall-e-2",
            "dall-e-3",

            # Whisper Models
            "whisper-1",
            "whisper-I",

            # TTS Models
            "tts-1",
            "tts-1-1106",
            "tts-1-hd",
            "tts-1-hd-1106",
            "tts-l-hd",

            # Embedding Models
            "text-embedding-3-large",
            "text-embedding-3-small",
            "text-embedding-ada-002",

            # Specialized Models
            "babbage-002",
            "chatgpt-4o-latest",
            "davinci-002",
            "gpt40-0806-loco-vm",

            # O1 Models
            "o1",
            "o1-mini",
            "o1-mini-2024-09-12",
            "o1-preview",
            "o1-preview-2024-09-12",

            # Omni Moderation
            "omni-moderation-2024-09-26",
            "omni-moderation-latest",

            # Future/Experimental
            "gpt-4.5-preview",
            "gpt-4.5-preview-2025-02-27",
            "o3-mini",
            "o3-mini-2025-01-31"
        ]

        #api_key = get_api_key("OPENAI_API_KEY", engine)
        if not api_key or api_key == "1234":
            print("Warning: Invalid OpenAI API key. Using fallback model list.")
            return fallback_models

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        api_url = "https://api.openai.com/v1/models"
        try:
            response = requests.get(api_url, headers=headers)
            response.raise_for_status()
            api_models = [model["id"] for model in response.json()["data"]]
            print(f"Successfully fetched {len(api_models)} models from OpenAI API")
            
            # Combine API models with fallback models, prioritizing API models
            combined_models = list(set(api_models + fallback_models))
            return combined_models
        except Exception as e:
            print(f"Failed to fetch models from OpenAI: {e}")
            if isinstance(e, requests.exceptions.RequestException) and hasattr(e, "response"):
                print(f"Response status code: {e.response.status_code}")
                print(f"Response content: {e.response.text}")
            print(f"Returning fallback list of {len(fallback_models)} OpenAI models")
            return fallback_models
    
    elif engine == "xai":
        fallback_models = [
            "grok-2",
            "grok-2-1212",
            "grok-2-latest",
            "grok-2-vision",
            "grok-2-vision-1212",
            "grok-2-vision-latest",
            "grok-beta",
            "grok-vision-beta"
        ]

        #api_key = get_api_key("XAI_API_KEY", engine)
        if not api_key or api_key == "1234":
            print("Warning: Invalid XAI API key. Using fallback model list.")
            return fallback_models

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        api_url = "https://api.x.ai/v1/models"
        try:
            response = requests.get(api_url, headers=headers)
            response.raise_for_status()
            api_models = [model["id"] for model in response.json()["data"]]
            print(f"Successfully fetched {len(api_models)} models from XAI API")
            
            # Combine API models with fallback models, prioritizing API models
            combined_models = list(set(api_models + fallback_models))
            return combined_models
        except Exception as e:
            print(f"Failed to fetch models from XAI: {e}")
            if isinstance(e, requests.exceptions.RequestException) and hasattr(e, "response"):
                print(f"Response status code: {e.response.status_code}")
                print(f"Response content: {e.response.text}")
            print(f"Returning fallback list of {len(fallback_models)} XAI models")
            return fallback_models

    elif engine == "mistral":
        fallback_models = [
            "codestral-2405",
            "codestral-2411-rc5",
            "codestral-2412",
            "codestral-2501",
            "codestral-latest",
            "codestral-mamba-2407",
            "codestral-mamba-latest",
            "ministral-3b-2410",
            "ministral-3b-latest",
            "ministral-8b-2410",
            "ministral-8b-latest",
            "mistral-embed",
            "mistral-large-2402",
            "mistral-large-2407",
            "mistral-large-2411",
            "mistral-large-latest",
            "mistral-large-pixtral-2411",
            "mistral-medium",
            "mistral-medium-2312",
            "mistral-medium-latest",
            "mistral-moderation-2411",
            "mistral-moderation-latest",
            "mistral-ocr-2503",
            "mistral-ocr-latest",
            "mistral-saba-2502",
            "mistral-saba-latest",
            "mistral-small",
            "mistral-small-2312",
            "mistral-small-2402",
            "mistral-small-2409",
            "mistral-small-2501",
            "mistral-small-latest",
            "mistral-tiny",
            "mistral-tiny-2312",
            "mistral-tiny-2407",
            "mistral-tiny-latest",
            "open-codestral-mamba",
            "open-mistral-7b",
            "open-mistral-nemo",
            "open-mistral-nemo-2407",
            "open-mixtral-8x22b",
            "open-mixtral-8x22b-2404",
            "open-mixtral-8x7b",
            "pixtral-12b",
            "pixtral-12b-2409",
            "pixtral-12b-latest",
            "pixtral-large-2411",
            "pixtral-large-latest"
        ]

        #api_key = get_api_key("MISTRAL_API_KEY", engine)
        if not api_key or api_key == "1234":
            print("Warning: Invalid Mistral API key. Using fallback model list.")
            return fallback_models

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        api_url = "https://api.mistral.ai/v1/models"
        try:
            response = requests.get(api_url, headers=headers)
            response.raise_for_status()
            api_models = [model["id"] for model in response.json()["data"]]
            print(f"Successfully fetched {len(api_models)} models from Mistral API")
            
            # Combine API models with fallback models, prioritizing API models
            combined_models = list(set(api_models + fallback_models))
            return combined_models
        except Exception as e:
            print(f"Failed to fetch models from Mistral: {e}")
            print(f"Returning fallback list of {len(fallback_models)} Mistral models")
            return fallback_models

    elif engine == "groq":
        fallback_models = [
            "deepseek-r1-distill-llama-70b",
            "deepseek-r1-distill-qwen-32b",
            "distil-whisper-large-v3-en",
            "gemma2-9b-it",
            "llama-guard-3-8b",
            "llama-3.1-70b-versatile",
            "llama-3.1-8b-instant",
            "llama-3.2-1b-preview",
            "llama-3.2-3b-preview",
            "llama-3.2-11b-vision-preview",
            "llama-3.2-90b-vision-preview",
            "llama-3.3-70b-specdec",
            "llama-3.3-70b-versatile",
            "llama3-8b-8192",
            "llama3-70b-8192",
            "llama3-groq-8b-8192-tool-use-preview",
            "llama3-groq-70b-8192-tool-use-preview",
            "llava-v1.5-7b-4096-preview",
            "mixtral-8x7b-32768",
            "mistral-saba-24b",
            "qwen-2.5-32b",
            "qwen-2.5-coder-32b",
            "qwen-qwq-32b",
            "whisper-large-v3",
            "whisper-large-v3-turbo"
        ]

        #api_key = get_api_key("GROQ_API_KEY", engine)
        if not api_key or api_key == "1234":
            print("Warning: Invalid GROQ API key. Using fallback model list.")
            return fallback_models

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        api_url = "https://api.groq.com/openai/v1/models"
        try:
            response = requests.get(api_url, headers=headers)
            response.raise_for_status()
            api_models = [model["id"] for model in response.json()["data"]]
            print(f"Successfully fetched {len(api_models)} models from GROQ API")
            
            # Combine API models with fallback models, prioritizing API models
            combined_models = list(set(api_models + fallback_models))
            return combined_models
        except Exception as e:
            print(f"Failed to fetch models from GROQ: {e}")
            print(f"Returning fallback list of {len(fallback_models)} GROQ models")
            return fallback_models

    elif engine == "anthropic":
        return [
            "claude-3-5-opus-latest",
            "claude-3-opus-20240229",
            "claude-3-5-sonnet-latest",
            "claude-3-5-sonnet-20240620",
            "claude-3-sonnet-20240229",
            "claude-3-haiku-20240307",
            "claude-3-5-haiku-latest",
            "claude-3-5-haiku-20241022"
        ]

    elif engine == "gemini":
        return [
            "learnlrn-1.5-pro-experimental",
            "gemini-2.0-flash-thinking-exp-1219",
            "gemini-2.0-flash-exp",
            "gemini-exp-1206",
            "gemini-exp-1121",
            "gemini-exp-1114",
            "gemini-1.5-pro-002",
            "gemini-1.5-flash-002",
            "gemini-1.5-flash-8b-exp-0924",
            "gemini-1.5-flash-latest",
            "gemini-1.5-flash",
            "gemini-1.5-pro-latest",
            "gemini-1.5-latest",
            "gemini-pro",
            "gemini-pro-vision",
        ]

    elif engine == "sentence_transformers":
        return [
            "sentence-transformers/all-MiniLM-L6-v2",
            "avsolatorio/GIST-small-Embedding-v0",
        ]

    elif engine == "transformers":
        # Standard list of transformers models to show
        fallback_models = [
            "Qwen/Qwen2.5-VL-3B-Instruct-AWQ",  # Default model we want to use
            "Qwen/Qwen2.5-VL-7B-Instruct-AWQ",
            "Qwen/QwQ-32B-AWQ",  # Keep QwQ-32B-AWQ model
            "Qwen/Qwen2.5-VL-3B-Instruct",
            "Qwen/Qwen2.5-VL-7B-Instruct",
            "Qwen/Qwen2.5-7B-Instruct",
            "Qwen/Qwen2-7B-Instruct",
            "Qwen/Qwen2-VL-7B-Instruct",
            "Qwen/Qwen2-72B-Instruct"
        ]
        
        # Check if we have a transformers model manager to list models
        try:
            from transformers_api import _transformers_manager
            
            # Get list of models from LLM directory
            try:
                # Get models directory dynamically
                try:
                    import folder_paths
                    models_dir = folder_paths.models_dir
                except (ImportError, AttributeError):
                    # Fallback to a default location if folder_paths is not available
                    models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models")
                    os.makedirs(models_dir, exist_ok=True)
                    print(f"Could not import folder_paths.models_dir, using fallback: {models_dir}")
                
                llm_path = os.path.join(models_dir, "LLM")
                if os.path.exists(llm_path) and os.path.isdir(llm_path):
                    # List directories in the LLM folder
                    local_models = []
                    for model_dir in os.listdir(llm_path):
                        model_path = os.path.join(llm_path, model_dir)
                        if os.path.isdir(model_path):
                            # Check if it has config.json to verify it's a model
                            if os.path.exists(os.path.join(model_path, "config.json")):
                                if "/" not in model_dir and "\\" not in model_dir:
                                    # For non-namespaced models, use the directory name
                                    local_models.append(model_dir)
                                else:
                                    # For models with namespaces, keep the structure
                                    local_models.append(model_dir)
                    
                    # If we found local models, add them to our list
                    if local_models:
                        print(f"Found {len(local_models)} local transformers models")
                        # Combine local models with fallback models (local models first)
                        combined_models = list(dict.fromkeys(local_models + fallback_models))
                        return combined_models
            except Exception as e:
                print(f"Error scanning local models directory: {e}")
            
            # If we couldn't find local models, return fallback list
            return fallback_models
        except ImportError:
            print("TransformersModelManager not available, using fallback models list")
            return fallback_models

    else:
        print(f"Unsupported engine - {engine}")
        return []

def validate_models(model, provider, model_type, base_ip, port, api_key):
        available_models = get_models(provider, base_ip, port, api_key)
        if available_models is None or model not in available_models:
            error_message = f"Invalid {model_type} model selected: {model} for provider {provider}. Available models: {available_models}"
            print(error_message)
            raise ValueError(error_message)

class EnhancedYAMLDumper(yaml.SafeDumper):
    def increase_indent(self, flow=False, indentless=False):
        return super(EnhancedYAMLDumper, self).increase_indent(flow, False)

def str_presenter(dumper, data):
    if len(data.splitlines()) > 1:  # check for multiline string
        return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
    return dumper.represent_scalar("tag:yaml.org,2002:str", data)

EnhancedYAMLDumper.add_representer(str, str_presenter)


def validate_huggingface_token(api_key):
    """Validate HuggingFace API token"""
    try:
        headers = {"Authorization": f"Bearer {api_key}"}
        # Try to access the API with the token
        response = requests.get(
            "https://huggingface.co/api/whoami",
            headers=headers
        )
        return response.status_code == 200
    except Exception as e:
        logger.error(f"Error validating HuggingFace token: {e}")
        return False

def get_huggingface_url(model_or_url):
    """Convert model name to full HuggingFace API URL if needed"""
    if model_or_url.startswith(('http://', 'https://')):
        return model_or_url
    return f'https://api-inference.huggingface.co/models/{model_or_url}'

def send_huggingface_request(endpoint, payload, api_key, max_retries=3):
    """Send request to HuggingFace Inference API with retry logic"""
    headers = {"Authorization": f"Bearer {api_key}"}
    url = get_huggingface_url(endpoint)
    
    for attempt in range(max_retries):
        try:
            response = requests.post(url, headers=headers, json=payload)
            
            if response.status_code == 200:
                return response
                
            elif 'estimated_time' in response.text:
                # Handle model loading
                estimated_time = response.json().get('estimated_time', 30)
                logger.info(f"Model loading, waiting {estimated_time} seconds...")
                time.sleep(estimated_time)
                continue
                
            else:
                raise Exception(f"HuggingFace API error: {response.text}")
                
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            logger.warning(f"Retry {attempt + 1}/{max_retries} after error: {e}")
            time.sleep(2 ** attempt)  # Exponential backoff

def numpy_int64_presenter(dumper, data):
    return dumper.represent_int(int(data))

EnhancedYAMLDumper.add_representer(np.int64, numpy_int64_presenter)

def dump_yaml(data, file_path):
    """
    Safely dumps a dictionary to a YAML file with custom formatting.
    Converts any numpy.int64 values to int to avoid YAML serialization errors.
    Uses multi-line string representation for better readability.
    """
    def convert_numpy_types(obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return obj

    # Convert numpy types in the entire data structure
    data = yaml.safe_load(yaml.dump(data, default_flow_style=False, allow_unicode=True))

    with open(file_path, "w") as yaml_file:
        yaml.dump(data, yaml_file, Dumper=EnhancedYAMLDumper, default_flow_style=False, 
                  sort_keys=False, allow_unicode=True, width=1000, indent=2)


def save_combo_settings(settings_dict, combo_presets_dir):
    """Save combo settings to the AutoCombo directory."""
    try:
        os.makedirs(combo_presets_dir, exist_ok=True)
        settings_path = os.path.join(combo_presets_dir, 'combo_settings.yaml')
        
        with open(settings_path, 'w') as f:
            yaml.safe_dump(settings_dict, f)
        logger.info(f"Saved combo settings to {settings_path}")
        return settings_dict
    except Exception as e:
        logger.error(f"Error saving combo settings: {str(e)}")
        return None

def load_combo_settings(combo_presets_dir):
    """Load combo settings from the AutoCombo directory."""
    try:
        settings_path = os.path.join(combo_presets_dir, 'combo_settings.yaml')
        
        if os.path.exists(settings_path):
            with open(settings_path, 'r') as f:
                settings = yaml.safe_load(f)
                logger.info(f"Loaded combo settings from {settings_path}")
                return settings
        else:
            logger.warning(f"Combo settings file not found at {settings_path}")
            return {}
    except Exception as e:
        logger.error(f"Error loading combo settings: {str(e)}")
        return {}

def create_settings_from_ui(ui_settings):
    """
    Create settings.yaml from UI settings with proper type conversion.
    Handles UI values that may be boolean or string.
    """
    import json

    def convert_to_bool(value):
        if isinstance(value, bool):
            return value
        if isinstance(value, str):
            return value.lower() == 'true'
        return bool(value)

    # Load profiles
    profiles_path = os.path.join(
        folder_paths.base_path,
        "custom_nodes",
        "ComfyUI-IF_LLM",
        "IF_AI",
        "presets",
        "profiles.json"
    )
    
    with open(profiles_path, 'r') as f:
        profiles = json.load(f)

    profile_name = ui_settings.get('profile', 'IF_PromptMKR')
    profile_content = profiles.get(profile_name, {}).get('instruction', '')

    # If 'prime_directives' is empty, use the profile content
    prime_directives = ui_settings.get('prime_directives')
    if not prime_directives or prime_directives in (None, '', 'None'):
        prime_directives = profile_content

    settings = {
        'base_ip': str(ui_settings.get('base_ip', 'localhost')),
        'port': str(ui_settings.get('port', '11434')),
        'user_prompt': str(ui_settings.get('user_prompt', 'Who helped Safiro infiltrate the Zaltar Organisation?')),
        'llm_provider': str(ui_settings.get('llm_provider', 'ollama')),
        'llm_model': str(ui_settings.get('llm_model', 'llama3.1:latest')),
        'prime_directives': prime_directives,
        'temperature': float(ui_settings.get('temperature', 0.7)),
        'max_tokens': int(ui_settings.get('max_tokens', 2048)),
        'stop_string': None if ui_settings.get('stop_string') in (None, 'None') else str(ui_settings.get('stop_string')),
        'keep_alive': convert_to_bool(ui_settings.get('keep_alive', False)),
        'clear_history': convert_to_bool(ui_settings.get('clear_history', False)),
        'history_steps': int(ui_settings.get('history_steps', 10)),
        'top_k': int(ui_settings.get('top_k', 40)),
        'top_p': float(ui_settings.get('top_p', 0.9)),
        'repeat_penalty': float(ui_settings.get('repeat_penalty', 1.2)),
        'seed': None if ui_settings.get('seed') in (None, 'None') else int(ui_settings.get('seed')),
        'external_api_key': str(ui_settings.get('external_api_key', '')),
        'random': convert_to_bool(ui_settings.get('random', False)),
        'aspect_ratio': str(ui_settings.get('aspect_ratio', '16:9')),
        'auto_combo': convert_to_bool(ui_settings.get('auto_combo', False)),
        'precision': str(ui_settings.get('precision', 'fp16')),
        'attention': str(ui_settings.get('attention', 'sdpa')),
        'batch_count': int(ui_settings.get('batch_count', 4)),
        'strategy': str(ui_settings.get('strategy', 'normal')),
        'profile': profile_name  # Include profile name
    }
    return settings

def format_response(self, response):
        """
        Format the response by adding appropriate line breaks and paragraph separations.
        """
        paragraphs = re.split(r"\n{2,}", response)

        formatted_paragraphs = []
        for para in paragraphs:
            if "```" in para:
                parts = para.split("```")
                for i, part in enumerate(parts):
                    if i % 2 == 1:  # This is a code block
                        parts[i] = f"\n```\n{part.strip()}\n```\n"
                para = "".join(parts)
            else:
                para = para.replace(". ", ".\n")

            formatted_paragraphs.append(para.strip())

        return "\n\n".join(formatted_paragraphs)

def print_available_models():
    """Print available models for each supported API engine"""
    
    # Test API key - using a dummy value since we'll mostly see fallback models
    test_api_key = "1234"
    
    # List of all supported engines
    engines = [
        "ollama",
        "huggingface",
        "deepseek", 
        "lmstudio",
        "textgen",
        "kobold",
        "llamacpp",
        "vllm",
        "openai",
        "xai",
        "mistral",
        "groq",
        "anthropic",
        "gemini",
        "sentence_transformers",
        "transformers"
    ]
    
    print("\n=== Available Models by Engine ===\n")
    
    for engine in engines:
        print(f"\n{engine.upper()} Models:")
        print("-" * (len(engine) + 8))
        
        try:
            # Get models for the current engine
            models = get_models(engine, "localhost", "11434", test_api_key)
            
            if models:
                # Print each model with an index
                for i, model in enumerate(models, 1):
                    print(f"{i}. {model}")
            else:
                print("No models available or engine requires valid API key/connection")
                
        except Exception as e:
            print(f"Error fetching models: {str(e)}")
            
        print() # Add blank line between engines

# Usage example:
if __name__ == "__main__":
    print_available_models()

def gemini2_process_images(images, max_input_images=5, target_size=(768, 768)):
    """
    Process a batch of images for Gemini 2.0 API.
    
    Args:
        images (torch.Tensor or list): Image batch in ComfyUI format [B,H,W,C] or list of tensors
        max_input_images (int): Maximum number of images to include (Gemini may have limits)
        target_size (tuple): Target size for images (width, height)
        
    Returns:
        list: List of processed PIL images ready for the Gemini API
    """
    import torch
    from PIL import Image
    import numpy as np
    
    # Handle different input types
    processed_images = []
    
    if isinstance(images, torch.Tensor):
        # Handle 4D tensor [B,H,W,C]
        if images.dim() == 4:
            # Limit to max_input_images
            batch_size = min(images.shape[0], max_input_images)
            
            for i in range(batch_size):
                # Get single image tensor [H,W,C]
                img_tensor = images[i].cpu()
                
                # Convert to numpy and scale to 0-255
                img_np = (img_tensor.numpy() * 255).clip(0, 255).astype(np.uint8)
                
                # Convert to PIL
                pil_img = Image.fromarray(img_np)
                
                # Resize to target size if needed
                if pil_img.size != target_size:
                    pil_img = pil_img.resize(target_size, Image.Resampling.LANCZOS)
                
                processed_images.append(pil_img)
                
        # Handle 3D tensor [H,W,C]
        elif images.dim() == 3:
            img_tensor = images.cpu()
            img_np = (img_tensor.numpy() * 255).clip(0, 255).astype(np.uint8)
            pil_img = Image.fromarray(img_np)
            
            if pil_img.size != target_size:
                pil_img = pil_img.resize(target_size, Image.Resampling.LANCZOS)
                
            processed_images.append(pil_img)
    
    # Handle list of tensors
    elif isinstance(images, list):
        # Limit to max_input_images
        num_images = min(len(images), max_input_images)
        
        for i in range(num_images):
            img = images[i]
            
            if isinstance(img, torch.Tensor):
                img_tensor = img.cpu()
                
                # Handle different tensor dimensions
                if img_tensor.dim() == 4 and img_tensor.shape[0] == 1:  # [1,H,W,C]
                    img_tensor = img_tensor.squeeze(0)
                
                img_np = (img_tensor.numpy() * 255).clip(0, 255).astype(np.uint8)
                pil_img = Image.fromarray(img_np)
                
                if pil_img.size != target_size:
                    pil_img = pil_img.resize(target_size, Image.Resampling.LANCZOS)
                    
                processed_images.append(pil_img)
    
    return processed_images

def gemini2_prepare_response(response, width=512, height=512):
    """
    Extract and prepare images from Gemini 2.0 API response.
    
    Args:
        response: Gemini API response object
        width (int): Target width for extracted images
        height (int): Target height for extracted images
        
    Returns:
        tuple: (list of image binaries, response text)
    """
    from io import BytesIO
    
    images = []
    response_text = ""
    
    # Handle empty response
    if not response or not hasattr(response, 'candidates') or not response.candidates:
        return images, "No response generated"
    
    # Process each candidate
    for candidate in response.candidates:
        if not hasattr(candidate, 'content') or not hasattr(candidate.content, 'parts'):
            continue
            
        for part in candidate.content.parts:
            # Process text parts
            if hasattr(part, 'text') and part.text:
                response_text += part.text + "\n"
                
            # Process image parts
            if hasattr(part, 'inline_data') and part.inline_data:
                try:
                    # Get binary image data
                    image_binary = part.inline_data.data
                    images.append(image_binary)
                except Exception as e:
                    print(f"Error extracting image from response: {e}")
    
    return images, response_text

def gemini2_create_client(api_key):
    """
    Create and return a Gemini API client.
    
    Args:
        api_key (str): The Gemini API key
        
    Returns:
        Client: Gemini API client object
    """
    try:
        from google import genai
        client = genai.Client(api_key=api_key)
        return client
    except ImportError:
        raise ImportError("The google-generativeai package is required. Please install it with: pip install google-generativeai")
    except Exception as e:
        raise RuntimeError(f"Failed to create Gemini client: {str(e)}")

def validate_gemini_key(api_key):
    """
    Validate a Gemini API key by making a simple test request.
    
    Args:
        api_key (str): The Gemini API key to validate
        
    Returns:
        bool: True if key is valid, False otherwise
    """
    try:
        from google import genai
        
        # Initialize client with the key
        client = genai.Client(api_key=api_key)
        
        # Try a simple models list request
        models = client.models.list()
        
        # If we get here, the key is valid
        return True
    except Exception as e:
        print(f"Invalid Gemini API key: {str(e)}")
        return False