import os
import spaces
import gradio as gr
import numpy as np
from PIL import Image
import random
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler
import torch
from transformers import pipeline as transformers_pipeline
import re
from cohere import ClientV2

# ------------------------------------------------------------
# COHERE CLIENT
# ------------------------------------------------------------
coh_api_key = os.getenv("COH_API")
if not coh_api_key:
    print("[WARNING] COH_API environment variable not found. LLM features will not work.")
    coh_client = None
else:
    try:
        coh_client = ClientV2(api_key=coh_api_key)
        print("[INFO] Cohere client initialized successfully.")
    except Exception as e:
        print(f"[ERROR] Failed to initialize Cohere client: {str(e)}")
        coh_client = None


# ------------------------------------------------------------
# MODEL CONFIGURATION
# ------------------------------------------------------------
MODEL_ID = "Heartsync/NSFW-Uncensored"
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1216


# ------------------------------------------------------------
# TRANSLATION FUNCTIONS (이전과 동일)
# ------------------------------------------------------------
non_english_regex = re.compile(r'[\uac00-\ud7a3\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]+')

def is_non_english(text):
    """Check if text contains non-English characters"""
    if re.search(r'[\uac00-\ud7a3]', text):
        print("[DETECT] Korean text detected")
        return True
    if re.search(r'[\u3040-\u30ff]', text):
        print("[DETECT] Japanese text detected")
        return True
    if re.search(r'[\u4e00-\u9fff]', text):
        print("[DETECT] Chinese/Kanji text detected")
        return True
    if re.search(r'[^\x00-\x7F]', text):
        print("[DETECT] Other non-English text detected")
        return True
    return False


def translate_with_cohere(text):
    """Translate non-English text to English"""
    if coh_client is None:
        print("[WARN] Cohere client not available, skipping translation")
        return text
    
    if not is_non_english(text):
        print("[INFO] English text detected, no translation needed")
        return text
    
    try:
        print(f"[INFO] Translating text: '{text}'")
        
        system_prompt = """You are a professional translator. Translate the input text to English accurately.
Provide ONLY the translated English text with no explanations."""
        
        messages = [
            {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
            {"role": "user", "content": [{"type": "text", "text": text}]}
        ]

        response = coh_client.chat(
            model="command-r-plus-08-2024",
            messages=messages,
            temperature=0.1
        )
        
        translated_text = response.text.strip() if hasattr(response, 'text') else str(response)
        translated_text = re.sub(r'^(Translation:|English:|Translated text:)\s*', '', translated_text, flags=re.IGNORECASE)
        
        print(f"[INFO] Original: '{text}'")
        print(f"[INFO] Translated: '{translated_text}'")
        
        return translated_text if len(translated_text) > 3 else text
            
    except Exception as e:
        print(f"[ERROR] Translation failed: {str(e)}")
        return text


def translate_prompt_if_needed(prompt):
    """Helper function to translate prompt"""
    if not is_non_english(prompt):
        return prompt
    
    if coh_client is None:
        return prompt
    
    try:
        trans_system = "Translate to English accurately. Only provide the translation."
        
        trans_response = coh_client.chat(
            model="command-r-plus-08-2024",
            messages=[
                {"role": "system", "content": [{"type": "text", "text": trans_system}]},
                {"role": "user", "content": [{"type": "text", "text": prompt}]}
            ],
            temperature=0.1
        )
        
        if hasattr(trans_response, 'text'):
            translated_prompt = trans_response.text.strip()
            print(f"[SUCCESS] Translated: '{prompt}' -> '{translated_prompt}'")
            return translated_prompt
            
    except Exception as e:
        print(f"[ERROR] Translation failed: {str(e)}")
    
    return prompt


# ------------------------------------------------------------
# EXAMPLES
# ------------------------------------------------------------
prompt_examples = [
    "The shy college girl, with glasses and a tight plaid skirt, nervously approaches her professor",
    "Her skirt rose a little higher with each gentle push, a soft blush of blush spreading across her cheeks",
    "Moody mature anime scene of two lovers under neon rain, sensual atmosphere",
    "The girl sits on the boy's lap by the window, his hands resting on her waist",
    "A woman in a business suit, elegant and confident pose",
    "Artistic portrait with dramatic lighting and soft shadows",
]


# ------------------------------------------------------------
# LLM PROMPT GENERATOR
# ------------------------------------------------------------
def generate_prompts(theme):
    """Generate optimal prompts using LLM"""
    try:
        if coh_client is None:
            return "Cohere API token not set."
        
        if non_english_regex.search(theme):
            theme = translate_with_cohere(theme)
        
        print(f"[INFO] Generating prompt for theme: {theme}")
        
        system_prefix = """You are an expert at creating detailed image generation prompts. Create ONE optimal prompt based on the theme.

Guidelines:
1. Generate only ONE high-quality prompt
2. 1-3 sentences long
3. Detailed and descriptive
4. ONLY respond in ENGLISH
5. NO prefixes or headers - just the prompt text"""

        messages = [
            {"role": "system", "content": [{"type": "text", "text": system_prefix}]},
            {"role": "user", "content": [{"type": "text", "text": theme}]}
        ]

        response = coh_client.chat(
            model="command-r-plus-08-2024",
            messages=messages,
            temperature=0.8
        )
        
        generated_prompt = response.text if hasattr(response, 'text') else str(response)
        
        if non_english_regex.search(generated_prompt):
            generated_prompt = translate_with_cohere(generated_prompt)
        
        generated_prompt = re.sub(r'^(AI🐼|Prompt|Response|Result|Output):\s*', '', generated_prompt)
        generated_prompt = re.sub(r'^["\']+|["\']+$', '', generated_prompt)
        generated_prompt = generated_prompt.strip()
        
        print(f"[INFO] Generated: {generated_prompt}")
        
        return generated_prompt if len(generated_prompt) > 10 else "Failed to generate prompt"
        
    except Exception as e:
        print(f"[ERROR] Prompt generation failed: {str(e)}")
        return f"Error: {str(e)}"


# ------------------------------------------------------------
# SDXL TEXT-TO-IMAGE (ZeroGPU 최적화)
# ------------------------------------------------------------
@spaces.GPU(duration=120)
def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    """Text-to-Image generation optimized for ZeroGPU"""
    
    print(f"[DEBUG] Original prompt: '{prompt}'")
    
    # Translate if needed
    if is_non_english(prompt):
        print(f"[ALERT] Non-English prompt detected")
        prompt = translate_prompt_if_needed(prompt)
        print(f"[INFO] Translated prompt: '{prompt}'")
    
    if is_non_english(negative_prompt):
        negative_prompt = translate_prompt_if_needed(negative_prompt)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    try:
        # GPU 컨텍스트 내에서 파이프라인 로드
        device = torch.device("cuda")
        print(f"[INFO] Loading pipeline on device: {device}")
        
        # FP16 파이프라인 로드
        pipe = StableDiffusionXLPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True,
        )
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe = pipe.to(device)
        
        # 메모리 효율성 개선
        try:
            pipe.enable_xformers_memory_efficient_attention()
            print("[INFO] xformers memory efficient attention enabled")
        except:
            print("[WARN] xformers not available, using default attention")
        
        # VAE 타일링 활성화 (메모리 절약)
        pipe.enable_vae_tiling()
        
        # 생성기 설정
        generator = torch.Generator(device=device).manual_seed(seed)
        
        # 이미지 생성
        print(f"[INFO] Generating image...")
        output_image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        
        # 메모리 정리
        del pipe
        torch.cuda.empty_cache()
        
        print("[SUCCESS] Image generated successfully")
        return output_image, seed
        
    except Exception as e:
        print(f"[ERROR] Generation failed: {str(e)}")
        import traceback
        traceback.print_exc()
        
        # 에러 발생 시 메모리 정리
        torch.cuda.empty_cache()
        
        return Image.new("RGB", (width, height), color=(0, 0, 0)), seed


