import os
import gradio as gr
import numpy as np
import random
from huggingface_hub import AsyncInferenceClient
from translatepy import Translator
from gradio_client import Client, handle_file
from PIL import Image
from pathlib import Path
MAX_SEED = np.iinfo(np.int32).max

HF_TOKEN = os.getenv('HF_TOKEN')
HF_TOKEN_UPSCALER = os.getenv('HF_TOKEN')

css2="""
/* Apply dark theme (black background) */
body {
    background-color: #000000;
    color: #FFFFFF;
}

/* Style the Gradio interface */
.gradio-container {
    background-color: #000000;
    border: 2px solid #FFFFFF;
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.1);
}

/* Title and markdown text */
.gradio-markdown h1, .gradio-markdown h2, .gradio-markdown h3 {
    color: #FFFFFF;
}

/* Input boxes (e.g., Textbox) */
.gradio-textbox input, .gradio-textbox textarea {
    background-color: #222222;
    color: #FFFFFF;
    border: 2px solid #444444;
    border-radius: 8px;
    padding: 10px;
    font-size: 16px;
    box-shadow: 0 0 5px rgba(255, 255, 255, 0.2);
    transition: 0.3s ease-in-out;
}

.gradio-textbox input:focus, .gradio-textbox textarea:focus {
    border-color: #ff00ff;
    box-shadow: 0 0 10px rgba(255, 0, 255, 0.7);
}

/* Buttons */
.gradio-button {
    background: linear-gradient(45deg, #ff007f, #ff00ff, #00ff00, #00ffff, #0000ff, #ff8c00);
    color: white;
    font-weight: bold;
    border: 2px solid #444444;
    border-radius: 10px;
    padding: 12px 20px;
    box-shadow: 0 0 15px rgba(255, 255, 255, 0.3);
    cursor: pointer;
    transition: 0.3s ease-in-out;
    font-size: 16px;
    text-transform: uppercase;
}

.gradio-button:hover {
    background: linear-gradient(45deg, #ff8c00, #00ffff, #ff00ff, #ff007f, #0000ff, #00ff00);
    box-shadow: 0 0 20px rgba(255, 255, 255, 0.5);
}

/* Dropdown */
.gradio-dropdown select {
    background-color: #222222;
    color: #FFFFFF;
    border: 2px solid #444444;
    border-radius: 8px;
    padding: 8px 12px;
    box-shadow: 0 0 5px rgba(255, 255, 255, 0.2);
}

.gradio-dropdown select:focus {
    border-color: #ff00ff;
    box-shadow: 0 0 10px rgba(255, 0, 255, 0.7);
}

/* Chatbot box */
.gradio-chatbot {
    background-color: #222222;
    border: 2px solid #444444;
    color: #FFFFFF;
    padding: 15px;
    border-radius: 12px;
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.2);
}

/* Slider */
.gradio-slider input {
    background-color: #222222;
    border: 2px solid #444444;
    color: #FFFFFF;
    border-radius: 8px;
    padding: 10px;
    box-shadow: 0 0 5px rgba(255, 255, 255, 0.2);
}

.gradio-slider input:focus {
    border-color: #ff00ff;
    box-shadow: 0 0 10px rgba(255, 0, 255, 0.7);
}

/* Accordion */
.gradio-accordion {
    background-color: #222222;
    border: 2px solid #444444;
    color: #FFFFFF;
    padding: 15px;
    border-radius: 12px;
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.2);
}

.gradio-accordion-button {
    background-color: #444444;
    color: #FFFFFF;
    border: none;
    border-radius: 8px;
    padding: 8px 16px;
    box-shadow: 0 0 5px rgba(255, 255, 255, 0.1);
    transition: 0.3s ease-in-out;
}

.gradio-accordion-button:hover {
    background-color: #ff00ff;
    box-shadow: 0 0 15px rgba(255, 0, 255, 0.5);
}

/* General hover effect for all Gradio elements */
.gradio-container *:hover {
    box-shadow: 0 0 10px rgba(255, 255, 255, 0.5);
}

/* Animation for glowing neon effect */
@keyframes neon {
    0% {
        text-shadow: 0 0 5px #ff0000, 0 0 10px #ff0000, 0 0 15px #ff0000, 0 0 20px #ff0000, 0 0 25px #ff0000, 0 0 30px #ff0000;
    }
    50% {
        text-shadow: 0 0 5px #00ff00, 0 0 10px #00ff00, 0 0 15px #00ff00, 0 0 20px #00ff00, 0 0 25px #00ff00, 0 0 30px #00ff00;
    }
    100% {
        text-shadow: 0 0 5px #0000ff, 0 0 10px #0000ff, 0 0 15px #0000ff, 0 0 20px #0000ff, 0 0 25px #0000ff, 0 0 30px #0000ff;
    }
}

/* Apply glowing text effect */
.gradio-container h1, .gradio-container h2, .gradio-container h3, .gradio-container p {
    animation: neon 1.5s ease-in-out infinite alternate;
}

"""

