import os
import io
import requests
from fastapi import FastAPI, HTTPException
from huggingface_hub import InferenceClient
from fastapi.responses import Response
from pydantic import BaseModel, HttpUrl

app = FastAPI()

client = InferenceClient(
    provider="fal-ai",
    api_key=os.environ["HF_TOKEN"]
)

class GenerateRequest(BaseModel):
    image_url: HttpUrl
    prompt: str
    model: str = "tencent/HunyuanImage-3.0-Instruct"

@app.get("/")
def greet_json():
    return {"Hello": "World"}

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/convert")
def convert(req: GenerateRequest):
    print(f"[convert] Recibida petición: prompt={req.prompt[:50]}...")

    # Descargar la imagen primero
    resp = requests.get(str(req.image_url))
    if resp.status_code != 200:
        raise HTTPException(status_code=400, detail="No se pudo descargar la imagen")
    input_image = resp.content
    print(f"[convert] Imagen descargada: {len(input_image)} bytes")

    try:
        out_img = client.image_to_image(
            input_image,
            prompt=req.prompt,
            model=req.model
        )
        print(f"[convert] Imagen generada: type={type(out_img)}, size={getattr(out_img, 'size', 'N/A')}")
    except Exception as e:
        print(f"[convert] Error en inferencia: {e}")
        raise HTTPException(status_code=502, detail=str(e))

    buf = io.BytesIO()
    out_img.save(buf, format="PNG")
    print(f"[convert] PNG guardado en buffer: {len(buf.getvalue())} bytes")
    return Response(content=buf.getvalue(), media_type="image/png")