import gradio as gr
import asyncio
import aiohttp
import time
from datetime import datetime
import plotly.graph_objects as go
from typing import Dict, List
import os
from dotenv import load_dotenv
import json
from PIL import Image, ImageDraw, ImageFont
import uuid
import threading

# Load environment variables first
load_dotenv()

# Constants
API_BASE_URL = "https://api.wavespeed.ai/api/v2"
API_KEY = os.getenv("WAVESPEED_API_KEY")  # Move API_KEY to global scope

if not API_KEY:
    raise ValueError("WAVESPEED_API_KEY not found in environment variables")

# Rest of constants
BACKENDS = {
    "flux-dev": {
        "endpoint": f"{API_BASE_URL}/wavespeed-ai/flux-dev-ultra-fast",
        "name": "Flux-dev",
        "color": "#FF9800",
    },
    "hidream-dev": {
        "endpoint": f"{API_BASE_URL}/wavespeed-ai/hidream-i1-dev",
        "name": "HiDream-dev",
        "color": "#2196F3",
    },
    "hidream-full": {
        "endpoint": f"{API_BASE_URL}/wavespeed-ai/hidream-i1-full",
        "name": "HiDream-full",
        "color": "#4CAF50",
    },
}


class BackendStatus:

    def __init__(self):
        self.reset()
        self.history: List[Dict] = []

    def reset(self):
        self.status = "idle"
        self.progress = 0
        self.start_time = None
        self.end_time = None

    def start(self):
        self.status = "processing"
        self.progress = 0
        self.start_time = time.time()
        self.end_time = None

    def complete(self):
        self.status = "completed"
        self.progress = 100
        self.end_time = time.time()
        self.history.append({
            "timestamp": datetime.now(),
            "duration": self.end_time - self.start_time
        })

    def fail(self):
        self.status = "failed"
        self.end_time = time.time()


class SessionManager:
    _instances = {}
    _lock = threading.Lock()

    @classmethod
    def get_manager(cls, session_id=None):
        if session_id is None:
            session_id = str(uuid.uuid4())

        with cls._lock:
            if session_id not in cls._instances:
                cls._instances[session_id] = GenerationManager()
            return session_id, cls._instances[session_id]

    @classmethod
    def cleanup_old_sessions(cls, max_age=3600):  # 1 hour default
        current_time = time.time()
        with cls._lock:
            to_remove = []
            for session_id, manager in cls._instances.items():
                if (hasattr(manager, "last_activity")
                        and current_time - manager.last_activity > max_age):
                    to_remove.append(session_id)

            for session_id in to_remove:
                del cls._instances[session_id]


class GenerationManager:

    def __init__(self):
        self.backend_statuses = {
            backend: BackendStatus()
            for backend in BACKENDS
        }
        self.last_activity = time.time()

    def update_activity(self):
        self.last_activity = time.time()

    def get_performance_plot(self):
        fig = go.Figure()

        has_data = False

        for backend, status in self.backend_statuses.items():
            durations = [h["duration"] for h in status.history]
            if durations:
                has_data = True
                avg_duration = sum(durations) / len(durations)
                # Use bar chart instead of box plot
                fig.add_trace(
                    go.Bar(
                        y=[avg_duration],  # Average duration
                        x=[BACKENDS[backend]["name"]],  # Backend name
                        name=BACKENDS[backend]["name"],
                        marker_color=BACKENDS[backend]["color"],
                        text=[f"{avg_duration:.2f}s"],  # Show time in seconds
                        textposition="auto",
                        width=[0.5],  # Make bars narrower
                    ))

        # Set a minimum y-axis range if we have data
        if has_data:
            max_duration = max([
                max([h["duration"] for h in status.history] or [0])
                for status in self.backend_statuses.values()
            ])
            # Add 20% padding to the top
            y_max = max_duration * 1.2
            # Ensure the y-axis always starts at 0
            fig.update_yaxes(range=[0, y_max])

        fig.update_layout(
            title="Average Generation Time",
            yaxis_title="Seconds",
            xaxis_title="",
            showlegend=False,
            template="simple_white",
            height=400,  # Increase height
            margin=dict(l=50, r=50, t=50, b=50),  # Add margins
            font=dict(size=14),  # Larger font
        )

        # Make sure we have a valid figure even if no data
        if not has_data:
            fig.add_annotation(
                text="No timing data available yet",
                xref="paper",
                yref="paper",
                x=0.5,
                y=0.5,
                showarrow=False,
                font=dict(size=16),
            )

        return fig

    async def submit_task(self, backend: str, prompt: str) -> str:
        status = self.backend_statuses[backend]
        status.start()

        try:
            url = BACKENDS[backend]["endpoint"]
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {API_KEY}",
            }
            payload = {
                "prompt": prompt,
                "enable_safety_checker": True,
                "enable_base64_output": True,  # Enable base64 output
                "size": "1024*1024",
                "seed": -1,
            }

            if backend == "flux-dev":
                payload.update({
                    "guidance_scale": 3.5,
                    "num_images": 1,
                    "num_inference_steps": 28,
                    "strength": 0.8,
                })

            print(f"Submitting task to {backend}")
            print(f"URL: {url}")
            print(f"Payload: {json.dumps(payload, indent=2)}")

            # Use aiohttp instead of requests for async
            async with aiohttp.ClientSession() as session:
                async with session.post(url, headers=headers,
                                        json=payload) as response:
                    if response.status == 200:
                        result = await response.json()
                        request_id = result["data"]["id"]
                        print(
                            f"Task submitted successfully. Request ID: {request_id}"
                        )
                        return request_id
                    else:
                        text = await response.text()
                        raise Exception(
                            f"API error: {response.status}, {text}")

        except Exception as e:
            status.fail()
            raise Exception(f"Failed to submit task: {str(e)}")

    # Add this method to reset history
    def reset_history(self):
        """Reset history for all backends"""
        for status in self.backend_statuses.values():
            status.history = []  # Clear history data
        return self


