"""
Module for extracting metadata from image files, particularly focusing on
Stable Diffusion metadata from PNG files.
"""

import io
from PIL import Image, PngImagePlugin
import re


class MetadataExtractor:
    """Class for extracting and parsing metadata from images."""
    
    @staticmethod
    def extract_metadata(image_path):
        """
        Extract metadata from an image file.
        
        Args:
            image_path: path to the image file
            
        Returns:
            dict: dictionary with extracted metadata
        """
        try:
            # Open image with PIL
            image = Image.open(image_path)
            
            # Extract metadata from PNG info
            metadata_text = image.info.get("parameters", "")
            
            # Parse the metadata
            parsed_metadata = MetadataExtractor.parse_metadata(metadata_text)
            
            # Add basic image info
            parsed_metadata.update({
                'width': image.width,
                'height': image.height,
                'format': image.format,
                'mode': image.mode,
            })
            
            return parsed_metadata
        except Exception as e:
            print(f"Error extracting metadata from {image_path}: {e}")
            return {'error': str(e)}
    
    @staticmethod
    def parse_metadata(metadata_text):
        """
        Parse Stable Diffusion metadata text into structured data.
        
        Args:
            metadata_text: raw metadata text from image
            
        Returns:
            dict: structured metadata
        """
        if not metadata_text:
            return {'raw_text': ''}
        
        result = {'raw_text': metadata_text}
        
        # Extract prompt
        prompt_end = metadata_text.find("Negative prompt:")
        if prompt_end > 0:
            result['prompt'] = metadata_text[:prompt_end].strip()
            negative_prompt_end = metadata_text.find("\n", prompt_end)
            if negative_prompt_end > 0:
                result['negative_prompt'] = metadata_text[prompt_end + len("Negative prompt:"):negative_prompt_end].strip()
        else:
            result['prompt'] = metadata_text.strip()
        
        # Extract model name
        model_match = re.search(r'Model: ([^,\n]+)', metadata_text)
        if model_match:
            result['model'] = model_match.group(1).strip()
        
        # Extract other parameters
        params = {
            'steps': r'Steps: (\d+)',
            'sampler': r'Sampler: ([^,\n]+)',
            'cfg_scale': r'CFG scale: ([^,\n]+)',
            'seed': r'Seed: ([^,\n]+)',
            'size': r'Size: ([^,\n]+)',
            'model_hash': r'Model hash: ([^,\n]+)',
        }
        
        for key, pattern in params.items():
            match = re.search(pattern, metadata_text)
            if match:
                result[key] = match.group(1).strip()
        
        return result
    
    @staticmethod
    def group_images_by_model(metadata_list):
        """
        Group images by model name.
        
        Args:
            metadata_list: list of (image_path, metadata) tuples
            
        Returns:
            dict: dictionary with model names as keys and lists of image paths as values
        """
        result = {}
        
        for image_path, metadata in metadata_list:
            model = metadata.get('model', 'unknown')
            if model not in result:
                result[model] = []
            result[model].append(image_path)
        
        return result
    
    @staticmethod
    def group_images_by_prompt(metadata_list):
        """
        Group images by prompt.
        
        Args:
            metadata_list: list of (image_path, metadata) tuples
            
        Returns:
            dict: dictionary with prompts as keys and lists of image paths as values
        """
        result = {}
        
        for image_path, metadata in metadata_list:
            prompt = metadata.get('prompt', 'unknown')
            # Use first 50 chars as key to avoid extremely long keys
            prompt_key = prompt[:50] + ('...' if len(prompt) > 50 else '')
            if prompt_key not in result:
                result[prompt_key] = []
            result[prompt_key].append((image_path, metadata.get('model', 'unknown')))
        
        return result
    
    @staticmethod
    def update_metadata(image_path, new_metadata, output_path=None):
        """
        Update metadata in an image file.
        
        Args:
            image_path: path to the input image file
            new_metadata: new metadata text to write
            output_path: path to save the updated image (if None, overwrites input)
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Open image with PIL
            image = Image.open(image_path)
            
            # Create a PngInfo object to store metadata
            pnginfo = PngImagePlugin.PngInfo()
            pnginfo.add_text("parameters", new_metadata)
            
            # Save the image with the updated metadata
            save_path = output_path if output_path else image_path
            image.save(save_path, format="PNG", pnginfo=pnginfo)
            
            return True
        except Exception as e:
            print(f"Error updating metadata: {e}")
            return False