# ============================
# FIREWORKS AI MODULE (HYBRID MODE)
# ============================
import os
import base64
import requests
import uuid

FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")

def fireworks_generate(prompt, model_name, log_callback):
    """
    Generate image using Fireworks AI.
    Returns filepath or None.
    """
    url = "https://api.fireworks.ai/inference/v1/image/generate"

    payload = {
        "model": model_name,
        "prompt": prompt,
        "size": "1024x1024",
        "num_images": 1
    }

    headers = {
        "Authorization": f"Bearer {FIREWORKS_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(url, json=payload, headers=headers, timeout=40)
        data = response.json()

        if "images" not in data:
            log_callback(f"  → Fireworks Error: {data}")
            return None

        b64 = data["images"][0]["image_base64"]
        img_bytes = base64.b64decode(b64)

        filename = f"fw_{uuid.uuid4().hex}.png"
        os.makedirs("outputs", exist_ok=True)
        filepath = os.path.join("outputs", filename)

        with open(filepath, "wb") as f:
            f.write(img_bytes)

        if os.path.getsize(filepath) < 10000:
            log_callback("  → Fireworks FAILED (image too small)")
            return None

        log_callback("  → Fireworks SUCCESS")
        return filepath

    except Exception as e:
        log_callback(f"  → Fireworks Exception: {str(e)}")
        return None


def fireworks_schnell(prompt, log_callback):
    return fireworks_generate(prompt, "black-forest-labs/FLUX.1-schnell", log_callback)


def fireworks_dev(prompt, log_callback):
    return fireworks_generate(prompt, "black-forest-labs/FLUX.1-dev", log_callback)