# ------------------------------------------------------------
# SDXL IMAGE-TO-IMAGE (ZeroGPU 최적화)
# ------------------------------------------------------------
@spaces.GPU(duration=120)
def img2img_infer(init_image, prompt, negative_prompt, strength, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
    """Image-to-Image generation optimized for ZeroGPU"""
    
    if init_image is None:
        return None, seed
    
    print(f"[DEBUG] Image-to-Image prompt: '{prompt}'")
    
    # Translate if needed
    if is_non_english(prompt):
        prompt = translate_prompt_if_needed(prompt)
    
    if is_non_english(negative_prompt):
        negative_prompt = translate_prompt_if_needed(negative_prompt)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    try:
        # GPU 컨텍스트 내에서 파이프라인 로드
        device = torch.device("cuda")
        print(f"[INFO] Loading img2img pipeline on device: {device}")
        
        img2img_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True,
        )
        img2img_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(img2img_pipe.scheduler.config)
        img2img_pipe = img2img_pipe.to(device)
        
        # 메모리 효율성 개선
        try:
            img2img_pipe.enable_xformers_memory_efficient_attention()
        except:
            pass
        
        img2img_pipe.enable_vae_tiling()
        
        # 이미지 전처리
        init_image = init_image.convert("RGB")
        init_image = init_image.resize((width, height), Image.Resampling.LANCZOS)
        
        generator = torch.Generator(device=device).manual_seed(seed)
        
        # 이미지 생성
        print(f"[INFO] Transforming image...")
        output_image = img2img_pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=init_image,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
        ).images[0]
        
        # 메모리 정리
        del img2img_pipe
        torch.cuda.empty_cache()
        
        print("[SUCCESS] Image transformed successfully")
        return output_image, seed
        
    except Exception as e:
        print(f"[ERROR] Transformation failed: {str(e)}")
        import traceback
        traceback.print_exc()
        
        torch.cuda.empty_cache()
        
        return None, seed


