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, MediaAction


@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,
        style: Optional[str] = None,
        aspect_ratio: Optional[str] = None,
        on_progress: Optional[Callable[[float, str], None]] = None
    ) -> List[GeneratedImage]:
        """
        Generate image from a text prompt.

        Args:
            prompt: Description of desired image
            style: Style preset (Photorealistic, Digital Art, etc.)
            aspect_ratio: Aspect ratio (1:1, 16:9, 9:16, etc.)
            on_progress: Callback for progress updates

        Returns:
            List of generated images
        """
        full_prompt = prompt
        if style and style != "Photorealistic":
            full_prompt = f"{prompt}, {style.lower()} style"

        options = {}
        if aspect_ratio:
            options["aspect_ratio"] = aspect_ratio

        result = await self.client.submit_media_task(
            action=MediaAction.ANALYZE_VISUAL,
            prompt=full_prompt,
            options=options if options else None,
            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.

        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
        """
        options = {
            "strength": strength,
            "edit_mode": True
        }

        result = await self.client.submit_media_task(
            action=MediaAction.ANALYZE_VISUAL,
            media_url=image_url,
            prompt=edit_prompt,
            options=options,
            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()