"""
Comprehensive free image generation API tester.
Tests Prodia, HuggingFace, and Google Gemini with API keys.

Usage:
  python test_with_keys.py --prodia YOUR_KEY --hf YOUR_TOKEN --gemini YOUR_KEY

Get free keys at:
  Prodia:      https://app.prodia.com/api      (1000 free images/month)
  HuggingFace: https://huggingface.co/settings/tokens  (free, rate-limited)
  Gemini:      https://aistudio.google.com/apikey       (free daily limit)
"""
import requests
import json
import time
import os
import sys
import base64
import argparse

OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))

ANIME_PROMPT = (
    "anime style illustration, young warrior girl with long silver hair and blue eyes, "
    "holding a glowing magical sword, standing in a fantasy forest with ethereal lighting, "
    "detailed anime key visual, masterpiece, best quality, sharp details"
)
NEGATIVE_PROMPT = (
    "low quality, worst quality, blurry, deformed, extra limbs, bad anatomy, "
    "bad hands, missing fingers, watermark, text, signature"
)


def test_prodia(api_key):
    """Test Prodia v1 API - 1000 free images/month"""
    print("\n" + "=" * 60)
    print("PRODIA API TEST")
    print("=" * 60)

    base = "https://api.prodia.com/v1"
    headers = {"X-Prodia-Key": api_key, "Content-Type": "application/json"}

    # 1. List anime models
    print("\n[1] Listing SD 1.5 models...")
    try:
        r = requests.get(f"{base}/sd/models", headers=headers, timeout=15)
        if r.status_code == 200:
            models = r.json()
            anime_models = [m for m in models if any(
                k in m.lower() for k in ["anime", "anything", "counterfeit", "waifu", "novel", "meinamix"]
            )]
            print(f"    Total models: {len(models)}")
            print(f"    Anime models: {len(anime_models)}")
            for m in anime_models:
                print(f"      - {m}")
        else:
            print(f"    Failed: {r.status_code} {r.text[:200]}")
    except Exception as e:
        print(f"    Error: {e}")

    # 2. List SDXL models
    print("\n[2] Listing SDXL models...")
    try:
        r = requests.get(f"{base}/sdxl/models", headers=headers, timeout=15)
        if r.status_code == 200:
            models = r.json()
            anime_models = [m for m in models if any(
                k in m.lower() for k in ["anime", "anything", "counterfeit", "animag", "pony", "novel"]
            )]
            print(f"    Total SDXL models: {len(models)}")
            print(f"    Anime SDXL models: {len(anime_models)}")
            for m in anime_models:
                print(f"      - {m}")
        else:
            print(f"    Failed: {r.status_code} {r.text[:200]}")
    except Exception as e:
        print(f"    Error: {e}")

    # 3. Generate an image with SD 1.5 anime model
    print("\n[3] Generating image (SD 1.5 anime model)...")
    try:
        start = time.time()
        r = requests.post(f"{base}/sd/generate", headers=headers, json={
            "prompt": ANIME_PROMPT,
            "model": "anything-v5-PrtRE.safetensors [7f96a1a9]",
            "negative_prompt": NEGATIVE_PROMPT,
            "steps": 25,
            "cfg_scale": 7,
            "seed": -1,
            "width": 768,
            "height": 512,
            "sampler": "DPM++ 2M Karras",
        }, timeout=30)

        if r.status_code == 200:
            job = r.json()
            job_id = job.get("job", "")
            print(f"    Job created: {job_id}")
            print(f"    Status: {job.get('status')}")

            # Poll for completion
            for attempt in range(30):
                time.sleep(2)
                status_r = requests.get(f"{base}/job/{job_id}", headers=headers, timeout=15)
                if status_r.status_code == 200:
                    status_data = status_r.json()
                    st = status_data.get("status", "unknown")
                    if st == "succeeded":
                        elapsed = time.time() - start
                        img_url = status_data.get("imageUrl", "")
                        print(f"    Completed in {elapsed:.1f}s!")
                        print(f"    Image URL: {img_url}")

                        # Download image
                        if img_url:
                            img_r = requests.get(img_url, timeout=30)
                            if img_r.status_code == 200:
                                path = os.path.join(OUTPUT_DIR, "test_prodia_sd15.png")
                                with open(path, "wb") as f:
                                    f.write(img_r.content)
                                print(f"    SAVED: {path} ({len(img_r.content)} bytes)")
                                return True
                        break
                    elif st == "failed":
                        print(f"    FAILED: {status_data}")
                        break
                    else:
                        print(f"    Waiting... ({st})", end="\r")
            else:
                print("    Timeout after 60 seconds")
        else:
            print(f"    Failed: {r.status_code} {r.text[:300]}")
    except Exception as e:
        print(f"    Error: {e}")

    # 4. Generate with SDXL
    print("\n[4] Generating image (SDXL)...")
    try:
        start = time.time()
        r = requests.post(f"{base}/sdxl/generate", headers=headers, json={
            "prompt": ANIME_PROMPT,
            "negative_prompt": NEGATIVE_PROMPT,
            "steps": 25,
            "cfg_scale": 7,
            "seed": -1,
            "width": 1024,
            "height": 768,
            "sampler": "DPM++ 2M Karras",
        }, timeout=30)

        if r.status_code == 200:
            job = r.json()
            job_id = job.get("job", "")
            print(f"    Job created: {job_id}")

            for attempt in range(30):
                time.sleep(2)
                status_r = requests.get(f"{base}/job/{job_id}", headers=headers, timeout=15)
                if status_r.status_code == 200:
                    status_data = status_r.json()
                    st = status_data.get("status", "unknown")
                    if st == "succeeded":
                        elapsed = time.time() - start
                        img_url = status_data.get("imageUrl", "")
                        print(f"    Completed in {elapsed:.1f}s!")

                        if img_url:
                            img_r = requests.get(img_url, timeout=30)
                            if img_r.status_code == 200:
                                path = os.path.join(OUTPUT_DIR, "test_prodia_sdxl.png")
                                with open(path, "wb") as f:
                                    f.write(img_r.content)
                                print(f"    SAVED: {path} ({len(img_r.content)} bytes)")
                                return True
                        break
                    elif st == "failed":
                        print(f"    FAILED: {status_data}")
                        break
                    else:
                        print(f"    Waiting... ({st})", end="\r")
            else:
                print("    Timeout")
        else:
            print(f"    Failed: {r.status_code} {r.text[:300]}")
    except Exception as e:
        print(f"    Error: {e}")
    return False