# Define base models
base_models = [
    "black-forest-labs/FLUX.1-schnell", 
    "black-forest-labs/FLUX.1-DEV", 
    "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro"
]

def load_local_loras(lora_directory="lora_models"):
    """Load loras from local safetensor files"""
    loras_list_custom = []
    
    if not os.path.exists(lora_directory):
        os.makedirs(lora_directory)
        print(f"[-] Created lora directory: {lora_directory}")
        
    lora_files = list(Path(lora_directory).glob("*.safetensors"))
    
    for lora_file in lora_files:
        lora_name = lora_file.stem
        lora_path = str(lora_file.absolute())
        loras_list_custom.append({
            "name": lora_name,
            "path": lora_path
        })
    
    print(f"[-] Loaded {len(loras_list_custom)} local loras")
    return loras_list_custom

# Function to enable LoRA if selected
def enable_lora(lora_path, basemodel):
    print(f"[-] Determining model: LoRA {'enabled' if lora_path else 'disabled'}, base model: {basemodel}")
    if not lora_path:
        return basemodel
        
    # Configure model with local lora
    return {
        "model": basemodel,
        "lora_weights": lora_path,
        "lora_scale": 0.75  # Adjust this value as needed
    }

# Function to upscale image
def get_upscale_finegrain(prompt, img_path, upscale_factor):
    try:
        print(f"[-] Starting upscaling process with factor {upscale_factor} for image {img_path}")
        client = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER)
        result = client.predict(
            input_image=handle_file(img_path), 
            prompt=prompt, 
            negative_prompt="worst quality, low quality, normal quality",
            upscale_factor=upscale_factor,
            controlnet_scale=0.6,
            controlnet_decay=1,
            condition_scale=6,
            denoise_strength=0.35, 
            num_inference_steps=18,
            solver="DDIM", 
            api_name="/process"
        )
        print(f"[-] Upscaling successful.")
        return result[1]  # Return upscale image path
    except Exception as e:
        print(f"[-] Error scaling image: {e}")
        return None

# Function to generate image
async def generate_image(prompt, model_config, lora_word, width, height, scales, steps, seed):
    try:
        if seed == -1:
            seed = random.randint(0, MAX_SEED)
        seed = int(seed)

        print(f"[-] Translating prompt: {prompt}")
        text = str(Translator().translate(prompt, 'English')) + "," + lora_word
        
        print(f"[-] Generating image with prompt: {text}")
        client = AsyncInferenceClient(token=HF_TOKEN)
        
        # Handle both simple model string and lora config
        if isinstance(model_config, dict):
            print(f"[-] Using model with LoRA: {model_config}")
            image = await client.text_to_image(
                prompt=text,
                height=height,
                width=width,
                guidance_scale=scales,
                num_inference_steps=steps,
                model=model_config["model"],
                lora_weights=model_config["lora_weights"],
                lora_scale=model_config["lora_scale"]
            )
        else:
            print(f"[-] Using base model: {model_config}")
            image = await client.text_to_image(
                prompt=text,
                height=height,
                width=width,
                guidance_scale=scales,
                num_inference_steps=steps,
                model=model_config
            )
            
        return image, seed
    except Exception as e:
        print(f"[-] Error generating image: {e}")
        return None, None

