import gradio as gr
import time
import logging
import threading
import uuid


# Set up logging
logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

def load_model(model_path):
    logger.info(f"Attempting to load model from {model_path}...")
    try:
        demo = gr.load(model_path)
        logger.info(f"Successfully loaded demo from {model_path}.")
        return demo
    except Exception as e:
        logger.error(f"Failed to load demo from {model_path}: {e}")
        return None

# Load the demo
demo = load_model("models/John6666/lustify-sdxl-nsfwsfw-v2-sdxl")

def generate_image(prompt, max_retries=5):
    call_id = uuid.uuid4()
    thread_id = threading.get_ident()
    logger.info(f"Call ID: {call_id}, Thread ID: {thread_id}, Attempt 1: Starting image generation with prompt: {prompt}")
    if demo is None:
        return None, "Error: Demo failed to load."

    for attempt in range(max_retries):
        try:
            logger.info(f"Attempt {attempt + 1}: Starting image generation with prompt: {prompt}")
            start_time = time.time()
            result = demo(prompt)
            generation_time = time.time() - start_time
            logger.info(f"Generation completed in {generation_time:.2f} seconds")
            logger.info(result)
            return result, f"Image generated successfully\nGeneration time: {generation_time:.2f} seconds"
        except Exception as e:
            logger.error(f"Error on attempt {attempt + 1}: {str(e)}")
            if attempt < max_retries - 1:
                logger.info(f"Retrying... ({attempt + 2}/{max_retries})")
            else:
                logger.error("Max retries reached. Giving up.")
                return None, f"Failed to generate image after {max_retries} attempts. Last error: {str(e)}"
    
    # This line should never be reached, but just in case:
    return None, "Unexpected error occurred"

# Set the root directory as a reusable variable
ROOT_DIR = "./tmp"

# Set a password for file explorer access
FILE_EXPLORER_PASSWORD = "secret123"  # You should use a more secure password in a real application

def show_image(file_path):
    logger.info(f"show_image function called with file_path: {file_path}")
    if file_path:
        result = file_path[0] if isinstance(file_path, list) else file_path
        logger.info(f"Returning file path: {result}")
        return result, f"Selected image: {result}"
    logger.warning("No file path provided")
    return None, "No image selected"

def check_seed(seed):
    logger.info("Seed check initiated")
    if seed == FILE_EXPLORER_PASSWORD:
        logger.info("Correct seed entered")
        return gr.update(visible=True)
    else:
        logger.info("Incorrect seed entered")
        return gr.update(visible=False)

logger.info("Starting Gradio app")
with gr.Blocks() as demo:
    gr.Markdown("# Text-to-Image Generation and File Explorer")
    gr.Markdown("Enter a prompt to generate an image or enter the correct seed to access additional features.")

    with gr.Row():
        with gr.Column(scale=1):
            prompt_input = gr.Textbox(label="Enter your prompt")
            seed_input = gr.Textbox(label="Enter seed", type="password")
            generate_button = gr.Button("Generate Image")

            file_explorer = gr.FileExplorer(
                glob="**/*.webp",
                file_count="single",
                root_dir=ROOT_DIR,
                label="Select a WebP image",
                visible=False
            )
            output_info = gr.Textbox(label="Output Information", lines=3)

        with gr.Column(scale=2):
            shared_image = gr.Image(label="Image Display")

    # Set up event handlers
    generate_button.click(
        generate_image,
        inputs=prompt_input,
        outputs=[shared_image, output_info]
    )

    seed_input.change(
        check_seed,
        inputs=seed_input,
        outputs=file_explorer
    )

    file_explorer.change(
        show_image,
        inputs=file_explorer,
        outputs=[shared_image, output_info]
    )

logger.info("Launching Gradio app")
demo.launch(max_threads=1)
logger.info("Gradio app launched")