"""ImageGenerator port — 图片生成能力契约。"""

from dataclasses import dataclass
from enum import Enum
from typing import Protocol, runtime_checkable


class TaskStatus(Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class ImageResult:
    status: TaskStatus
    task_id: str
    image_url: str | None = None
    image_bytes: bytes | None = None
    error: str | None = None


@runtime_checkable
class ImageGenerator(Protocol):
    def generate(self, prompt: str, size: str, style: str) -> ImageResult:
        """Submit generation, poll until done, return result (sync, may block)."""
        ...

    def check_status(self, task_id: str) -> ImageResult:
        """Check current status of a generation task."""
        ...

    def download_image(self, image_url: str) -> bytes:
        """Download image bytes from a generated image URL (e.g. from ImageResult.image_url)."""
        ...