import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import os
import time
import json
from datetime import datetime

# Load API Token from environment variable
API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment variable

# Hugging Face Inference API URL
API_URL = "https://api-inference.huggingface.co/models/enhanceaiteam/Flux-uncensored"

class ImageGenerator:
    def __init__(self):
        self.headers = {"Authorization": f"Bearer {API_TOKEN}"}
        self.generation_history = []
    
    def generate_image(self, prompt, negative_prompt="", num_inference_steps=50, guidance_scale=7.5, seed=None, progress=gr.Progress()):
        """
        Generate an image with advanced parameters
        """
        if not API_TOKEN:
            return None, "Error: HF_API_TOKEN environment variable not set"
        
        if not prompt or prompt.strip() == "":
            return None, "Error: Please enter a prompt"
        
        # Prepare the payload with additional parameters
        payload = {
            "inputs": prompt,
            "parameters": {
                "negative_prompt": negative_prompt if negative_prompt else None,
                "num_inference_steps": num_inference_steps,
                "guidance_scale": guidance_scale,
                "seed": seed if seed else None
            }
        }
        
        # Remove None values
        payload["parameters"] = {k: v for k, v in payload["parameters"].items() if v is not None}
        
        try:
            progress(0.1, desc="Initializing generation...")
            
            # Make API request
            response = requests.post(API_URL, headers=self.headers, json=payload, timeout=60)
            
            progress(0.5, desc="Processing response...")
            
            if response.status_code == 200:
                # Parse the response
                image_bytes = BytesIO(response.content)
                image = Image.open(image_bytes)
                
                # Save to history
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                self.generation_history.append({
                    "timestamp": timestamp,
                    "prompt": prompt,
                    "negative_prompt": negative_prompt,
                    "seed": seed
                })
                
                progress(1.0, desc="Complete!")
                return image, f"Success! Image generated at {timestamp}"
            
            elif response.status_code == 503:
                # Model is loading
                progress(0.3, desc="Model is loading, please wait...")
                time.sleep(5)
                return self.generate_image(prompt, negative_prompt, num_inference_steps, 
                                         guidance_scale, seed, progress)
            
            else:
                error_msg = f"Error {response.status_code}: {response.text}"
                return None, error_msg
                
        except requests.exceptions.Timeout:
            return None, "Error: Request timed out. Please try again."
        except requests.exceptions.ConnectionError:
            return None, "Error: Connection error. Please check your internet connection."
        except Exception as e:
            return None, f"Error: {str(e)}"
    
    def get_history(self):
        """Return generation history as formatted text"""
        if not self.generation_history:
            return "No generations yet"
        
        history_text = "### Generation History\n\n"
        for i, item in enumerate(reversed(self.generation_history[-10:]), 1):
            history_text += f"{i}. **{item['timestamp']}**\n"
            history_text += f"   Prompt: {item['prompt'][:50]}...\n"
            if item['negative_prompt']:
                history_text += f"   Negative: {item['negative_prompt'][:30]}...\n"
            if item['seed']:
                history_text += f"   Seed: {item['seed']}\n"
            history_text += "\n"
        
        return history_text
    
    def clear_history(self):
        """Clear generation history"""
        self.generation_history = []
        return "History cleared"

