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

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 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"

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 image_to_base64(file_path):
    try:
        with open(file_path, "rb") as image_file:
            ext = Path(file_path).suffix.lower().lstrip('.')
            mime_map = {
                'jpg': 'jpeg',
                'jpeg': 'jpeg',
                'png': 'png',
                'webp': 'webp',
                'gif': 'gif'
            }
            mime_type = mime_map.get(ext, 'jpeg')
            
            raw_data = image_file.read()
            encoded = base64.b64encode(raw_data)
            missing_padding = len(encoded) % 4
            if missing_padding:
                encoded += b'=' * (4 - missing_padding)
                
            return f"data:image/{mime_type};base64,{encoded.decode('utf-8')}"
    except Exception as e:
        raise ValueError(f"Base64编码失败: {str(e)}")

def generate_video(
    image,
    prompt,
    enable_safety,
    flow_shift,
    guidance_scale,
    negative_prompt,
    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("每小时限制20次请求")
        yield "❌ 请求过于频繁，请稍后再试", error_img
        return

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

    API_KEY = os.getenv("WAVESPEED_API_KEY")
    if not API_KEY:
        error_img = create_error_image("API密钥缺失")
        yield "❌ Error: Missing API Key", error_img
        return

    try:
        base64_image = image_to_base64(image)
    except Exception as e:
        error_img = create_error_image(str(e))
        yield f"❌ 文件上传失败: {str(e)}", error_img
        return

    payload = {
        "enable_safety_checker": enable_safety,
        "flow_shift": flow_shift,
        "guidance_scale": guidance_scale,
        "image": base64_image,
        "negative_prompt": negative_prompt,
        "prompt": prompt,
        "seed": seed if seed != -1 else random.randint(0, 999999),
        "size": size
    }

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
    }

    try:
        response = requests.post(
            "https://api.wavespeed.ai/api/v2/wavespeed-ai/hunyuan-custom-ref2v-480p",
            headers=headers,
            data=json.dumps(payload)
        )
        
        if response.status_code != 200:
            error_img = create_error_image(response.text)
            yield f"❌ API错误 ({response.status_code}): {response.text}", error_img
            return
            
        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(0.5)
        try:
            response = requests.get(result_url, headers=headers)
            if response.status_code != 200:
                error_img = create_error_image(response.text)
                yield f"❌ 轮询错误 ({response.status_code}): {response.text}", error_img
                return

            data = response.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}秒\n"
                       f"下载链接: {video_url}"), video_url
                return
                
            elif status == "failed":
                error_img = create_error_image(data.get('error', '未知错误'))
                yield f"❌ 任务失败: {data.get('error', '未知错误')}", error_img
                return
                
            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("# 🌊Hunyuan-Custom-Ref2v Run On [WaveSpeedAI](https://wavespeed.ai/)")
    gr.Markdown("""HunyuanCustom, a multi-modal, conditional, and controllable generation model centered on subject consistency, built upon the Hunyuan Video generation framework. It enables the generation of subject-consistent videos conditioned on text, images, audio, and video inputs.""")

    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(type="filepath", label="Input Image")
            prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Prompt...")
            negative_prompt = gr.Textbox(label="Negative Prompt", lines=2)
            size = gr.Dropdown(["832*480", "480*832"], value="832*480", label="Size")
            seed = gr.Number(-1, label="Seed")
            random_seed_btn = gr.Button("Random🎲Seed", variant="secondary")
            guidance = gr.Slider(1, 20, value=7.5, step=0.1, label="Guidance")
            flow_shift = gr.Slider(1, 20, value=13, step=1, label="Shift")
            enable_safety = gr.Checkbox(True, label="Enable Safety Checker", interactive=False)

        with gr.Column(scale=1):
            video_output = gr.Video(label="Video Output", format="mp4", interactive=False,  elem_classes=["video-preview"])
            generate_btn = gr.Button("Generate", variant="primary")
            status_output = gr.Textbox(label="status", interactive=False, lines=4)
                
            gr.Examples(
                examples=[
                    [
                        "A dog is chasing a cat in the park. ",
                        "https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_poodle.png?raw=true"
                    ],
                    [
                        "A single person, in the dressing room. A woman is holding a lipstick, trying it on, and introducing it. ",
                        "https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_boy.png?raw=true"
                    ],
                    [
                        "A man is drinking Moutai in the pavilion. ",
                        "https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_man_03.png?raw=true"
                    ],
                    [ 
                        "A woman is boxing with a panda, and they are at a stalemate. ",
                        "https://github.com/Tencent/HunyuanCustom/blob/main/assets/images/seg_woman_01.png?raw=true"
                    ]
                ],
                inputs=[prompt, img_input],
                label="Examples Prompt",
                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,
            enable_safety,
            flow_shift,
            guidance,
            negative_prompt,
            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
    )