# Helper function to create error images as data URIs
def create_error_image(backend, error_message):
    try:
        import base64
        from io import BytesIO

        # Create an in-memory image
        img = Image.new("RGB", (512, 512), color="#ffdddd")
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.truetype("Arial", 20)
        except:
            font = ImageFont.load_default()

        # Wrap and draw error message
        words = error_message.split(" ")
        lines = []
        line = ""
        for word in words:
            if len(line + word) < 40:
                line += word + " "
            else:
                lines.append(line)
                line = word + " "
        if line:
            lines.append(line)

        y_position = 100
        for line in lines:
            draw.text((50, y_position), line, fill="black", font=font)
            y_position += 30

        # Save to a BytesIO object instead of a file
        buffer = BytesIO()
        img.save(buffer, format="JPEG")
        img_bytes = buffer.getvalue()

        # Convert to base64 and return as data URI
        return f"data:image/jpeg;base64,{base64.b64encode(img_bytes).decode('utf-8')}"
    except Exception as e:
        print(f"Failed to create error image: {e}")
        # Return a simple error message as fallback
        return "Error: " + error_message


# Fix the poll_once function to accept a manager parameter
async def poll_once(manager, backend, request_id):
    """Poll once and return result if complete, otherwise None"""
    headers = {"Authorization": f"Bearer {API_KEY}"}
    url = f"{API_BASE_URL}/predictions/{request_id}/result"

    async with aiohttp.ClientSession() as session:
        async with session.get(url, headers=headers) as response:
            if response.status == 200:
                result = await response.json()
                data = result["data"]
                current_status = data["status"]

                if current_status == "completed":
                    # IMPORTANT: Update status BEFORE returning - using the passed manager
                    manager.backend_statuses[backend].complete()
                    manager.update_activity()

                    # Handle base64 output
                    output = data["outputs"][0]

                    # Check if it's a base64 string or URL
                    if isinstance(output, str) and output.startswith("http"):
                        # It's a URL - return as is
                        return output
                    else:
                        # It's base64 data - format it as a data URI if needed
                        try:
                            # Format as data URI for Gradio to display directly
                            if isinstance(
                                    output, str
                            ) and not output.startswith("data:image"):
                                # Convert raw base64 to data URI format
                                return f"data:image/jpeg;base64,{output}"
                            else:
                                # Already in data URI format
                                return output
                        except Exception as e:
                            print(f"Error processing base64 image: {e}")
                            raise Exception(
                                f"Failed to process base64 image: {str(e)}")

                elif current_status == "failed":
                    manager.backend_statuses[backend].fail()
                    manager.update_activity()
                    error = data.get("error", "Unknown error")
                    raise Exception(error)

                # Still processing
                return None
            else:
                raise Exception(f"Poll error: {response.status}")


# Store recent generations
recent_generations = []

# Example prompts
example_prompts = [
    "A deep sea diver exploring an underwater city ruins, using a palette of deep blues and silvers.",
    "A Martian greenhouse complex that uses genetically modified crops designed to thrive in low gravity. Outside, rovers fitted with AI guidance systems patrol dusty red plains, ensuring each pressurized dome remains airtight.",
    "A sleek, futuristic sports car with glowing blue accents, racing through a virtual reality landscape, 3D render",
]


