import gradio as gr
import requests
import os
from PIL import Image, ImageOps
from io import BytesIO
from tqdm import tqdm
import time
import numpy as np
import base64
from gradio.components import DownloadButton  # Updated import


# Defining the repository information and the trigger word
repo = "artificialguybr/LineAniRedmond-LinearMangaSDXL-V2"
trigger_word = "lineart,LineAniAF,"

# Hard set prompt template
hard_set_prompt = (
    "minimalist black and white single-line art illustration of [subject], created in one continuous, unbroken stroke from start to finish. "
    "The design uses true black and true white values only, with no gradients or gray tones. The artwork emphasizes clean, connected, and flowing lines, "
    "resulting in a sleek, modern aesthetic. The [subject] is depicted in a dynamic and recognizable pose, optimized for artistic purposes and ready for 3D printing applications."
)

constraints = (
    "MUST: be only black and white, must be single-line art, must never use gray values/shading/tint or hue of any kind, only true black and true white. "
    "NEVER use anything besides a single black line of varying widths as necessary to achieve a result of the user's request."
)

def generate_image(prompt):
    print("Generating image with prompt:", prompt)
    api_url = f"https://api-inference.huggingface.co/models/{repo}"
    token = os.getenv("HF_TOKEN")
    headers = {
        "Authorization": f"Bearer {token}"
    }
    
    # Incorporate the hard set prompt and constraints
    full_prompt = (
        f"{hard_set_prompt.replace('[subject]', prompt)} {constraints} {trigger_word}"
    )
    payload = {
        "inputs": full_prompt,
        "parameters": {
            "negative_prompt": "(non-minimalist, best quality, high quality, normal quality, hires, details, anything but basic 1-line-art inferred from user submitted text, overexposed, underexposed, grayscale, bw, bad photo, bad photography, bad art:1.4), (watermark, signature, text font, username, error, logo, words, letters, digits, autograph, trademark, name:1.2), (blur, blurry, grainy), morbid, ugly, asymmetrical, mutated malformed, mutilated, poorly lit, bad shadow, draft, cropped, out of frame, cut off, censored, jpeg artifacts, out of focus, glitch, duplicate, (airbrushed, cartoon, anime, semi-realistic, cgi, render, blender, digital art, manga, amateur:1.3), , :1.3)",
            "num_inference_steps": 30,
            "scheduler": "DPMSolverMultistepScheduler"
        },
    }

    error_count = 0
    pbar = tqdm(total=None, desc="Loading model")
    while True:
        print("Sending request to API...")
        response = requests.post(api_url, headers=headers, json=payload)
        print("API response status code:", response.status_code)
        if response.status_code == 200:
            print("Image generation successful!")
            # Open the image and convert it to grayscale
            image = Image.open(BytesIO(response.content)).convert("L")
            # Convert the grayscale image to a binary image
            image = image.point(lambda p: 255 if p > 128 else 0, mode='1')
            # Invert the image for potrace (black on white)
            image = ImageOps.invert(image)
            # Convert the PIL image to a numpy array
            array = np.array(image)
            # Trace the bitmap to SVG
            bitmap = potrace.Bitmap(array)
            path = bitmap.trace()
            # Generate SVG data
            svg_data = path.to_svg()
            return svg_data
        elif response.status_code == 503:
            time.sleep(1)
            pbar.update(1)
        elif response.status_code == 500 and error_count < 5:
            time.sleep(1)
            error_count += 1
        else:
            print("API Error:", response.status_code)
            raise Exception(f"API Error: {response.status_code}")

with gr.Blocks() as demo:
    gr.Markdown("# LineArt XL Image Generator")
    prompt = gr.Textbox(lines=2, placeholder="Describe the subject here...")
    generate_button = gr.Button("Generate Image")
    
    with gr.Row():
        svg_display = gr.HTML()
        download_button = gr.Download(label="Download SVG")  # Updated component usage

    def display_svg(svg_data):
        # Encode SVG data in base64 for data URI
        svg_bytes = svg_data.encode('utf-8')
        svg_b64 = base64.b64encode(svg_bytes).decode('utf-8')
        # Create HTML content with an iframe
        html_content = f'<iframe src="data:image/svg+xml;base64,{svg_b64}" width="500" height="500"></iframe>'
        return html_content, svg_bytes

    generate_button.click(
        generate_image,
        inputs=prompt,
        outputs=None
    ).then(
        display_svg,
        inputs=generate_image,
        outputs=[svg_display, download_button]
    )

demo.launch()