agents/specialized/image_agent.py

# agents/specialized/image_agent.py
#
# Cascade de generación de imágenes:
#   1. HuggingFace FLUX.1-schnell  (genera imagen)
#      → sube a HF Dataset vfven/mission-control-images (persistente)
#   2. Gemini Imagen API            (pendiente, requiere billing)
#   3. Picsum Photos                (fallback visual garantizado)

import os
import re
import json
import asyncio
from pathlib import Path
from datetime import datetime
from io import BytesIO

from agents.specialized.base import build_result, call_llm
from agents.registry import AGENT_REGISTRY

ROLE = """
Eres el agente de imágenes.
Cuando se te pida generar, crear, mostrar o visualizar una imagen, responde SOLO con un JSON:
{"image_queries": ["término en inglés 1", "término en inglés 2", "término en inglés 3"]}
Los términos deben ser específicos, detallados y en inglés para obtener mejores resultados.
Si la tarea no es sobre imágenes, responde: {"skip":"no image task"}
"""

DOCS_DIR       = Path("data/docs")
HF_MODEL       = "black-forest-labs/FLUX.1-schnell"
HF_DATASET     = "vfven/mission-control-images"
IMAGE_PROVIDERS = ["fal-ai", "hf-inference", "replicate", "together"]


# ─────────────────────────────────────────────────────────────────────────────
# Subir imagen a HF Dataset (persistente)
# ─────────────────────────────────────────────────────────────────────────────

async def _upload_to_dataset(image, file_name: str) -> str | None:
    """
    Sube una imagen PIL al dataset vfven/mission-control-images.
    Devuelve la URL pública permanente o None si falla.
    El dataset se crea automáticamente si no existe.
    """
    try:
        from huggingface_hub import HfApi
        import inspect

        token = os.getenv("HF_API_TOKEN")
        api   = HfApi(token=token)

        # Crear dataset si no existe
        try:
            api.create_repo(
                repo_id=HF_DATASET,
                repo_type="dataset",
                private=False,
                exist_ok=True,
            )
            print(f"[Dataset] repo {HF_DATASET} listo")
        except Exception as e:
            print(f"[Dataset] advertencia al crear repo: {e}")

        # Serializar imagen PIL → bytes en memoria
        buf = BytesIO()
        await asyncio.to_thread(image.save, buf, format="PNG")
        buf.seek(0)

        # Subir al dataset
        path_in_repo = f"images/{file_name}"

        # upload_file acepta file-like object
        kwargs = dict(
            path_or_fileobj=buf,
            path_in_repo=path_in_repo,
            repo_id=HF_DATASET,
            repo_type="dataset",
            token=token,
        )

        # commit_message solo en versiones que lo soporten
        if "commit_message" in inspect.signature(api.upload_file).parameters:
            kwargs["commit_message"] = f"Add {file_name}"

        await asyncio.to_thread(api.upload_file, **kwargs)

        # URL pública del dataset
        url = f"https://huggingface.co/datasets/{HF_DATASET}/resolve/main/{path_in_repo}"
        print(f"[Dataset] subida OK → {url}")
        return url

    except Exception as e:
        print(f"[Dataset] error subiendo imagen: {e}")
        return None


# ─────────────────────────────────────────────────────────────────────────────
# 1. HuggingFace FLUX
# ─────────────────────────────────────────────────────────────────────────────