# Use a state variable to store session ID
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    session_id = gr.State(None)  # Add this to store session ID

    gr.Markdown("# 🌊 WaveSpeedAI HiDream Arena")

    # Add the introduction with link to WaveSpeedAI
    gr.Markdown(
        """
        [WaveSpeedAI](https://wavespeed.ai/) is the global pioneer in accelerating AI-powered video and image generation.
        "Our in-house inference accelerator provides lossless speedup on image & video generation based on our rich inference optimization software stack, including our in-house inference compiler, CUDA kernel libraries and parallel computing libraries."
        """
    )
    gr.Markdown(
        """
        This demo showcases the performance and outputs of leading image generation models, including HiDream and Flux, on our accelerated inference platform.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            example_dropdown = gr.Dropdown(
                choices=example_prompts,
                label="Choose an example prompt",
                interactive=True
            )
            input_text = gr.Textbox(
                example_prompts[0],
                label="Enter your prompt",
                placeholder="Type here...",
                lines=3,
            )
        with gr.Column(scale=1):
            generate_btn = gr.Button("Generate", variant="primary")

    example_dropdown.change(lambda ex: ex, inputs=[example_dropdown], outputs=[input_text])

    # Two status boxes - small (default) and big (during generation)
    small_status_box = gr.Markdown("Ready to generate images",
                                   elem_id="small-status")

    # Big status box in its own row with styling
    with gr.Row(elem_id="big-status-row"):
        big_status_box = gr.Markdown("",
                                     elem_id="big-status",
                                     visible=False,
                                     elem_classes="big-status-box")

    with gr.Row():
        with gr.Column():
            draft_output = gr.Image(label="Flux-dev")
        with gr.Column():
            quick_output = gr.Image(label="HiDream-dev")
        with gr.Column():
            best_output = gr.Image(label="HiDream-full")

    performance_plot = gr.Plot(label="Performance Metrics")

    with gr.Accordion("Recent Generations (last 8)", open=False):
        recent_gallery = gr.Gallery(label="Prompt and Output", columns=3)

    def update_recent_gallery(prompt, results):
        recent_generations.append({
            "prompt": prompt,
            "flux-dev": results["flux-dev"],
            "hidream-dev": results["hidream-dev"],
            "hidream-full": results["hidream-full"],
        })
        if len(recent_generations) > 8:
            recent_generations.pop(0)
        gallery_items = []
        for r in reversed(recent_generations):
            gallery_items.append(
                (r["flux-dev"], f"FLUX-dev: {r['prompt']}"))
            gallery_items.append(
                (r["hidream-dev"], f"HiDream-dev: {r['prompt']}"))
            gallery_items.append(
                (r["hidream-full"], f"HiDream-full: {r['prompt']}"))
        return gr.update(value=gallery_items)

    # Add custom CSS for the big status box
    css = """
    #big-status-row {
        margin: 20px 0;
    }
    #big-status {
        font-size: 28px;  /* Even larger font size */
        font-weight: bold;
        padding: 30px;    /* More padding */
        background-color: #0D47A1;  /* Deeper blue background */
        color: white;     /* White text */
        border-radius: 10px;
        text-align: center;
        margin: 0 auto;
        box-shadow: 0 6px 12px rgba(0, 0, 0, 0.2);  /* Stronger shadow */
        animation: deep-breath 3s infinite;  /* Slower, deeper breathing animation */
        width: 100%;      /* Full width */
        max-width: 800px; /* Maximum width */
        transition: all 0.3s ease;  /* Smooth transitions */
        border-left: 6px solid #64B5F6;  /* Add a colored border */
        border-right: 6px solid #64B5F6;  /* Add a colored border */
    }

    /* Deeper breathing animation */
    @keyframes deep-breath {
        0% { 
            opacity: 0.7; 
            transform: scale(0.98);
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
        }
        50% { 
            opacity: 1; 
            transform: scale(1.01);
            box-shadow: 0 8px 16px rgba(0, 0, 0, 0.3);
        }
        100% { 
            opacity: 0.7; 
            transform: scale(0.98);
            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
        }
    }
    """
    gr.HTML(f"<style>{css}</style>")

    # Update the generation function to use session manager
    async def generate_all_backends_with_status_boxes(prompt,
                                                      current_session_id):
        """Generate images with big status box during generation"""
        # Get or create a session manager
        session_id, manager = SessionManager.get_manager(current_session_id)
        manager.update_activity()

        # IMPORTANT: Reset history when starting a new generation
        if prompt and prompt.strip() != "":
            manager.reset_history()  # Clear previous performance metrics

        if not prompt or prompt.strip() == "":
            # Handle empty prompt case
            yield (
                "⚠️ Please enter a prompt first",
                "⚠️ Please enter a prompt first",
                gr.update(visible=True),
                gr.update(visible=False),
                None,
                None,
                None,
                None,
                session_id,  # Return the session ID
                None,
            )
            return

        # Status message
        status_message = f"🔄 PROCESSING: '{prompt}'"

        # Initial state - clear all images, show big status box
        yield (
            status_message,
            status_message,
            gr.update(visible=True),
            gr.update(visible=False),
            None,
            None,
            None,
            None,
            session_id,  # Return the session ID
            None,
        )

        # For production mode:
        completed_backends = set()
        results = {"flux-dev": None, "hidream-dev": None, "hidream-full": None}

        try:
            # Submit all tasks
            request_ids = {}
            for backend in BACKENDS:
                try:
                    request_id = await manager.submit_task(backend, prompt)
                    request_ids[backend] = request_id
                except Exception as e:
                    # Handle submission error
                    print(f"Error submitting task for {backend}: {e}")
                    results[backend] = create_error_image(backend, str(e))
                    completed_backends.add(backend)

            # Poll all backends until they complete
            max_poll_attempts = 300
            poll_attempt = 0

            # Main polling loop
            while len(completed_backends
                      ) < 3 and poll_attempt < max_poll_attempts:
                poll_attempt += 1

                # Poll each pending backend
                for backend in list(BACKENDS.keys()):
                    if backend in completed_backends:
                        continue

                    try:
                        # Only do actual API calls every few attempts to reduce load
                        if poll_attempt % 2 == 0 or backend == "flux-dev":
                            # Use the session manager instead of global manager
                            result = await poll_once(manager, backend,
                                                     request_ids[backend])
                            if result:  # Backend completed
                                results[backend] = result
                                completed_backends.add(backend)

                                # Yield updated state when an image completes
                                yield (
                                    status_message,
                                    status_message,
                                    gr.update(visible=True),
                                    gr.update(visible=False),
                                    results["flux-dev"],
                                    results["hidream-dev"],
                                    results["hidream-full"],
                                    (manager.get_performance_plot()
                                     if any(completed_backends) else None),
                                    session_id,
                                    None,
                                )
                    except Exception as e:
                        print(f"Error polling {backend}: {str(e)}")

                # Wait between poll attempts
                await asyncio.sleep(0.1)

            # Final status
            final_status = ("✅ All generations completed!"
                            if len(completed_backends) == 3 else
                            "⚠️ Some generations timed out")

            gallery_update = update_recent_gallery(prompt, results)

            # Final yield
            yield (
                final_status,
                final_status,
                gr.update(visible=False),
                gr.update(visible=True),
                results["flux-dev"],
                results["hidream-dev"],
                results["hidream-full"],
                manager.get_performance_plot(),
                session_id,
                gallery_update,
            )

        except Exception as e:
            # Error handling
            error_message = f"❌ Error: {str(e)}"
            yield (
                error_message,
                error_message,
                gr.update(visible=False),
                gr.update(visible=True),
                None,
                None,
                None,
                None,
                session_id,
                None,
            )

    # Schedule periodic cleanup of old sessions
    def cleanup_task():
        SessionManager.cleanup_old_sessions()
        # Schedule the next cleanup
        threading.Timer(3600, cleanup_task).start()  # Run every hour

    # Start the cleanup task
    cleanup_task()

    # Update the click handler to include session_id
    generate_btn.click(
        fn=generate_all_backends_with_status_boxes,
        inputs=[input_text, session_id],
        outputs=[
            small_status_box,
            big_status_box,
            big_status_box,  # visibility
            small_status_box,  # visibility
            draft_output,
            quick_output,
            best_output,
            performance_plot,
            session_id,  # Update the session ID
            recent_gallery,  # Update the gallery
        ],
        api_name="generate",
        max_batch_size=10,  # Process up to 10 requests at once
        concurrency_limit=20,  # Allow up to 20 concurrent requests
        concurrency_id="generation",  # Group concurrent requests under this ID
    )

# Launch with increased max_threads
if __name__ == "__main__":
    demo.queue(max_size=50).launch(
        server_name="0.0.0.0",
        max_threads=16,  # Increase thread count for better concurrency
    )