"""
Image Generation Tool
Generate images using ComfyUI.
"""
from typing import Dict, Any, Optional
from loguru import logger

from tools.base import BaseTool, ToolResult
from tools.comfyui.base import ComfyUIClient


class ImageGenerationTool(BaseTool):
    """Generate images using ComfyUI."""
    
    def __init__(self, config: Optional[Dict] = None):
        self.client = ComfyUIClient()
        super().__init__(config)
    
    @property
    def name(self) -> str:
        return "generate_image"
    
    @property
    def description(self) -> str:
        return "Generate an image from a text description. Creates visual content based on your prompt."
    
    @property
    def parameters(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "Description of the image to generate"
                },
                "negative_prompt": {
                    "type": "string",
                    "description": "What to avoid in the image (optional)",
                    "default": ""
                },
                "size": {
                    "type": "string",
                    "description": "Image size (e.g., '512x512', '1024x768')",
                    "default": "512x512"
                },
                "steps": {
                    "type": "integer",
                    "description": "Number of generation steps",
                    "default": 20
                },
                "cfg_scale": {
                    "type": "number",
                    "description": "CFG scale for prompt adherence",
                    "default": 7.0
                },
                "seed": {
                    "type": "integer",
                    "description": "Random seed for reproducibility (optional)"
                }
            },
            "required": ["prompt"]
        }
    
    async def execute(
        self,
        prompt: str,
        negative_prompt: str = "",
        size: str = "512x512",
        steps: int = 20,
        cfg_scale: float = 7.0,
        seed: Optional[int] = None,
        **kwargs
    ) -> ToolResult:
        """Generate an image."""
        self._log_execution({"prompt": prompt[:100], "size": size, "steps": steps})
        
        # Reload config to get latest settings
        self.client.reload_config()
        
        # Load the image workflow
        workflow = self.client.load_workflow("image")
        
        if not workflow:
            return ToolResult(
                success=False,
                error="Image generation workflow not configured. Please upload a workflow JSON in the admin panel."
            )
        
        try:
            # Modify workflow with parameters
            modified_workflow = self.client.modify_workflow(
                workflow,
                prompt=prompt,
                workflow_type="image",
                negative_prompt=negative_prompt,
                size=size,
                steps=steps,
                cfg_scale=cfg_scale,
                seed=seed
            )
            
            # Queue the prompt
            prompt_id = await self.client.queue_prompt(modified_workflow)
            logger.info(f"Queued image generation: {prompt_id}")
            
            # Wait for completion
            outputs = await self.client.wait_for_completion(
                prompt_id,
                timeout=300  # 5 minutes for image generation
            )
            
            # Get output images
            images = await self.client.get_output_images(outputs)
            
            if not images:
                return ToolResult(
                    success=False,
                    error="No images were generated"
                )
            
            # Return info about generated images
            result_parts = [f"Successfully generated {len(images)} image(s):"]
            for img in images:
                result_parts.append(f"  - {img['filename']}")
            
            result = "\n".join(result_parts)
            
            self._log_success(result)
            return ToolResult(success=True, data=result)
        
        except TimeoutError as e:
            self._log_error(str(e))
            return ToolResult(success=False, error="Image generation timed out")
        
        except Exception as e:
            self._log_error(str(e))
            return ToolResult(success=False, error=str(e))