def create_enhanced_ui():
    """Create an enhanced Gradio interface with more features"""
    
    # Initialize generator
    generator = ImageGenerator()
    
    # Custom CSS for better styling
    custom_css = """
    .gradio-container {
        max-width: 1200px !important;
        margin: auto !important;
    }
    .generate-btn {
        background: linear-gradient(90deg, #6366f1 0%, #8b5cf6 100%) !important;
        color: white !important;
        border: none !important;
    }
    .generate-btn:hover {
        background: linear-gradient(90deg, #4f46e5 0%, #7c3aed 100%) !important;
    }
    .history-panel {
        background: #f3f4f6;
        border-radius: 8px;
        padding: 10px;
    }
    """
    
    with gr.Blocks(theme="hev832/Applio", css=custom_css, title="Flux Uncensored Enhanced") as ui:
        gr.Markdown("""
        # 🎨 Flux Uncensored Image Generator
        ### Advanced image generation with Hugging Face
        """)
        
        with gr.Tabs():
            # Main Generation Tab
            with gr.TabItem("Generate"):
                with gr.Row():
                    with gr.Column(scale=2):
                        # Main input
                        prompt = gr.Textbox(
                            label="📝 Prompt",
                            placeholder="Describe the image you want to generate in detail...",
                            lines=4
                        )
                        
                        # Advanced options (collapsible)
                        with gr.Accordion("⚙️ Advanced Options", open=False):
                            negative_prompt = gr.Textbox(
                                label="Negative Prompt",
                                placeholder="What to avoid in the image...",
                                lines=2
                            )
                            
                            with gr.Row():
                                steps = gr.Slider(
                                    label="Inference Steps",
                                    minimum=20,
                                    maximum=100,
                                    value=50,
                                    step=1
                                )
                                
                                guidance = gr.Slider(
                                    label="Guidance Scale",
                                    minimum=1.0,
                                    maximum=20.0,
                                    value=7.5,
                                    step=0.5
                                )
                            
                            seed = gr.Number(
                                label="Seed (optional)",
                                value=None,
                                precision=0
                            )
                        
                        # Generate button
                        generate_btn = gr.Button(
                            "🎨 Generate Image",
                            variant="primary",
                            elem_classes="generate-btn"
                        )
                    
                    with gr.Column(scale=1):
                        # Status and info
                        status = gr.Textbox(
                            label="Status",
                            value="Ready to generate",
                            interactive=False
                        )
                
                # Output
                with gr.Row():
                    output_image = gr.Image(
                        label="Generated Image",
                        type="pil",
                        height=400
                    )
                
                # Download button
                with gr.Row():
                    download_btn = gr.File(
                        label="Download Image",
                        interactive=False,
                        visible=False
                    )
            
            # History Tab
            with gr.TabItem("📜 History"):
                with gr.Row():
                    history_text = gr.Markdown("No generations yet")
                
                with gr.Row():
                    refresh_history_btn = gr.Button("🔄 Refresh History")
                    clear_history_btn = gr.Button("🗑️ Clear History", variant="stop")
            
            # Info Tab
            with gr.TabItem("ℹ️ Info"):
                gr.Markdown("""
                ## About Flux Uncensored
                
                This is an unofficial Gradio interface for the Flux Uncensored model on Hugging Face.
                
                ### Tips for better results:
                - Be specific and detailed in your prompts
                - Use negative prompts to avoid unwanted elements
                - Experiment with different guidance scales (7.5 is a good starting point)
                - More inference steps generally produce better quality but take longer
                
                ### Parameters explained:
                - **Prompt**: What you want to generate
                - **Negative Prompt**: What you don't want in the image
                - **Inference Steps**: Number of denoising steps (higher = better quality but slower)
                - **Guidance Scale**: How closely to follow the prompt (higher = more adherence)
                - **Seed**: Random seed for reproducibility
                
                ### Note:
                Make sure to set your `HF_API_TOKEN` environment variable before running.
                """)
        
        # Event handlers
        def on_generate(prompt, negative_prompt, steps, guidance, seed):
            img, msg = generator.generate_image(
                prompt, 
                negative_prompt, 
                steps, 
                guidance, 
                seed if seed != 0 else None
            )
            if img:
                return img, msg, gr.update(visible=True, value=img)
            return None, msg, gr.update(visible=False)
        
        generate_btn.click(
            fn=on_generate,
            inputs=[prompt, negative_prompt, steps, guidance, seed],
            outputs=[output_image, status, download_btn]
        )
        
        # History handlers
        def update_history():
            return generator.get_history()
        
        refresh_history_btn.click(
            fn=update_history,
            outputs=[history_text]
        )
        
        clear_history_btn.click(
            fn=generator.clear_history,
            outputs=[history_text]
        ).then(
            fn=lambda: "History cleared",
            outputs=[history_text]
        )
        
        # Auto-refresh history when generating
        generate_btn.click(
            fn=update_history,
            outputs=[history_text]
        )
        
        # Clear inputs
        def clear_inputs():
            return "", "", 50, 7.5, None
        
        with gr.Row():
            clear_btn = gr.Button("🗑️ Clear Inputs")
            clear_btn.click(
                fn=clear_inputs,
                outputs=[prompt, negative_prompt, steps, guidance, seed]
            )
    
    return ui

# Run the interface
if __name__ == "__main__":
    # Check for API token
    if not API_TOKEN:
        print("⚠️  Warning: HF_API_TOKEN environment variable is not set!")
        print("Please set it using: export HF_API_TOKEN='your_token_here'")
    
    # Create and launch the UI
    ui = create_enhanced_ui()
    ui.launch(
        server_name="0.0.0.0",  # Allow external connections
        server_port=7860,        # Default Gradio port
        share=False,             # Set to True to create a public link
        debug=True
    )