"""Image generation — text-to-image, img-to-img, upscaling, and more.

In this example, we:
  - Parse a model ID from a SeaArt URL
  - Estimate the coin cost before generating
  - Submit text-to-image tasks and wait for results
  - Generate multiple images in a single task
  - Transform an existing image with img-to-img
  - Upscale a generated image
  - Remove background from an image
  - Upload a local file and use it as input
  - Fetch artwork details and download URLs

All generation methods return a TaskEnvelope with .wait() for polling.
Call envelope.wait(timeout=...) to block until the task completes.

Requires:
    SEAART_EMAIL and SEAART_PASSWORD environment variables.
"""

from __future__ import annotations

import os

from seaart import SyncClient
from seaart.helpers import extract_model_link

MODEL_URL = "https://www.seaart.ai/create/image?id=cu8d8ble878c738hejd0&model_ver_no=f1ebdcc036b767a672db420eb367688d"


# ---------------------------------------------------------------------------
# Cost estimation
# ---------------------------------------------------------------------------


def estimate_cost() -> None:
    """Check how many coins a generation will cost before running it."""
    model = extract_model_link(MODEL_URL)
    assert model is not None

    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        est = client.images.estimate(
            prompt="a serene mountain landscape at sunset, oil painting",
            model_no=model.model_no,
            # model_ver_no is optional — resolved automatically from cache if omitted.
            model_ver_no=model.model_ver_no,
        )
        print(f"Estimated cost: {est.value.calculate_coins} coins")


# ---------------------------------------------------------------------------
# Text-to-image
# ---------------------------------------------------------------------------


def text_to_image() -> None:
    """Generate a single image from a text prompt."""
    model = extract_model_link(MODEL_URL)
    assert model is not None

    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        envelope = client.images.text_to_img(
            prompt="a serene mountain landscape at sunset, oil painting",
            model_no=model.model_no,
            model_ver_no=model.model_ver_no,
        )
        print(f"Task submitted: {envelope.value.id}")

        # .wait() polls until the task finishes and returns a TaskItem.
        item = envelope.wait(timeout=120)
        print(f"Image URL: {item.image_url}")


def text_to_image_batch() -> None:
    """Generate multiple images in a single task."""
    model = extract_model_link(MODEL_URL)
    assert model is not None

    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        envelope = client.images.text_to_img(
            prompt="cyberpunk cityscape, neon lights, rain",
            model_no=model.model_no,
            num_images=2,
        )
        item = envelope.wait(timeout=180)

        # .image_urls contains all generated images.
        for url in item.image_urls:
            print(f"Image: {url}")


# ---------------------------------------------------------------------------
# Image-to-image
# ---------------------------------------------------------------------------


def image_to_image() -> None:
    """Transform an existing image with a text prompt.

    strength controls how much the output differs from the input:
    low values (0.1–0.3) preserve more of the original,
    high values (0.7–1.0) allow greater creative freedom.
    """
    model = extract_model_link(MODEL_URL)
    assert model is not None

    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        # Upload a local image first (or use an existing URL).
        uploaded = client.images.upload(path="input.jpg")
        image_url = uploaded.value.url

        envelope = client.images.img_to_img(
            prompt="same scene but in watercolor style",
            model_no=model.model_no,
            image_url=image_url,
            strength=0.45,
        )
        item = envelope.wait(timeout=120)
        print(f"Result: {item.image_url}")


# ---------------------------------------------------------------------------
# Upscaling
# ---------------------------------------------------------------------------


def upscale_image() -> None:
    """Upscale a generated image to higher resolution."""
    model = extract_model_link(MODEL_URL)
    assert model is not None

    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        # Generate an image first.
        envelope = client.images.text_to_img(
            prompt="a detailed portrait of a knight",
            model_no=model.model_no,
        )
        item = envelope.wait(timeout=120)

        # Upscale it.
        assert item.image_url is not None
        upscale_envelope = client.images.upscale(image_url=item.image_url)
        upscaled = upscale_envelope.wait(timeout=120)
        print(f"Upscaled: {upscaled.image_url}")


# ---------------------------------------------------------------------------
# Background removal
# ---------------------------------------------------------------------------


def remove_background() -> None:
    """Remove the background from an image."""
    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        uploaded = client.images.upload(path="photo.jpg")
        envelope = client.images.remove_bg(image_url=uploaded.value.url)
        item = envelope.wait(timeout=60)
        print(f"Result: {item.image_url}")


# ---------------------------------------------------------------------------
# Artwork details & download
# ---------------------------------------------------------------------------


def artwork_details() -> None:
    """Fetch full artwork info and get a direct download URL."""
    with SyncClient() as client:
        client.auth.login(
            email=os.environ["SEAART_EMAIL"],
            password=os.environ["SEAART_PASSWORD"],
        )

        model = extract_model_link(MODEL_URL)
        assert model is not None

        envelope = client.images.text_to_img(
            prompt="a fox in autumn forest",
            model_no=model.model_no,
        )
        item = envelope.wait(timeout=120)

        # Fetch detailed artwork metadata.
        details = client.images.detail(item.task_id)
        print(f"Model used: {details.value.art_model_name}")
        for uri in details.value.img_uris or []:
            print(f"  {uri.url}")

        # Get a direct download link.
        assert item.image_url is not None
        download = client.images.get_download_url(item.image_url)
        print(f"Download: {download.value.url}")


if __name__ == "__main__":
    text_to_image()
