import os
import requests
import time
import functools
import threading
import uuid
import base64
from pathlib import Path
from dotenv import load_dotenv
import gradio as gr
import random
import torch
import io
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoTokenizer, AutoModelForSequenceClassification

load_dotenv()
API_KEY = os.getenv("WAVESPEED_API_KEY")
if not API_KEY:
    raise ValueError("WAVESPEED_API_KEY is not set in environment variables")


MODEL_URL = "TostAI/nsfw-text-detection-large"
TITLE = "🖼️🔍 Image Prompt Safety Classifier 🛡️"
DESCRIPTION = "✨ Enter an image generation prompt to classify its safety level! ✨"

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)

# Define class names with emojis and detailed descriptions
CLASS_NAMES = {
    0: "✅ SAFE - This prompt is appropriate and harmless.",
    1: "⚠️ QUESTIONABLE - This prompt may require further review.",
    2: "🚫 UNSAFE - This prompt is likely to generate inappropriate content."
}


@functools.lru_cache(maxsize=128)
def classify_text(text):
    inputs = tokenizer(text,
                       return_tensors="pt",
                       truncation=True,
                       padding=True,
                       max_length=1024)

    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=1).item()

    return predicted_class, CLASS_NAMES[predicted_class]


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

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

        with cls._lock:
            if client_id not in cls._instances:
                cls._instances[client_id] = ClientGenerationManager()
            return cls._instances[client_id]

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

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


class ClientGenerationManager:

    def __init__(self):
        self.lock = threading.Lock()
        self.last_activity = time.time()
        self.request_timestamps = []  # Track timestamps of requests

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

    def add_request_timestamp(self):
        with self.lock:
            self.request_timestamps.append(time.time())

    def has_exceeded_limit(self, limit=20):
        with self.lock:
            current_time = time.time()
            # Filter timestamps to only include those within the last hour
            self.request_timestamps = [
                ts for ts in self.request_timestamps
                if current_time - ts <= 3600
            ]
            return len(self.request_timestamps) >= limit


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.last_activity = time.time()
        self.request_timestamps = []  # Track timestamps of requests

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

    def add_request_timestamp(self):
        self.request_timestamps.append(time.time())

    def has_exceeded_limit(self,
                           limit=10):  # Default limit: 10 requests per hour
        current_time = time.time()
        # Filter timestamps to only include those within the last hour
        self.request_timestamps = [
            ts for ts in self.request_timestamps if current_time - ts <= 3600
        ]
        return len(self.request_timestamps) >= limit


@torch.no_grad()
def classify_prompt(prompt):
    inputs = tokenizer(prompt,
                       return_tensors="pt",
                       truncation=True,
                       max_length=512)
    outputs = model(**inputs)
    return torch.argmax(outputs.logits).item()


def image_to_base64(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def decode_base64_to_image(base64_str):
    image_data = base64.b64decode(base64_str)
    return Image.open(io.BytesIO(image_data))


def generate_image(
    image_file,
    prompt,
    seed,
    session_id,
    enable_safety_checker,
    request: gr.Request,
):
    try:
        client_ip = request.client.host
        x_forwarded_for = request.headers.get('x-forwarded-for')
        if x_forwarded_for:
            client_ip = x_forwarded_for
        print(f"Client IP: {client_ip}")
        client_generation_manager = ClientManager.get_manager(client_ip)
        client_generation_manager.update_activity()
        if client_generation_manager.has_exceeded_limit(limit=10):
            error_message = "❌ Your network has exceeded the limit of 10 requests per hour. Please try again later."
            yield error_message, None, "", None
            return

        client_generation_manager.add_request_timestamp()
        """Generate images with big status box during generation"""
        # Get or create a session manager
        session_id, manager = SessionManager.get_manager(session_id)
        manager.update_activity()

        # # Check if the user has exceeded the request limit
        # if manager.has_exceeded_limit(
        #         limit=10):  # Set the limit to 10 requests per hour
        #     error_message = "❌ You have exceeded the limit of 10 requests per hour. Please try again later."
        #     yield error_message, None, "", None
        #     return

        # Add the current request timestamp
        manager.add_request_timestamp()

        if not prompt or prompt.strip() == "":
            # Handle empty prompt case
            error_message = "⚠️ Please enter a prompt first"
            yield error_message, None, "", None
            return

        error_messages = []
        if not image_file:
            error_messages.append("Please upload an image file")
        elif not Path(image_file).exists():
            error_messages.append("File does not exist")
        if not prompt.strip():
            error_messages.append("Prompt cannot be empty")
        if error_messages:
            error_message = "❌ Input validation failed: " + ", ".join(
                error_messages)
            yield error_message, None, "", None
            return

        # Check if the prompt is safe
        classification, message = classify_text(prompt)
        if classification == 2:  # UNSAFE
            yield "❌ NSFW prompt detected", None, "", None
            return

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

        try:
            base64_image = image_to_base64(image_file)
            input_image = decode_base64_to_image(base64_image)
        except Exception as e:
            error_message = f"❌ File processing failed: {str(e)}"
            yield error_message, None, "", None
            return

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}",
        }
        payload = {
            "enable_safety_checker": enable_safety_checker,
            "image": base64_image,
            "prompt": prompt,
            "seed": int(seed) if seed != -1 else random.randint(0, 999999)
        }

        response = requests.post(
            "https://api.wavespeed.ai/api/v3/wavespeed-ai/flux-kontext-dev-ultra-fast",
            headers=headers,
            json=payload,
            timeout=30)
        response.raise_for_status()

        request_id = response.json()["data"]["id"]
        result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
        start_time = time.time()

        for _ in range(60):
            time.sleep(1.0)
            resp = requests.get(result_url, headers=headers)
            resp.raise_for_status()

            data = resp.json()["data"]
            status = data["status"]

            if status == "completed":
                elapsed = time.time() - start_time
                output_url = data["outputs"][0]
                has_nsfw_content = data["has_nsfw_contents"][0]
                if has_nsfw_content:
                    error_message = "❌ NSFW content detected in the output"
                    yield error_message, None, "", None
                else:
                    yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", output_url, output_url, update_recent_gallery(
                        prompt, input_image, output_url)
                return
            elif status == "failed":
                raise Exception(data.get("error", "Unknown error"))
            else:
                error_message = f"⏳ Current status: {status.capitalize()}..."
                yield error_message, None, "", None

        raise Exception("Generation timed out")

    except Exception as e:
        error_message = f"❌ Generation failed: {str(e)}"
        yield error_message, None, "", None


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


