import base64
import io
import os
import re
import time
from PIL import Image
from g4f.client import Client
from g4f.Provider import RetryProvider, PollinationsAI, ImageLabs, Blackbox, HuggingSpace, Airforce
from g4f.Provider.hf_space.BlackForestLabsFlux1Schnell import BlackForestLabsFlux1Schnell
from g4f.Provider.hf_space.VoodoohopFlux1Schnell import VoodoohopFlux1Schnell

NEGATIVE_PROMPT = (
    "low quality, blurry, pixelated, bad anatomy, bad hands, three hands, three legs, bad arms, missing legs, "
    "missing arms, poorly drawn face, poorly rendered hands, bad face, fused face, cloned face, worst face, "
    "three crus, extra crus, fused crus, worst feet, three feet, fused feet, fused thigh, three thigh, extra thigh, "
    "worst thigh, missing fingers, extra fingers, ugly fingers, long fingers, bad composition, horn, extra eyes, huge eyes, "
    "2girl, amputation, disconnected limbs, cartoon, cg, 3d, unreal, animate, cgi, render, artwork, illustration, "
    "3d render, cinema 4d, artstation, octane render, mutated body parts, painting, oil painting, 2d, sketch, bad photography, "
    "bad photo, deviant art, aberrations, abstract, anime, black and white, collapsed, conjoined, creative, drawing, extra windows, "
    "harsh lighting, jpeg artifacts, low saturation, monochrome, multiple levels, overexposed, oversaturated, photoshop, rotten, surreal, "
    "twisted, UI, underexposed, unnatural, unreal engine, unrealistic, video game, deformed body features, NSFW, NUDE, vulgar, negative, "
    "unsuitable, inappropriate, offensive, revealing, sexual, explicit"
)

def extract_summary(text):
    """
    Clean and extract the summary portion from the text.
    """
    text = text.replace("#", "").strip().lower()
    match = re.search(r"summary(.*?)highlights", text, re.DOTALL)
    return match.group(1).strip() if match else text

def fix_base64_padding(data):
    """
    Ensure that the base64 string has the proper padding.
    """
    missing_padding = len(data) % 4
    if missing_padding:
        data += "=" * (4 - missing_padding)
    return data

def generate_image(title, category, summary):
    print("Generating image...")
    start = time.time()
    prompt = f"Generate accurate image representing the {category} concept: ```{title.strip()}: {summary.strip()}```"
    client = Client()
    attempts = [
        ([BlackForestLabsFlux1Schnell, VoodoohopFlux1Schnell, HuggingSpace], "flux-schnell"),
        ([HuggingSpace, PollinationsAI], "flux-dev"),
        ([PollinationsAI], "flux-pro"),
        ([Airforce, PollinationsAI, Blackbox], "flux"),
        ([PollinationsAI], "dall-e-3"),
        ([HuggingSpace], "sd-3")
    ]
    
    for providers, model in attempts:
        try:
            provider = RetryProvider(
                providers=providers,
                shuffle=True,
                single_provider_retry=True,
                max_retries=3,
            )
            response = client.images.generate(
                provider=provider,
                model=model,
                prompt=prompt,
                negative_prompt=NEGATIVE_PROMPT,
                response_format="b64_json",
                width=1024,
                height=576,
            )
            img_data = response.data[0].b64_json
            elapsed = time.time() - start
            print(f"Image generated in {elapsed:.2f} seconds using model {model}")
            if img_data:
                return f"data:image/png;base64,{img_data}"
        except Exception as e:
            print(f"Attempt with model {model} failed: {e}")
    
    return None

def verify_image(image_data):
    try:
        image_stream = io.BytesIO(image_data)
        img = Image.open(image_stream)
        img.verify()
        return True
    except Exception as e:
        print(f"Error verifying image: {e}")
        return False

def fetch_image(title, category, summary):
    summary = extract_summary(summary)
    fallback_url = "https://i.ibb.co/TBJqggw/Image-Not-Found.jpg"
    try:
        data_uri = generate_image(title, category, summary)
        if data_uri:
            base64_str = data_uri.split(",")[1]
            base64_str = fix_base64_padding(base64_str)
            decoded = base64.b64decode(base64_str, validate=True)
            if verify_image(decoded):
                return f"data:image/png;base64,{base64_str}"
        return fallback_url
    except Exception as e:
        print(f"Error fetching image: {e}")
        return fallback_url
    finally:
        if os.path.exists("image.png"):
            os.remove("image.png")

if __name__ == "__main__":
    title = "Exposition: Enumerative Geometry and Tree-Level Gromov-Witten Invariants"
    category = "Mathematics"
    summary = (
        "The text discusses the Kontsevich-Manin formula for enumerating degree d rational curves via Gromov-Witten invariants. "
        "It details the calculation of these invariants using moduli spaces of stable maps and explores their implications in enumerative geometry."
    )
    image_url = fetch_image(title, category, summary)
    print(image_url)