async def _generate_with_hf(queries: list) -> list:
    import inspect
    from huggingface_hub import InferenceClient

    token = os.getenv("HF_API_TOKEN")
    urls  = []
    supports_provider = "provider" in inspect.signature(InferenceClient.__init__).parameters

    for prompt in queries[:2]:
        generated = False

        if supports_provider:
            for prov in IMAGE_PROVIDERS:
                try:
                    print(f"[HF] provider={prov} → '{prompt[:45]}...'")
                    client = InferenceClient(provider=prov, api_key=token)
                    image  = await asyncio.to_thread(
                        client.text_to_image, prompt, model=HF_MODEL
                    )

                    file_name = f"hf_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.png"

                    # Intentar subir al dataset (persistente)
                    dataset_url = await _upload_to_dataset(image, file_name)

                    if dataset_url:
                        urls.append(dataset_url)
                        print(f"[HF] imagen en dataset → {dataset_url}")
                    else:
                        # Fallback: guardar local (efímero)
                        DOCS_DIR.mkdir(parents=True, exist_ok=True)
                        file_path = DOCS_DIR / file_name
                        await asyncio.to_thread(image.save, str(file_path))
                        urls.append(f"/docs/{file_name}")
                        print(f"[HF] imagen local (efímera) → {file_name}")

                    generated = True
                    break
                except Exception as e:
                    print(f"[HF] {prov} falló: {e}")
        else:
            try:
                print(f"[HF] versión antigua sin provider → '{prompt[:45]}...'")
                client = InferenceClient(token=token)
                image  = await asyncio.to_thread(client.text_to_image, prompt, model=HF_MODEL)
                file_name   = f"hf_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}.png"
                dataset_url = await _upload_to_dataset(image, file_name)
                if dataset_url:
                    urls.append(dataset_url)
                else:
                    DOCS_DIR.mkdir(parents=True, exist_ok=True)
                    await asyncio.to_thread(image.save, str(DOCS_DIR / file_name))
                    urls.append(f"/docs/{file_name}")
                generated = True
            except Exception as e:
                print(f"[HF] excepción: {e}")

        if not generated:
            print(f"[HF] todos los providers fallaron para este prompt")

    return urls


# ─────────────────────────────────────────────────────────────────────────────
# 2. Gemini Imagen API  (pendiente)
# ─────────────────────────────────────────────────────────────────────────────

async def _generate_with_gemini(queries: list) -> list:
    print("[Gemini] imagen API no habilitada aún, saltando...")
    return []


# ─────────────────────────────────────────────────────────────────────────────
# Cascade: Gemini → HF → picsum
# ─────────────────────────────────────────────────────────────────────────────

async def _generate_images(queries: list, provider: str) -> list:
    image_urls = []

    if provider == "gemini":
        try:
            image_urls = await _generate_with_gemini(queries)
        except Exception as e:
            print(f"[cascade] Gemini falló: {e}")

    if not image_urls:
        try:
            image_urls = await _generate_with_hf(queries)
        except Exception as e:
            print(f"[cascade] HF falló: {e}")

    if not image_urls:
        print("[cascade] usando picsum como fallback final")
        image_urls = [
            f"https://picsum.photos/seed/{abs(hash(q)) % 9999}/400/300"
            for q in queries[:2]
        ]

    return image_urls


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────

async def run(task: str, context: dict = None) -> dict:
    result = build_result("image_agent")
    try:
        response = await call_llm("image_agent", ROLE, task, context)
        result["response"] = response

        agent_config = AGENT_REGISTRY.get("image_agent", {})
        provider     = agent_config.get("provider", "huggingface")
        print(f"[image_agent] provider={provider}")

        match = re.search(r'\{.*"image_queries".*\}', response, re.DOTALL)

        if match:
            try:
                data    = json.loads(match.group(0))
                queries = data.get("image_queries", [])

                if not queries:
                    result["response"] = "El agente no generó queries de imagen."
                    return result

                result["queries"]  = queries
                result["response"] = "Generando imágenes para:\n" + "\n".join(f"• {q}" for q in queries)

                image_urls = await _generate_images(queries, provider)

                result["image_urls"] = image_urls
                result["response"]  += "\n\nImagen generada ✔" if image_urls else "\n\n⚠️ No se pudo generar imagen"

            except json.JSONDecodeError as e:
                result["response"] = f"Error parseando JSON del agente: {e}"
        else:
            result["response"] = response

    except Exception as e:
        result["success"] = False
        result["error"]   = str(e)
        print(f"[image_agent] excepción general: {e}")

    return result