#!/usr/bin/env python3
"""
Image Integration Module

Integrates with image-generation and image-edit skills.
Handles product vs non-product image workflows.
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path
from typing import Optional, List


class ImageIntegration:
    """Integrate with image-generation and image-edit skills"""
    
    def __init__(self, skills_base_path: str = None):
        """
        Initialize image integration
        
        Args:
            skills_base_path: Base path to skills directory
        """
        if skills_base_path is None:
            # Default: assume we're in skills/seo-multi-channel/scripts/
            base = Path(__file__).parent.parent.parent
            self.skills_base = str(base)
        else:
            self.skills_base = skills_base
        
        self.image_gen_script = os.path.join(self.skills_base, 'image-generation/scripts/image_gen.py')
        self.image_edit_script = os.path.join(self.skills_base, 'image-edit/scripts/image_edit.py')
    
    def generate_image(self, prompt: str, output_dir: str, width: int = 1024, 
                      height: int = 1024, topic: str = None, channel: str = None) -> str:
        """
        Generate image using image-generation skill
        
        Args:
            prompt: Image generation prompt
            output_dir: Directory to save image
            width: Image width
            height: Image height
            topic: Topic name (for filename)
            channel: Channel name (for subfolder)
            
        Returns:
            Path to generated image
        """
        # Create output directory
        if topic and channel:
            output_path = os.path.join(output_dir, topic, channel, 'images')
        else:
            output_path = output_dir
        
        os.makedirs(output_path, exist_ok=True)
        
        # Build command
        cmd = [
            sys.executable,
            self.image_gen_script,
            'generate',
            prompt,
            '--width', str(width),
            '--height', str(height)
        ]
        
        print(f"\n🎨 Generating image...")
        print(f"  Prompt: {prompt[:100]}...")
        print(f"  Size: {width}x{height}")
        
        try:
            # Run image generation
            result = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(self.image_gen_script))
            
            if result.returncode == 0:
                # Parse output (format: "filename.png [id]")
                output_line = result.stdout.strip().split('\n')[-1]
                image_path = output_line.split(' ')[0]
                
                # Move to our output directory if needed
                if image_path and os.path.exists(image_path):
                    dest_path = os.path.join(output_path, os.path.basename(image_path))
                    if image_path != dest_path:
                        import shutil
                        shutil.copy(image_path, dest_path)
                    print(f"  ✓ Saved: {dest_path}")
                    return dest_path
            
            print(f"  ✗ Generation failed: {result.stderr}")
            return None
            
        except Exception as e:
            print(f"  ✗ Error: {e}")
            return None
    
    def edit_product_image(self, base_image_path: str, edit_prompt: str, 
                          output_dir: str, topic: str = None, channel: str = None) -> str:
        """
        Edit product image using image-edit skill
        
        Args:
            base_image_path: Path to existing product image
            edit_prompt: Edit instructions
            output_dir: Directory to save edited image
            topic: Topic name
            channel: Channel name
            
        Returns:
            Path to edited image
        """
        if not os.path.exists(base_image_path):
            print(f"  ✗ Base image not found: {base_image_path}")
            return None
        
        # Create output directory
        if topic and channel:
            output_path = os.path.join(output_dir, topic, channel, 'images')
        else:
            output_path = output_dir
        
        os.makedirs(output_path, exist_ok=True)
        
        # Build command
        cmd = [
            sys.executable,
            self.image_edit_script,
            edit_prompt,
            base_image_path
        ]
        
        print(f"\n✏️  Editing product image...")
        print(f"  Base: {base_image_path}")
        print(f"  Edit: {edit_prompt[:100]}...")
        
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, cwd=os.path.dirname(self.image_edit_script))
            
            if result.returncode == 0:
                output_line = result.stdout.strip().split('\n')[-1]
                image_path = output_line.split(' ')[0]
                
                if image_path and os.path.exists(image_path):
                    dest_path = os.path.join(output_path, os.path.basename(image_path))
                    if image_path != dest_path:
                        import shutil
                        shutil.copy(image_path, dest_path)
                    print(f"  ✓ Saved: {dest_path}")
                    return dest_path
            
            print(f"  ✗ Edit failed: {result.stderr}")
            return None
            
        except Exception as e:
            print(f"  ✗ Error: {e}")
            return None
    
    def find_product_images(self, product_name: str, website_repo: str) -> List[str]:
        """
        Find existing product images in website repo
        
        Args:
            product_name: Product name to search for
            website_repo: Path to website repository
            
        Returns:
            List of image paths
        """
        import glob
        
        extensions = ['.jpg', '.jpeg', '.png', '.webp']
        found_images = []
        
        # Search patterns
        patterns = [
            f"**/*{product_name}*{{ext}}",
            f"public/images/**/*{{ext}}",
            f"src/assets/**/*{{ext}}"
        ]
        
        for pattern in patterns:
            for ext in extensions:
                search_pattern = pattern.format(ext=ext)
                matches = glob.glob(os.path.join(website_repo, search_pattern), recursive=True)
                found_images.extend(matches[:5])  # Limit per pattern
        
        return list(set(found_images))[:10]  # Return unique, max 10
    
    def handle_product_content(self, product_name: str, website_repo: str, 
                              edit_prompt: str, output_dir: str, 
                              topic: str, channel: str) -> Optional[str]:
        """
        Handle image for product content
        
        Workflow:
        1. Browse website repo for product images
        2. If found: edit with image-edit
        3. If not found: ask user to provide
        
        Args:
            product_name: Product name
            website_repo: Path to website repo
            edit_prompt: Edit instructions
            output_dir: Output directory
            topic: Topic name
            channel: Channel name
            
        Returns:
            Path to image or None
        """
        print(f"\n🔍 Looking for product images: {product_name}")
        
        # Step 1: Find existing images
        images = self.find_product_images(product_name, website_repo)
        
        if images:
            print(f"  ✓ Found {len(images)} image(s)")
            best_image = images[0]  # Use first/best match
            
            # Step 2: Edit image
            return self.edit_product_image(
                best_image,
                edit_prompt,
                output_dir,
                topic,
                channel
            )
        else:
            print(f"  ✗ No product images found in repo")
            print(f"  Please provide product image manually")
            return None
    
    def handle_non_product_content(self, content_type: str, topic: str, 
                                  output_dir: str, channel: str) -> Optional[str]:
        """
        Generate fresh image for non-product content
        
        Args:
            content_type: Type (service, stats, knowledge)
            topic: Topic name
            output_dir: Output directory
            channel: Channel name
            
        Returns:
            Path to generated image
        """
        # Create prompt based on content type
        prompts = {
            'service': f"Professional illustration of {topic}, modern flat design, business context, Thai-friendly aesthetic",
            'stats': f"Data visualization infographic for {topic}, clean charts, professional style",
            'knowledge': f"Educational illustration for {topic}, clear visual metaphor, engaging style",
            'default': f"Professional image for {topic}, modern design, high quality"
        }
        
        prompt = prompts.get(content_type, prompts['default'])
        
        # Generate image
        return self.generate_image(
            prompt,
            output_dir,
            topic=topic,
            channel=channel
        )


def main():
    """Test image integration"""
    parser = argparse.ArgumentParser(description='Test Image Integration')
    parser.add_argument('--action', choices=['generate', 'edit', 'find'], required=True)
    parser.add_argument('--prompt', help='Image prompt or edit instructions')
    parser.add_argument('--topic', help='Topic name')
    parser.add_argument('--channel', help='Channel name')
    parser.add_argument('--output-dir', default='./output', help='Output directory')
    parser.add_argument('--product-name', help='Product name (for find action)')
    parser.add_argument('--website-repo', help='Website repo path (for find action)')
    
    args = parser.parse_args()
    
    integration = ImageIntegration()
    
    if args.action == 'generate':
        result = integration.handle_non_product_content(
            'service', args.topic, args.output_dir, args.channel
        )
        print(f"\nResult: {result}")
    
    elif args.action == 'edit':
        if not args.product_name or not args.website_repo:
            print("Error: --product-name and --website-repo required for edit")
            return
        
        result = integration.handle_product_content(
            args.product_name, args.website_repo, args.prompt, 
            args.output_dir, args.topic, args.channel
        )
        print(f"\nResult: {result}")
    
    elif args.action == 'find':
        if not args.product_name or not args.website_repo:
            print("Error: --product-name and --website-repo required for find")
            return
        
        images = integration.find_product_images(args.product_name, args.website_repo)
        print(f"\nFound {len(images)} images:")
        for img in images:
            print(f"  - {img}")


if __name__ == '__main__':
    main()