import os
import requests
import json
import time
import threading
import uuid
import shutil
from datetime import datetime
from pathlib import Path
from http.server import HTTPServer, SimpleHTTPRequestHandler
import base64
from dotenv import load_dotenv
import gradio as gr
import random
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from functools import lru_cache

load_dotenv()

MODEL_URL = "TostAI/nsfw-text-detection-large"
CLASS_NAMES = {
    0: "✅ SAFE",
    1: "⚠️ QUESTIONABLE",
    2: "🚫 UNSAFE"
}

tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)

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

    @classmethod
    def get_session(cls, session_id):
        with cls._lock:
            if session_id not in cls._instances:
                cls._instances[session_id] = {
                    'count': 0,
                    'history': [],
                    'last_active': time.time()
                }
            return cls._instances[session_id]

    @classmethod
    def cleanup_sessions(cls):
        with cls._lock:
            now = time.time()
            expired = [k for k, v in cls._instances.items() if now - v['last_active'] > 3600]
            for k in expired:
                del cls._instances[k]

class RateLimiter:
    def __init__(self):
        self.clients = {}
        self.lock = threading.Lock()

    def check(self, client_id):
        with self.lock:
            now = time.time()
            if client_id not in self.clients:
                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
                return True
                
            if now > self.clients[client_id]['reset']:
                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
                return True
                
            if self.clients[client_id]['count'] >= 20:
                return False
                
            self.clients[client_id]['count'] += 1
            return True

session_manager = SessionManager()
rate_limiter = RateLimiter()

def image_to_base64(file_path):
    try:
        with open(file_path, "rb") as f:
            ext = Path(file_path).suffix.lower()[1:]
            mime_map = {'jpg':'jpeg','jpeg':'jpeg','png':'png','webp':'webp','gif':'gif'}
            mime = mime_map.get(ext, 'jpeg')
            
            encoded = base64.b64encode(f.read())
            if len(encoded) % 4:
                encoded += b'=' * (4 - len(encoded) % 4)
                
            return f"data:image/{mime};base64,{encoded.decode()}"
    except Exception as e:
        raise ValueError(f"Base64 Error: {str(e)}")

def create_error_image(message):
    img = Image.new("RGB", (832, 480), "#ffdddd")
    try:
        font = ImageFont.truetype("arial.ttf", 24)
    except:
        font = ImageFont.load_default()
        
    draw = ImageDraw.Draw(img)
    text = f"Error: {message[:60]}..." if len(message) > 60 else message
    draw.text((50, 200), text, fill="#ff0000", font=font)
    img.save("error.jpg")
    return "error.jpg"


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

def generate_video(
    image,
    prompt,
    duration,
    enable_safety,
    flow_shift,
    guidance,
    negative_prompt,
    steps,
    seed,
    size,
    session_id
):

    safety_level = classify_prompt(prompt)
    if safety_level != 0:
        error_img = create_error_image(CLASS_NAMES[safety_level])
        yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img
        return

    if not rate_limiter.check(session_id):
        error_img = create_error_image("Hourly limit exceeded (20 requests)")
        yield "❌ 请求过于频繁，请稍后再试", error_img
        return

    session = session_manager.get_session(session_id)
    session['last_active'] = time.time()
    session['count'] += 1

    try:
        api_key = os.getenv("WAVESPEED_API_KEY")
        if not api_key:
            raise ValueError("API key missing")
            
        base64_img = image_to_base64(image)
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        
        guidance_scale = guidance
        inference_steps = steps
        
        payload = {
            "image": base64_img,
            "enable_safety_checker": True,
            "prompt": prompt,
            "duration": duration,
            "flow_shift": flow_shift,
            "guidance_scale": guidance_scale,
            "negative_prompt": negative_prompt,
            "num_inference_steps": inference_steps,
            "seed": seed if seed != -1 else random.randint(0, 999999),
            "size": "832*480"
        }

        # 提交任务
        response = requests.post(
            "https://api.wavespeed.ai/api/v2/wavespeed-ai/wan-2.1/i2v-480p-ultra-fast",
            headers=headers,
            json=payload
        )
        
        if response.status_code != 200:
            raise Exception(f"API Error {response.status_code}: {response.text}")
            
        request_id = response.json()["data"]["id"]
        yield f"✅ 任务已提交 (ID: {request_id})", None
        
    except Exception as e:
        error_img = create_error_image(str(e))
        yield f"❌ 提交失败: {str(e)}", error_img
        return

    result_url = f"https://api.wavespeed.ai/api/v2/predictions/{request_id}/result"
    start_time = time.time()
    
    while True:
        time.sleep(1)
        try:
            resp = requests.get(result_url, headers=headers)
            if resp.status_code != 200:
                raise Exception(f"状态查询失败: {resp.text}")
                
            data = resp.json()["data"]
            status = data["status"]
            
            if status == "completed":
                elapsed = time.time() - start_time
                video_url = data["outputs"][0]
                session["history"].append(video_url)
                yield f"🎉 生成成功! 耗时 {elapsed:.1f}s", video_url
                return
                
            elif status == "failed":
                raise Exception(data.get("error", "Unknown error"))
                
            else:
                yield f"⏳ 当前状态: {status.capitalize()}...", None
                
        except Exception as e:
            error_img = create_error_image(str(e))
            yield f"❌ 生成失败: {str(e)}", error_img
            return