# Main function to generate images and optionally upscale
async def gen(prompt, basemodel, width, height, scales, steps, seed, upscale_factor, process_upscale, lora_model, process_lora):
    print(f"[-] Starting image generation with prompt: {prompt}")
    
    model = enable_lora(lora_model, basemodel) if process_lora else basemodel
    print(f"[-] Using model: {model}")

    image, seed = await generate_image(prompt, model, "", width, height, scales, steps, seed)
    
    if image is None:
        print("[-] Image generation failed.")
        return []

    image_path = "temp_image.jpg"
    print(f"[-] Saving temporary image to: {image_path}")
    image.save(image_path, format="JPEG")

    upscale_image_path = None
    if process_upscale:
        print(f"[-] Processing upscaling with factor: {upscale_factor}")
        upscale_image_path = get_upscale_finegrain(prompt, image_path, upscale_factor)
        if upscale_image_path is not None and os.path.exists(upscale_image_path):
            print(f"[-] Upscaling complete. Image saved at: {upscale_image_path}")
            return [image_path, upscale_image_path]  # Return both images
        else:
            print("[-] Upscaling failed, upscaled image path not found.")

    return [image_path]

# Load local loras
local_loras = load_local_loras()

# Creating Gradio interface
with gr.Blocks(css=css2, theme=IndonesiaTheme()) as WallpaperFluxMaker:
    gr.HTML('<div id="banner">✨ Flux MultiMode Generator + Upscaler ✨</div>')

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            output_res = gr.Gallery(
                label="⚡ Flux / Upscaled Image ⚡", 
                elem_id="output-res",
                columns=2,
                height="auto"
            )

        with gr.Row():
            with gr.Column(scale=1, elem_id="col-left"):
                prompt = gr.Textbox(
                    label="📜 Description", 
                    placeholder="Write your prompt in any language, it will be translated to English.",
                    elem_id="textbox-prompt"
                )

                basemodel_choice = gr.Dropdown(
                    label="🖼️ Select Model", 
                    choices=base_models, 
                    value=base_models[0]
                )

                # Updated to use local lora paths
                lora_model_choice = gr.Dropdown(
                    label="🎨 Select LoRA",
                    choices=[lora["path"] for lora in local_loras],
                    value=local_loras[0]["path"] if local_loras else None
                )

                process_lora = gr.Checkbox(label="🎨 Enable LoRA")
                process_upscale = gr.Checkbox(label="🔍 Enable Upscaling")
                upscale_factor = gr.Radio(
                    label="🔍 Upscale Factor",
                    choices=[2, 4, 8],
                    value=2
                )

            with gr.Column(scale=1, elem_id="col-right"):
                with gr.Accordion(label="⚙️ Advanced Options", open=True):
                    width = gr.Slider(label="Width", minimum=512, maximum=1280, step=8, value=1280)
                    height = gr.Slider(label="Height", minimum=512, maximum=1280, step=8, value=768)
                    scales = gr.Slider(label="Scale", minimum=1, maximum=20, step=1, value=8)
                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, step=1, value=8)
                    seed = gr.Number(label="Seed", value=-1)

        btn = gr.Button("🚀 Generate Image", elem_id="generate-btn")
        
        btn.click(
            fn=gen,
            inputs=[
                prompt, basemodel_choice, width, height, scales, steps, seed,
                upscale_factor, process_upscale, lora_model_choice, process_lora
            ],
            outputs=output_res
        )

WallpaperFluxMaker.queue(api_open=True).launch(show_api=True)