# ------------------------------------------------------------
# HELPER FUNCTIONS
# ------------------------------------------------------------
def get_random_prompt():
    return random.choice(prompt_examples)


def boost_prompt(keyword):
    if not keyword or keyword.strip() == "":
        return "Please enter a keyword or theme first"
    
    if coh_client is None:
        return "Cohere API token not set"
    
    prompt = generate_prompts(keyword)
    return prompt.strip() if len(prompt) > 10 else "Failed to generate prompt"


# ------------------------------------------------------------
# UI LAYOUT
# ------------------------------------------------------------
css = """
body {background: linear-gradient(135deg, #f2e6ff 0%, #e6f0ff 100%); color: #222; font-family: 'Noto Sans', sans-serif;}
#col-container {margin: 0 auto; max-width: 768px; padding: 15px; background: rgba(255, 255, 255, 0.8); border-radius: 15px; box-shadow: 0 8px 32px rgba(31, 38, 135, 0.2);}
.gr-button {background: #7fbdf6; color: #fff; border-radius: 8px; transition: all 0.3s ease; font-weight: bold;}
.gr-button:hover {background: #5a9ae6; transform: translateY(-2px); box-shadow: 0 5px 15px rgba(0,0,0,0.1);}
#prompt-box textarea {font-size: 1.1rem; height: 9rem !important; background: #fff; color: #222; border-radius: 10px;}
.boost-btn {background: #ff7eb6; margin-top: 5px;}
.boost-btn:hover {background: #ff5aa5;}
.random-btn {background: #9966ff; margin-top: 5px;}
.random-btn:hover {background: #8040ff;}
.title {color: #6600cc; text-shadow: 1px 1px 2px rgba(0,0,0,0.1);}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        ## 🖌️ NSFW Uncensored Text & Imagery: AI Limits Explorer  
        
        **ZeroGPU Optimized | Multi-language Support**
        """, elem_classes=["title"]
    )
    
    with gr.Tabs():
        # Text-to-Image Tab
        with gr.TabItem("Text to Image"):
            with gr.Column(elem_id="col-container"):
                with gr.Row():
                    keyword_input = gr.Text(
                        label="Keyword Input",
                        placeholder="Enter keyword in any language",
                        value="random",
                    )
                    boost_button = gr.Button("BOOST", elem_classes=["boost-btn"])
                    random_button = gr.Button("RANDOM", elem_classes=["random-btn"])
                
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        elem_id="prompt-box",
                        max_lines=3,
                        placeholder="Enter prompt in any language",
                    )
                    run_button = gr.Button("Generate", scale=0)

                result = gr.Image(label="Generated Image")

                with gr.Accordion("Advanced Settings", open=False):
                    negative_prompt = gr.Text(
                        label="Negative prompt",
                        value="low quality, watermark, signature",
                    )
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    with gr.Row():
                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                    with gr.Row():
                        guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=20.0, step=0.1, value=7)
                        num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)

        # Image-to-Image Tab
        with gr.TabItem("Image to Image"):
            with gr.Column(elem_id="col-container"):
                input_image = gr.Image(label="Input Image", type="pil")
                
                with gr.Row():
                    img2img_prompt = gr.Text(
                        label="Prompt",
                        placeholder="Describe transformation (any language)",
                    )
                    img2img_run_button = gr.Button("Transform", scale=0)
                
                img2img_result = gr.Image(label="Transformed Image")
                
                with gr.Accordion("Advanced Settings", open=False):
                    img2img_negative_prompt = gr.Text(
                        label="Negative prompt",
                        value="low quality, watermark",
                    )
                    strength = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.75)
                    img2img_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    img2img_randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    with gr.Row():
                        img2img_width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                        img2img_height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                    with gr.Row():
                        img2img_guidance_scale = gr.Slider(label="Guidance", minimum=0.0, maximum=20.0, step=0.1, value=7.5)
                        img2img_num_inference_steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=30)

    # Event handlers
    boost_button.click(fn=boost_prompt, inputs=[keyword_input], outputs=[prompt])
    random_button.click(fn=get_random_prompt, outputs=[prompt])
    
    run_button.click(
        fn=infer,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed]
    )
    
    img2img_run_button.click(
        fn=img2img_infer,
        inputs=[input_image, img2img_prompt, img2img_negative_prompt, strength, img2img_seed, 
                img2img_randomize_seed, img2img_width, img2img_height, img2img_guidance_scale, img2img_num_inference_steps],
        outputs=[img2img_result, img2img_seed]
    )

demo.queue().launch()