"""
candy.imagegen — Génération d'images via l'API Candy v20
=========================================================

Nouveautés v20 :
  - ImageGen.edit()    — modifier une zone précise (inpainting)
  - ImageGen.upscale() — améliorer la résolution d'une image

Usage :
    from candy import ImageGen

    # Génère une image
    ImageGen.create("un chat astronaute sur la lune", save="chat.png")

    # Avec paramètres avancés
    ImageGen.create(
        "portrait d'un robot steampunk",
        style="photorealistic",
        size="1024x1024",
        save="robot.png"
    )

    # Variation
    ImageGen.variation("original.png", strength=0.6, save="variation.png")

    # Description
    desc = ImageGen.describe("photo.jpg", lang="FR")

    # Batch
    images = ImageGen.batch(["forêt", "ville", "portrait"], save_dir="./output/")

    # Styles disponibles
    print(ImageGen.styles())

    # Édition / inpainting (v20)
    ImageGen.edit("photo.jpg", mask="mask.png", prompt="remplace le fond par un coucher de soleil", save="edited.png")

    # Upscale (v20)
    ImageGen.upscale("small.png", scale=2, save="big.png")
"""

import base64
import json
import time
from pathlib import Path
from typing import Optional, Union
import httpx

CANDY_JSON_URL = "https://client.hubworld.net/candy.json"

_url_cache: Optional[str] = None
_url_cache_time: float = 0
URL_CACHE_TTL = 300


def _get_api_base() -> str:
    global _url_cache, _url_cache_time
    now = time.time()
    if _url_cache and (now - _url_cache_time) < URL_CACHE_TTL:
        return _url_cache
    try:
        r = httpx.get(CANDY_JSON_URL, timeout=10)
        r.raise_for_status()
        d = r.json()
        if d.get("is_closed", True):
            raise RuntimeError("API Candy hors ligne.")
        link = d.get("link", "").strip().rstrip("/")
        if not link:
            raise RuntimeError("Aucun lien API dans candy.json.")
        _url_cache      = link
        _url_cache_time = now
        return link
    except Exception as e:
        raise RuntimeError(f"Impossible de joindre l'API Candy : {e}")


class ImageGenError(Exception):
    pass