def test_huggingface(hf_token):
    """Test HuggingFace Inference API - free with token"""
    print("\n" + "=" * 60)
    print("HUGGINGFACE INFERENCE API TEST")
    print("=" * 60)

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

    models = [
        ("black-forest-labs/FLUX.1-schnell", "FLUX.1 Schnell"),
        ("stabilityai/stable-diffusion-xl-base-1.0", "SDXL Base"),
        ("ByteDance/Hyper-SD", "Hyper-SD"),
    ]

    for model_id, model_name in models:
        print(f"\n  Testing {model_name} ({model_id})...")
        try:
            start = time.time()
            r = requests.post(
                f"https://router.huggingface.co/hf-inference/models/{model_id}",
                headers=headers,
                json={
                    "inputs": ANIME_PROMPT,
                    "parameters": {
                        "negative_prompt": NEGATIVE_PROMPT,
                        "width": 1024,
                        "height": 768,
                    }
                },
                timeout=120,
            )
            elapsed = time.time() - start
            ct = r.headers.get("content-type", "")

            print(f"    Status: {r.status_code} | Type: {ct} | Size: {len(r.content)} | Time: {elapsed:.1f}s")

            if r.status_code == 200 and "image" in ct:
                slug = model_id.split("/")[-1].replace(".", "_")
                path = os.path.join(OUTPUT_DIR, f"test_hf_{slug}.png")
                with open(path, "wb") as f:
                    f.write(r.content)
                print(f"    SUCCESS! Saved: {path}")
            elif r.status_code == 200:
                # Might be JSON with image URL
                try:
                    data = r.json()
                    print(f"    JSON response: {json.dumps(data)[:300]}")
                except:
                    print(f"    Non-image 200 response, size: {len(r.content)}")
            else:
                body = r.text[:400] if len(r.text) < 1000 else r.text[:400]
                print(f"    FAILED: {body}")
        except Exception as e:
            print(f"    Error: {e}")


