src/services/image.py

"""
Image Service

High-level service for image generation and editing.
Abstracts all API complexity from the UI layer.
"""

from typing import Callable, Optional, List
from dataclasses import dataclass

from ..api.client import StackNetClient


@dataclass
class GeneratedImage:
    """Generated image result."""
    image_url: str
    image_path: Optional[str] = None
    prompt: Optional[str] = None
    width: Optional[int] = None
    height: Optional[int] = None


class ImageService:
    """
    Service for image generation and editing.

    Provides clean interfaces for:
    - Text-to-image generation
    - Image-to-image editing/transformation
    """

    def __init__(self, client: Optional[StackNetClient] = None):
        self.client = client or StackNetClient()

    async def generate_image(
        self,
        prompt: str,
        format_type: str = "image",
        on_progress: Optional[Callable[[float, str], None]] = None
    ) -> List[GeneratedImage]:
        """
        Generate image from a text prompt.

        Args:
            prompt: Description of desired image
            format_type: Format type - "image" (generate_image_5),
                        "multi" (generate_image_multi_5), or "3d" (generate_image_3d)
            on_progress: Callback for progress updates

        Returns:
            List of generated images
        """
        # Select tool based on format type
        if format_type == "multi":
            tool_name = "generate_image_multi_5"
        elif format_type == "3d":
            tool_name = "generate_image_3d"
        else:
            tool_name = "generate_image_5"

        result = await self.client.submit_tool_task(
            tool_name=tool_name,
            parameters={
                "prompt": prompt
            },
            on_progress=on_progress
        )

        if not result.success:
            raise Exception(result.error or "Image generation failed")

        return self._parse_image_result(result.data, prompt)

    async def edit_image(
        self,
        image_url: str,
        edit_prompt: str,
        strength: float = 0.5,
        on_progress: Optional[Callable[[float, str], None]] = None
    ) -> List[GeneratedImage]:
        """
        Edit/transform an existing image using generate_image_edit_5 tool.

        Args:
            image_url: URL to source image
            edit_prompt: Edit instructions
            strength: Edit strength (0.1 to 1.0)
            on_progress: Progress callback

        Returns:
            List of edited images
        """
        result = await self.client.submit_tool_task(
            tool_name="generate_image_edit_5",
            parameters={
                "prompt": edit_prompt,
                "image_url": image_url,
                "strength": strength
            },
            on_progress=on_progress
        )

        if not result.success:
            raise Exception(result.error or "Image editing failed")

        return self._parse_image_result(result.data, edit_prompt)

    def _parse_image_result(self, data: dict, prompt: str) -> List[GeneratedImage]:
        """Parse API response into GeneratedImage objects."""
        images = []

        # Handle various response formats
        raw_images = data.get("images", [])

        if not raw_images:
            # Check for single image URL
            image_url = (
                data.get("image_url") or
                data.get("imageUrl") or
                data.get("url") or
                data.get("content")
            )
            if image_url:
                raw_images = [{"url": image_url}]

        for img_data in raw_images:
            if isinstance(img_data, str):
                # Raw URL string
                image_url = img_data
            else:
                image_url = (
                    img_data.get("url") or
                    img_data.get("image_url") or
                    img_data.get("imageUrl")
                )

            if image_url:
                images.append(GeneratedImage(
                    image_url=image_url,
                    prompt=prompt,
                    width=img_data.get("width") if isinstance(img_data, dict) else None,
                    height=img_data.get("height") if isinstance(img_data, dict) else None
                ))

        return images

    async def download_image(self, image: GeneratedImage) -> str:
        """Download an image to local file."""
        if image.image_path:
            return image.image_path

        # Determine extension from URL
        url = image.image_url
        if ".png" in url:
            ext = ".png"
        elif ".jpg" in url or ".jpeg" in url:
            ext = ".jpg"
        else:
            ext = ".png"

        filename = f"image_{hash(url) % 10000}{ext}"
        image.image_path = await self.client.download_file(url, filename)
        return image.image_path

    def cleanup(self):
        """Clean up temporary files."""
        self.client.cleanup()