import gradio as gr
import requests
import io
from PIL import Image
import json
import os
import logging
import time
from tqdm import tqdm

# Placeholder class for processed images
class SomeClass:
    def __init__(self):
        self.images = []

# Import your existing modules
from pixelart import Script as PixelArtScript
from postprocessing_pixelart import ScriptPostprocessingUpscale
from utils import downscale_image, limit_colors, resize_image, convert_to_grayscale, convert_to_black_and_white

with open('loras.json', 'r') as f:
    loras = json.load(f)

def update_selection(selected_state: gr.SelectData):
    selected_lora_index = selected_state.index
    selected_lora = loras[selected_lora_index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨"
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        selected_state
    )

def run_lora(prompt, selected_state, pixel_art_options, postprocess_options, progress=gr.Progress(track_tqdm=True)):
    selected_lora_index = selected_state.index
    selected_lora = loras[selected_lora_index]
    api_url = f"https://api-inference.huggingface.co/models/{selected_lora['repo']}"
    payload = {
        "inputs": f"{prompt} {selected_lora['trigger_word']}",
        "parameters": {"negative_prompt": "bad art, ugly, watermark, deformed"},
    }
    response = requests.post(api_url, json=payload)
    if response.status_code == 200:
        original_image = Image.open(io.BytesIO(response.content))
        
        processed = SomeClass()
        processed.images = [original_image]
        
        pixel_art_script = PixelArtScript()
        postprocess_script = ScriptPostprocessingUpscale()
        
        pixel_art_script.postprocess(
            processed,
            **pixel_art_options
        )
        
        postprocess_script.process(
            processed,
            **postprocess_options
        )
        
        refined_image = processed.images[-1]
        
        return original_image, refined_image

def apply_post_processing(image, image_processing_options):
    processed_image = image.copy()
    
    if image_processing_options['downscale'] > 1:
        processed_image = downscale_image(processed_image, image_processing_options['downscale'])
        
    if image_processing_options['limit_colors']:
        processed_image = limit_colors(processed_image)
        
    if image_processing_options['grayscale']:
        processed_image = convert_to_grayscale(processed_image)
        
    if image_processing_options['black_and_white']:
        processed_image = convert_to_black_and_white(processed_image)
        
    return processed_image

with gr.Blocks() as app:
    title = gr.Markdown("# artificialguybr LoRA portfolio")
    description = gr.Markdown("### This is a Pixel Art Generator using SD Loras.")
    selected_state = gr.State()
    
    with gr.Row():
        gallery = gr.Gallery(
            [(item["image"], item["title"]) for item in loras],
            label="LoRA Gallery",
            allow_preview=False,
            columns=3
        )
        
        with gr.Column():
            prompt_title = gr.Markdown("### Click on a LoRA in the gallery to create with it")
            selected_info = gr.Markdown("")
            
            with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1, placeholder="Type a prompt after selecting a LoRA")
                button = gr.Button("Run")
            
            result = gr.Image(interactive=False, label="Generated Image")
            refined_result = gr.Image(interactive=False, label="Refined Generated Image")
            
            # New Output for Post-Processed Image
            post_processed_result = gr.Image(interactive=False, label="Post-Processed Image")
            
            # New UI elements for pixel art options
            with gr.Row():
                pixel_art_options = PixelArtScript().ui(True)
                postprocess_options = ScriptPostprocessingUpscale().ui()
                
            # New UI elements for image processing options
            with gr.Row():
                downscale = gr.Slider(minimum=1, maximum=10, step=1, label="Downscale")
                limit_colors = gr.Checkbox(label="Limit Colors")
                grayscale = gr.Checkbox(label="Grayscale")
                black_and_white = gr.Checkbox(label="Black and White")
                
            image_processing_options = {
                'downscale': downscale,
                'limit_colors': limit_colors,
                'grayscale': grayscale,
                'black_and_white': black_and_white
            }
            
            post_process_button = gr.Button("Apply Post-Processing")
    
    gallery.select(
        update_selection,
        outputs=[prompt, selected_info, selected_state]
    )
    
    prompt.submit(
        fn=run_lora,
        inputs=[prompt, selected_state, pixel_art_options, postprocess_options],
        outputs=[result, refined_result]
    )
    
    post_process_button.click(
        fn=apply_post_processing,
        inputs=[refined_result, image_processing_options],
        outputs=[post_processed_result]
    )

app.queue(max_size=20, concurrency_count=5)
app.launch()