class _ImageGen:
    """
    Client de génération d'images via l'API Candy v20.
    Utilise le même tunnel Cloudflare que le reste de candy.
    """

    def create(
        self,
        prompt: str,
        *,
        style:    str   = "default",
        size:     str   = "512x512",
        steps:    int   = 20,
        guidance: float = 7.5,
        negative: str   = "",
        save:     Optional[str] = None,
        timeout:  int   = 120,
    ) -> Union[str, bytes]:
        """
        Génère une image à partir d'un prompt texte.

        Args:
            prompt:   Description de l'image
            style:    "default" "photorealistic" "anime" "painting"
                      "sketch" "3d" "pixel" "watercolor" "oil" "comic"
            size:     "512x512" "768x768" "1024x1024" "1024x768"
            steps:    Nombre de pas de débruitage (10-50)
            guidance: Échelle de guidage (1.0-20.0)
            negative: Prompt négatif (ce qu'on ne veut PAS)
            save:     Chemin de sauvegarde (.png / .jpg)
            timeout:  Timeout en secondes (défaut 120)
        Returns:
            bytes si save=None, sinon str (chemin du fichier sauvegardé)
        """
        base = _get_api_base()
        payload = {
            "prompt":   prompt,
            "style":    style,
            "size":     size,
            "steps":    steps,
            "guidance": guidance,
            "negative": negative,
        }
        try:
            resp = httpx.post(f"{base}/image/create", json=payload, timeout=timeout)
            resp.raise_for_status()
            data = resp.json()
        except httpx.TimeoutException:
            raise ImageGenError(f"Timeout après {timeout}s.")
        except Exception as e:
            raise ImageGenError(f"Erreur API : {e}")

        b64    = data.get("image_base64", "")
        fmt    = data.get("format", "png")
        img_bytes = base64.b64decode(b64) if b64 else b""

        if save:
            Path(save).write_bytes(img_bytes)
            print(f"\033[90m[candy] 🎨 Image sauvegardée → {save}\033[0m")
            return save
        return img_bytes

    def variation(
        self,
        image_path: str,
        *,
        strength: float = 0.7,
        prompt:   str   = "",
        save:     Optional[str] = None,
        timeout:  int   = 120,
    ) -> Union[str, bytes]:
        """
        Génère une variation d'une image existante.

        Args:
            image_path: Chemin vers l'image source
            strength:   Force de la variation (0.0 = identique, 1.0 = totalement différent)
            prompt:     Prompt optionnel pour guider la variation
            save:       Chemin de sauvegarde
        """
        base     = _get_api_base()
        img_data = Path(image_path).read_bytes()
        b64_in   = base64.b64encode(img_data).decode()
        payload  = {
            "image_base64": b64_in,
            "strength":     strength,
            "prompt":       prompt,
        }
        try:
            resp = httpx.post(f"{base}/image/variation", json=payload, timeout=timeout)
            resp.raise_for_status()
            data = resp.json()
        except Exception as e:
            raise ImageGenError(f"Erreur variation : {e}")

        img_bytes = base64.b64decode(data.get("image_base64", ""))
        if save:
            Path(save).write_bytes(img_bytes)
            print(f"\033[90m[candy] 🎨 Variation sauvegardée → {save}\033[0m")
            return save
        return img_bytes

    def describe(
        self,
        image_path: str,
        *,
        lang:    str = "FR",
        detail:  str = "medium",
        timeout: int = 60,
    ) -> str:
        """
        Décrit une image en langage naturel.

        Args:
            image_path: Chemin vers l'image
            lang:       Langue de la description ("FR", "EN", etc.)
            detail:     Niveau de détail — "low" "medium" "high"
        """
        base     = _get_api_base()
        img_data = Path(image_path).read_bytes()
        b64_in   = base64.b64encode(img_data).decode()
        suffix   = Path(image_path).suffix.lstrip(".").lower()
        mime     = "image/jpeg" if suffix in ("jpg", "jpeg") else f"image/{suffix}"
        payload  = {
            "image_base64": b64_in,
            "mime_type":    mime,
            "lang":         lang,
            "detail":       detail,
        }
        try:
            resp = httpx.post(f"{base}/image/describe", json=payload, timeout=timeout)
            resp.raise_for_status()
            return resp.json().get("description", "")
        except Exception as e:
            raise ImageGenError(f"Erreur description : {e}")

    def batch(
        self,
        prompts:     list,
        *,
        style:       str            = "default",
        size:        str            = "512x512",
        save_dir:    Optional[str]  = None,
        timeout:     int            = 120,
    ) -> list:
        """
        Génère plusieurs images d'un coup.

        Args:
            prompts:  Liste de prompts
            style:    Style visuel commun
            size:     Taille commune
            save_dir: Dossier de sauvegarde (sinon retourne liste de bytes)
        Returns:
            Liste de chemins si save_dir, sinon liste de bytes
        """
        results = []
        if save_dir:
            Path(save_dir).mkdir(parents=True, exist_ok=True)
        for i, prompt in enumerate(prompts):
            try:
                save_path = str(Path(save_dir) / f"image_{i+1:03d}.png") if save_dir else None
                result    = self.create(prompt, style=style, size=size,
                                        save=save_path, timeout=timeout)
                results.append(result)
            except Exception as e:
                results.append(f"[ERREUR] {e}")
        return results

    # ── Nouveautés v20 ────────────────────────────────────────────────────────

    def edit(
        self,
        image_path: str,
        *,
        mask:    Optional[str] = None,
        prompt:  str           = "",
        save:    Optional[str] = None,
        timeout: int           = 120,
    ) -> Union[str, bytes]:
        """
        Édite une image existante via inpainting (v20).

        Args:
            image_path: Chemin vers l'image à éditer
            mask:       Chemin vers le masque (zones blanches = à modifier)
                        Si None, le modèle détecte automatiquement les zones
            prompt:     Description de ce qu'on veut mettre à la place
            save:       Chemin de sauvegarde

        Example:
            ImageGen.edit(
                "photo.jpg",
                mask="mask.png",
                prompt="remplace le fond par un coucher de soleil",
                save="edited.png"
            )
        """
        base     = _get_api_base()
        img_data = Path(image_path).read_bytes()
        b64_in   = base64.b64encode(img_data).decode()
        payload: dict = {
            "image_base64": b64_in,
            "prompt":       prompt,
        }
        if mask:
            mask_data      = Path(mask).read_bytes()
            payload["mask_base64"] = base64.b64encode(mask_data).decode()
        try:
            resp = httpx.post(f"{base}/image/edit", json=payload, timeout=timeout)
            resp.raise_for_status()
            data = resp.json()
        except Exception as e:
            raise ImageGenError(f"Erreur edit : {e}")
        img_bytes = base64.b64decode(data.get("image_base64", ""))
        if save:
            Path(save).write_bytes(img_bytes)
            print(f"\033[90m[candy] ✏️ Image éditée sauvegardée → {save}\033[0m")
            return save
        return img_bytes

    def upscale(
        self,
        image_path: str,
        *,
        scale:   int           = 2,
        save:    Optional[str] = None,
        timeout: int           = 120,
    ) -> Union[str, bytes]:
        """
        Améliore la résolution d'une image (v20).

        Args:
            image_path: Chemin vers l'image basse résolution
            scale:      Facteur d'agrandissement — 2 ou 4
            save:       Chemin de sauvegarde

        Example:
            ImageGen.upscale("small.png", scale=4, save="big.png")
        """
        if scale not in (2, 4):
            raise ValueError("scale doit être 2 ou 4.")
        base     = _get_api_base()
        img_data = Path(image_path).read_bytes()
        b64_in   = base64.b64encode(img_data).decode()
        payload  = {"image_base64": b64_in, "scale": scale}
        try:
            resp = httpx.post(f"{base}/image/upscale", json=payload, timeout=timeout)
            resp.raise_for_status()
            data = resp.json()
        except Exception as e:
            raise ImageGenError(f"Erreur upscale : {e}")
        img_bytes = base64.b64decode(data.get("image_base64", ""))
        if save:
            Path(save).write_bytes(img_bytes)
            print(f"\033[90m[candy] 🔍 Image upscalée (x{scale}) → {save}\033[0m")
            return save
        return img_bytes

    def styles(self) -> list:
        """Retourne la liste des styles disponibles."""
        return [
            "default", "photorealistic", "anime", "painting",
            "sketch", "3d", "pixel", "watercolor", "oil", "comic",
        ]


ImageGen = _ImageGen()