def test_gemini(api_key):
    """Test Google Gemini image generation - free tier"""
    print("\n" + "=" * 60)
    print("GOOGLE GEMINI IMAGE GENERATION TEST")
    print("=" * 60)

    models_to_try = [
        "gemini-2.0-flash-exp",
        "gemini-2.5-flash-preview-04-17",
    ]

    for model in models_to_try:
        print(f"\n  Testing {model}...")
        try:
            start = time.time()
            r = requests.post(
                f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}",
                headers={"Content-Type": "application/json"},
                json={
                    "contents": [{
                        "parts": [{"text": f"Generate an anime style illustration: {ANIME_PROMPT}"}]
                    }],
                    "generationConfig": {
                        "responseModalities": ["IMAGE", "TEXT"],
                        "temperature": 1.0,
                    }
                },
                timeout=60,
            )
            elapsed = time.time() - start
            print(f"    Status: {r.status_code} | Time: {elapsed:.1f}s")

            if r.status_code == 200:
                data = r.json()
                found_image = False
                for candidate in data.get("candidates", []):
                    for part in candidate.get("content", {}).get("parts", []):
                        if "inlineData" in part:
                            img_data = part["inlineData"]
                            mime = img_data.get("mimeType", "image/png")
                            img_bytes = base64.b64decode(img_data["data"])
                            ext = "png" if "png" in mime else "jpg"
                            path = os.path.join(OUTPUT_DIR, f"test_gemini_{model.replace('-','_')}.{ext}")
                            with open(path, "wb") as f:
                                f.write(img_bytes)
                            print(f"    SUCCESS! Image saved: {path} ({len(img_bytes)} bytes)")
                            found_image = True
                        elif "text" in part:
                            print(f"    Text response: {part['text'][:200]}")
                if not found_image:
                    print(f"    No image in response")
                    print(f"    Full response: {json.dumps(data)[:500]}")
            else:
                print(f"    Failed: {r.text[:400]}")
        except Exception as e:
            print(f"    Error: {e}")


def test_gemini_imagen(api_key):
    """Test Google Imagen via Gemini API - separate image gen model"""
    print("\n" + "=" * 60)
    print("GOOGLE IMAGEN TEST (via Gemini API)")
    print("=" * 60)

    print("\n  Testing imagen-3.0-generate-002...")
    try:
        start = time.time()
        r = requests.post(
            f"https://generativelanguage.googleapis.com/v1beta/models/imagen-3.0-generate-002:predict?key={api_key}",
            headers={"Content-Type": "application/json"},
            json={
                "instances": [{"prompt": ANIME_PROMPT}],
                "parameters": {
                    "sampleCount": 1,
                    "aspectRatio": "16:9",
                }
            },
            timeout=60,
        )
        elapsed = time.time() - start
        print(f"  Status: {r.status_code} | Time: {elapsed:.1f}s")

        if r.status_code == 200:
            data = r.json()
            predictions = data.get("predictions", [])
            for i, pred in enumerate(predictions):
                img_b64 = pred.get("bytesBase64Encoded", "")
                if img_b64:
                    img_bytes = base64.b64decode(img_b64)
                    path = os.path.join(OUTPUT_DIR, f"test_imagen_{i}.png")
                    with open(path, "wb") as f:
                        f.write(img_bytes)
                    print(f"  SUCCESS! Saved: {path} ({len(img_bytes)} bytes)")
        else:
            print(f"  Response: {r.text[:400]}")
    except Exception as e:
        print(f"  Error: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test free image generation APIs")
    parser.add_argument("--prodia", help="Prodia API key (from app.prodia.com/api)")
    parser.add_argument("--hf", help="HuggingFace token (from huggingface.co/settings/tokens)")
    parser.add_argument("--gemini", help="Google Gemini API key (from aistudio.google.com/apikey)")
    args = parser.parse_args()

    print("=" * 60)
    print("FREE IMAGE GENERATION API TESTER")
    print("=" * 60)
    print(f"Testing with: Prodia={'YES' if args.prodia else 'NO'}, "
          f"HF={'YES' if args.hf else 'NO'}, Gemini={'YES' if args.gemini else 'NO'}")

    if not any([args.prodia, args.hf, args.gemini]):
        print("\nNo API keys provided! Get them free at:")
        print("  Prodia:      https://app.prodia.com       → Dashboard → API → Create Key")
        print("  HuggingFace: https://huggingface.co/settings/tokens → New Token")
        print("  Gemini:      https://aistudio.google.com/apikey → Create API Key")
        print("\nThen run:")
        print("  python test_with_keys.py --prodia YOUR_KEY --hf YOUR_TOKEN --gemini YOUR_KEY")
        sys.exit(1)

    if args.prodia:
        test_prodia(args.prodia)

    if args.hf:
        test_huggingface(args.hf)

    if args.gemini:
        test_gemini(args.gemini)
        test_gemini_imagen(args.gemini)

    print("\n" + "=" * 60)
    print("ALL TESTS COMPLETE")
    print("=" * 60)
    print("\nCheck the generated images in:", OUTPUT_DIR)