def cleanup_task():
    while True:
        session_manager.cleanup_sessions()
        time.sleep(3600)

with gr.Blocks(
    theme=gr.themes.Soft(),
    css="""
    .video-preview { max-width: 600px !important; }
    .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("# 🌊 Wan-2.1-i2v-480p-Ultra-Fast Run On 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("""
        The Wan2.1 14B model is an advanced image-to-video model that offers accelerated inference capabilities, enabling high-res video generation with high visual quality and motion diversity.
        """)
    
    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(type="filepath", label="Upload Image")
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Prompt...")
            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
            
            with gr.Row():
                size = gr.Dropdown(["832*480", "480*832"], value="832*480", interactive=True, label="Resolution")
                steps = gr.Slider(1, 50, value=30, label="Inference Steps")
            with gr.Row():
                duration = gr.Slider(1, 10, value=5, step=1, label="时长(秒)")
                guidance = gr.Slider(1, 20, value=7, label="Guidance Scale")
            with gr.Row():
                seed = gr.Number(-1, label="Seed")
                random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
            with gr.Row():
                enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",value=True, interactive=False)
                flow_shift = gr.Number(3, label="Flow Shift",interactive=False)
                
        with gr.Column(scale=1):
            video_output = gr.Video(label="Generated Video", format="mp4", elem_classes=["video-preview"])
            status_output = gr.Textbox(label="System Status", interactive=False, lines=4)
            generate_btn = gr.Button("Generated", variant="primary")
            
 #           with gr.Accordion("Generation History", open=False):
 #               history_gallery = gr.Gallery(label="History", columns=3)
                
            with gr.Accordion("Safety Status", open=True):
                gr.Markdown("""
                <div class="status-box safe">
                    ✅ Content safety check passed
                </div>
                """)
                
    gr.Examples(
        examples=[
            [
                "Victorian era, 19th-century gentleman wearing a black top hat and tuxedo, standing on a cobblestone street, dim gaslight lamps, passersby in vintage clothing, gentle breeze moving his coat, slow cinematic pan around him, nostalgic retro film style, realistic textures",
                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745725874603980753_95mFCAxu.jpg"
            ],
            [
                "A cyberpunk female warrior with short silver hair and glowing green eyes, wearing a futuristic armored suit, standing in a neon-lit rainy city street, camera slowly circling around her, raindrops falling in slow motion, neon reflections on wet pavement, cinematic atmosphere, highly detailed, ultra realistic, 4K",
                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745726299175719855_pFO0WSRM.jpg"
            ],
            [
                "Wide shot of a brave medieval female knight in shining silver armor and a red cape, standing on a castle rooftop at sunset, slowly drawing a large ornate sword from its scabbard, seen from a distance with the vast castle and surrounding landscape in the background, golden light bathing the scene, hair and cape flowing gently in the wind, cinematic epic atmosphere, dynamic motion, majestic clouds drifting, ultra realistic, high fantasy world, 4K ultra-detailed",
                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745727436576834405_rtsokheb.jpg"
            ],
            [
                "A girl stands in a lively 17th-century market. She holds a red tomato, looks gently into the camera and smiles briefly. Then, she glances at the tomato in her hand, slowly sets it back into the basket, turns around gracefully, and walks away with her back to the camera. The market around her is rich with colorful vegetables, meats hanging above, and bustling townsfolk. Golden-hour painterly lighting, subtle facial expressions, smooth cinematic motion, ultra-realistic detail, Vermeer-inspired style",
                "https://d2g64w682n9w0w.cloudfront.net/media/images/1745079024013078406_QT6jKNPZ.png"
            ],
            [
                "A calming video explaining diabetes management and prevention tips to reduce anxiety.",
                "https://d2g64w682n9w0w.cloudfront.net/predictions/517d518c28ef49ed9464610af48528f5/1.jpg"
            ],
            [
                "Girl dancing and spinning with friends.",
                "https://d2g64w682n9w0w.cloudfront.net/media/d45e0d4893d44712b359f3ad0b3c2795/images/1745449961409630099_KISOKGEB.jpg"
            ]
        ],
        inputs=[prompt, img_input],
        label="Example Inputs",
        examples_per_page=3
    )

    random_seed_btn.click(
        fn=lambda: random.randint(0, 999999),
        outputs=seed
    )
    
    generate_btn.click(
        generate_video,
        inputs=[
        img_input,
        prompt,
        duration,
        enable_safety,
        flow_shift,
        guidance,
        negative_prompt,
        steps,
        seed,
        size,
        session_id 
    ],
        outputs=[
            status_output, 
            video_output
        ]
    )

if __name__ == "__main__":
    threading.Thread(target=cleanup_task, daemon=True).start()
    app.queue(max_size=4).launch(
        server_name="0.0.0.0",
        max_threads=16,
        share=False
    )