# Store recent generations
recent_generations = []

with gr.Blocks(theme=gr.themes.Soft(),
               css="""
    .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
    .safe { background: #e8f5e9; border: 1px solid #a5d6a7; }
    .warning { background: #fff3e0; border: 1px solid #ffcc80; }
    .error { background: #ffebee; border: 1px solid #ef9a9a; }
    """) as app:

    session_id = gr.State(str(uuid.uuid4()))

    gr.Markdown("# 🖼️FLUX Kontext Dev Ultra Fast Live")
    gr.Markdown(
        "FLUX Kontext dev is a new SOTA image editing model published by Black Forest Labs. We have deployed it on [WaveSpeedAI](https://wavespeed.ai/) for ultra-fast image editing. You can use it to edit images in various styles, add objects, or even change the mood of the image. It supports both text prompts and image inputs."
    )
    gr.Markdown(
        "- [FLUX Kontext dev on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev)"
        "\n"
        "- [FLUX Kontext dev LoRA on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora)"
        "\n"
        "- [FLUX Kontext dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
        "\n"
        "- [FLUX Kontext dev LoRA Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-lora-ultra-fast)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_file = gr.Image(label="Upload Image",
                                  type="filepath",
                                  sources=["upload", "clipboard"],
                                  interactive=True,
                                  image_mode="RGB",
                                  value="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png")
            prompt = gr.Textbox(label="Prompt",
                                placeholder="Please enter your prompt...",
                                lines=3,
                                value="Convert the image into Claymation style.")
            seed = gr.Number(label="seed",
                             value=-1,
                             minimum=-1,
                             maximum=999999,
                             step=1)
            random_btn = gr.Button("random🎲seed", variant="secondary")
            enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",
                                        value=True,
                                        interactive=False)
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Result")
            status = gr.Textbox(label="Status", elem_classes=["status-box"])
            output_url = gr.Textbox(label="Image URL",
                                    interactive=True,
                                    visible=False)
            submit_btn = gr.Button("Start Generation", variant="primary")
    gr.Examples(
        examples=[
            [
                "Convert the image into Claymation style.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png"
            ],
            [
                "Convert the image into Ghibli style.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
            ],
            [
                "Add sunglasses to the face of the statue.",
                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
            ],
            # [
            #     'Convert the image into an ink sketch style.',
            #     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
            # ],
            # [
            #     'Add a butterfly to the scene.',
            #     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_result.png"
            # ]
        ],
        inputs=[prompt, image_file],
        label="Examples")

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

    def get_recent_gallery_items():
        gallery_items = []
        for r in reversed(recent_generations):
            if any(x is None for x in r.values()):
                continue
            gallery_items.append((r["input"], f"Input: {r['prompt']}"))
            gallery_items.append((r["output"], f"Output: {r['prompt']}"))
        return gr.update(value=gallery_items)

    def update_recent_gallery(prompt, input_image, output_image):
        recent_generations.append({
            "prompt": prompt,
            "input": input_image,
            "output": output_image,
        })
        if len(recent_generations) > 16:
            recent_generations.pop(0)
        return get_recent_gallery_items()

    random_btn.click(fn=lambda: random.randint(0, 999999), outputs=seed)

    submit_btn.click(
        generate_image,
        inputs=[image_file, prompt, seed, session_id, enable_safety],
        outputs=[status, output_image, output_url, recent_gallery],
        api_name=False,
        max_batch_size=10,
        concurrency_limit=20,
        concurrency_id="generation",
    )

if __name__ == "__main__":
    # Start the cleanup task
    cleanup_task()
    app.queue(max_size=20).launch(
        server_name="0.0.0.0",
        max_threads=10,
        share=False,
    )