# ===== CRITICAL: Import spaces FIRST before any CUDA operations =====
try:
    import spaces
    HF_SPACES = True
except ImportError:
    # If running locally, create a dummy decorator
    def spaces_gpu_decorator(duration=60):
        def decorator(func):
            return func
        return decorator
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()
    HF_SPACES = False
    print("Warning: Running without Hugging Face Spaces GPU allocation")

# ===== Now import other libraries =====
import random
import os
import uuid
import re
import time
from datetime import datetime

import gradio as gr
import numpy as np
import requests
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# ===== OpenAI 설정 =====
from openai import OpenAI

# Add error handling for API key
try:
    client = OpenAI(api_key=os.getenv("LLM_API"))
except Exception as e:
    print(f"Warning: OpenAI client initialization failed: {e}")
    client = None

# ===== 프롬프트 증강용 스타일 프리셋 =====
STYLE_PRESETS = {
    "None": "",
    "Realistic Photo": "photorealistic, 8k, ultra-detailed, cinematic lighting, realistic skin texture",
    "Oil Painting": "oil painting, rich brush strokes, canvas texture, baroque lighting",
    "Comic Book": "comic book style, bold ink outlines, cel shading, vibrant colors",
    "Watercolor": "watercolor illustration, soft gradients, splatter effect, pastel palette",
}

# ===== 저장 폴더 =====
SAVE_DIR = "saved_images"
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR, exist_ok=True)

# ===== 디바이스 & 모델 로드 =====
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

repo_id = "black-forest-labs/FLUX.1-dev"
adapter_id = "seawolf2357/chocs"

# Add error handling for model loading
try:
    pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
    pipeline.load_lora_weights(adapter_id)
    pipeline = pipeline.to(device)
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    pipeline = None

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

# ===== 한글 여부 판별 =====
HANGUL_RE = re.compile(r"[\u3131-\u318E\uAC00-\uD7A3]+")

def is_korean(text: str) -> bool:
    return bool(HANGUL_RE.search(text))

# ===== 번역 & 증강 함수 =====

def openai_translate(text: str, retries: int = 3) -> str:
    """한글을 영어로 번역 (OpenAI GPT-4o-mini 사용). 영어 입력이면 그대로 반환."""
    if not is_korean(text):
        return text
    
    if client is None:
        print("Warning: OpenAI client not available, returning original text")
        return text

    for attempt in range(retries):
        try:
            res = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": "Translate the following Korean prompt into concise, descriptive English suitable for an image generation model. Keep the meaning, do not add new concepts."
                    },
                    {"role": "user", "content": text}
                ],
                temperature=0.3,
                max_tokens=256,
            )
            return res.choices[0].message.content.strip()
        except Exception as e:
            print(f"[translate] attempt {attempt + 1} failed: {e}")
            time.sleep(2)
    return text  # 번역 실패 시 원문 그대로

def enhance_prompt(text: str, retries: int = 3) -> str:
    """OpenAI를 통해 프롬프트를 증강하여 고품질 이미지 생성을 위한 상세한 설명으로 변환."""
    if client is None:
        print("Warning: OpenAI client not available, returning original text")
        return text

    for attempt in range(retries):
        try:
            res = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": """You are an expert prompt engineer for image generation models. Enhance the given prompt to create high-quality, detailed images.

Guidelines:
- Add specific visual details (lighting, composition, colors, textures)
- Include technical photography terms (depth of field, focal length, etc.)
- Add atmosphere and mood descriptors
- Specify image quality terms (4K, ultra-detailed, professional, etc.)
- Keep the core subject and meaning intact
- Make it comprehensive but not overly long
- Focus on visual elements that will improve image generation quality

Example:
Input: "A man giving a speech"
Output: "A professional man giving an inspiring speech at a podium, dramatic lighting with warm spotlights, confident posture and gestures, high-resolution 4K photography, sharp focus, cinematic composition, bokeh background with audience silhouettes, professional event setting, detailed facial expressions, realistic skin texture"
"""
                    },
                    {"role": "user", "content": f"Enhance this prompt for high-quality image generation: {text}"}
                ],
                temperature=0.7,
                max_tokens=512,
            )
            return res.choices[0].message.content.strip()
        except Exception as e:
            print(f"[enhance] attempt {attempt + 1} failed: {e}")
            time.sleep(2)
    return text  # 증강 실패 시 원문 그대로

def prepare_prompt(user_prompt: str, style_key: str, enhance_prompt_enabled: bool = False) -> str:
    """한글이면 번역하고, 프롬프트 증강 옵션이 활성화되면 증강하고, 선택한 스타일 프리셋을 붙여서 최종 프롬프트를 만든다."""
    # 1. 번역 (한글인 경우)
    prompt_en = openai_translate(user_prompt)
    
    # 2. 프롬프트 증강 (활성화된 경우)
    if enhance_prompt_enabled:
        prompt_en = enhance_prompt(prompt_en)
        print(f"Enhanced prompt: {prompt_en}")
    
    # 3. 스타일 프리셋 적용
    style_suffix = STYLE_PRESETS.get(style_key, "")
    if style_suffix:
        final_prompt = f"{prompt_en}, {style_suffix}"
    else:
        final_prompt = prompt_en
    
    return final_prompt

# ===== 이미지 저장 =====

def save_generated_image(image: Image.Image, prompt: str) -> str:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    unique_id = str(uuid.uuid4())[:8]
    filename = f"{timestamp}_{unique_id}.png"
    filepath = os.path.join(SAVE_DIR, filename)
    image.save(filepath)

    # 메타데이터 저장
    metadata_file = os.path.join(SAVE_DIR, "metadata.txt")
    with open(metadata_file, "a", encoding="utf-8") as f:
        f.write(f"{filename}|{prompt}|{timestamp}\n")
    return filepath

# ===== Diffusion 호출 =====

def run_pipeline(prompt: str, seed: int, width: int, height: int, guidance_scale: float, num_steps: int, lora_scale: float):
    if pipeline is None:
        raise ValueError("Model pipeline not loaded")
    
    generator = torch.Generator(device=device).manual_seed(int(seed))
    result = pipeline(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_steps,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]
    return result

# ===== Gradio inference 래퍼 =====

@spaces.GPU(duration=60)
def generate_image(
    user_prompt: str,
    style_key: str,
    enhance_prompt_enabled: bool = False,
    seed: int = 42,
    randomize_seed: bool = True,
    width: int = 1024,
    height: int = 768,
    guidance_scale: float = 3.5,
    num_inference_steps: int = 30,
    lora_scale: float = 1.0,
    progress=None,
):
    try:
        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        # 1) 번역 + 증강
        final_prompt = prepare_prompt(user_prompt, style_key, enhance_prompt_enabled)
        print(f"Final prompt: {final_prompt}")

        # 2) 파이프라인 호출
        image = run_pipeline(final_prompt, seed, width, height, guidance_scale, num_inference_steps, lora_scale)

        # 3) 저장
        save_generated_image(image, final_prompt)

        return image, seed
    
    except Exception as e:
        print(f"Error generating image: {e}")
        # Return a placeholder or error message
        error_image = Image.new('RGB', (width, height), color='red')
        return error_image, seed

# ===== 예시 프롬프트 (한국어/영어 혼용 허용) =====

examples = [
    "Mr. cho 두 손으로 'Healing !' 현수막을 들고 있는 모습, 환경보호와 지속가능한 임업 발전에 대한 의지를 보여주고 있다.",
    "Mr. cho 양팔을 들어 올리며 기쁜 표정으로 환호하는 모습, 조림 사업 성공과 미래 임업에 대한 희망을 보여주고 있다.",
    "Mr. cho 운동복을 입고 산림 속에서 트레킹하는 모습, 건강한 생활습관과 활기찬 리더십을 보여주고 있다.",
    "Mr. cho 산촌 마을에서 여성 임업인들과 따뜻하게 악수하는 모습, 여성 임업종사자들에 대한 진정한 관심과 소통을 보여주고 있다.",
    "Mr. cho 임업박람회장에서 울창한 숲을 향해 손가락으로 가리키며 영감을 주는 제스처를 취하고 있고, 여성들과 아이들이 박수를 치고 있다.",
    "Mr. cho 산림축제에 참여하여 열정적으로 응원하는 여성 임업인들에게 둘러싸여 있는 모습.",
    "Mr. cho 목재시장을 방문하여 여성 목재상들과 목공예 장인들과 친근하게 대화하는 모습.",
    "Mr. cho 산림과학원을 둘러보며 여성 연구원들과 교수들과 함께 임업 정책에 대해 토론하는 모습.",
    "Mr. cho 대규모 임업인 대회에서 자신감 있는 제스처와 결연한 표정으로 역동적인 연설을 하는 모습.",
    "Mr. cho 활기찬 인터뷰 현장에서 미래 임업 발전에 대한 비전을 열정적으로 설명하는 모습.",
    "Mr. cho 중요한 임업정책 회의를 준비하며 서류들에 둘러싸여 집중하고 단호한 모습을 보이는 모습."
]

# ===== 커스텀 CSS (진한 붉은색 고급 디자인) =====
custom_css = """
:root {
    --color-primary: #E91E63;
    --color-secondary: #FCE4EC;
    --color-accent: #F8BBD9;
    --color-rose: #F06292;
    --color-gold: #FFB74D;
    --color-warm-gray: #F5F5F5;
    --color-dark-gray: #424242;
    --background-primary: linear-gradient(135deg, #FAFAFA 0%, #F5F5F5 50%, #EEEEEE 100%);
    --background-accent: linear-gradient(135deg, #FCE4EC 0%, #F8BBD9 100%);
    --text-primary: #212121;
    --text-secondary: #757575;
    --shadow-soft: 0 4px 20px rgba(0, 0, 0, 0.08);
    --shadow-medium: 0 8px 30px rgba(0, 0, 0, 0.12);
    --border-radius: 16px;
}

/* 전체 배경 */
footer {visibility: hidden;}
.gradio-container {
    background: var(--background-primary) !important;
    min-height: 100vh;
    font-family: 'Inter', 'Noto Sans KR', sans-serif;
}

/* 타이틀 스타일 */
.title {
    color: var(--text-primary) !important;
    font-size: 3rem !important;
    font-weight: 700 !important;
    text-align: center;
    margin: 2rem 0;
    background: linear-gradient(135deg, var(--color-primary) 0%, var(--color-rose) 50%, var(--color-gold) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    letter-spacing: -0.02em;
}

.subtitle {
    color: var(--text-secondary) !important;
    font-size: 1.2rem !important;
    text-align: center;
    margin-bottom: 2rem;
    font-weight: 400;
}

.collection-link {
    text-align: center;
    margin-bottom: 2rem;
    font-size: 1rem;
}

.collection-link a {
    color: var(--color-primary);
    text-decoration: none;
    transition: all 0.3s ease;
    font-weight: 500;
    border-bottom: 1px solid transparent;
}

.collection-link a:hover {
    color: var(--color-rose);
    border-bottom-color: var(--color-rose);
}

/* 심플한 카드 스타일 */
.model-description {
    background: rgba(255, 255, 255, 0.9);
    border: 1px solid rgba(233, 30, 99, 0.1);
    border-radius: var(--border-radius);
    padding: 2rem;
    margin: 1.5rem 0;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(10px);
    -webkit-backdrop-filter: blur(10px);
}

.model-description p {
    color: var(--text-primary) !important;
    font-size: 1rem;
    line-height: 1.6;
    margin: 0;
}

/* 버튼 스타일 */
button.primary {
    background: var(--background-accent) !important;
    color: var(--color-primary) !important;
    border: 1px solid var(--color-accent) !important;
    border-radius: 12px !important;
    box-shadow: var(--shadow-soft) !important;
    transition: all 0.2s ease !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
}

button.primary:hover {
    background: linear-gradient(135deg, var(--color-accent) 0%, var(--color-secondary) 100%) !important;
    transform: translateY(-1px) !important;
    box-shadow: var(--shadow-medium) !important;
}

/* 입력 컨테이너 */
.input-container {
    background: rgba(255, 255, 255, 0.8);
    border: 1px solid rgba(233, 30, 99, 0.15);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-bottom: 1.5rem;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(10px);
    -webkit-backdrop-filter: blur(10px);
}

/* 고급 설정 */
.advanced-settings {
    background: rgba(255, 255, 255, 0.6);
    border: 1px solid rgba(233, 30, 99, 0.1);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
    backdrop-filter: blur(8px);
    -webkit-backdrop-filter: blur(8px);
}

/* 예시 영역 */
.example-region {
    background: rgba(252, 228, 236, 0.3);
    border: 1px solid rgba(233, 30, 99, 0.15);
    border-radius: var(--border-radius);
    padding: 1.5rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 프롬프트 입력칸 스타일 */
.large-prompt textarea {
    min-height: 120px !important;
    font-size: 15px !important;
    line-height: 1.5 !important;
    background: rgba(255, 255, 255, 0.9) !important;
    border: 2px solid rgba(233, 30, 99, 0.2) !important;
    border-radius: 12px !important;
    color: var(--text-primary) !important;
    transition: all 0.3s ease !important;
    padding: 1rem !important;
}

.large-prompt textarea:focus {
    border-color: var(--color-primary) !important;
    box-shadow: 0 0 0 3px rgba(233, 30, 99, 0.1) !important;
    outline: none !important;
}

.large-prompt textarea::placeholder {
    color: var(--text-secondary) !important;
    font-style: italic;
}

/* 생성 버튼 */
.small-generate-btn {
    max-width: 140px !important;
    height: 48px !important;
    font-size: 15px !important;
    padding: 12px 24px !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
}

/* 프롬프트 증강 섹션 */
.prompt-enhance-section {
    background: linear-gradient(135deg, rgba(255, 183, 77, 0.1) 0%, rgba(252, 228, 236, 0.2) 100%);
    border: 1px solid rgba(255, 183, 77, 0.3);
    border-radius: var(--border-radius);
    padding: 1.2rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 스타일 프리셋 섹션 */
.style-preset-section {
    background: linear-gradient(135deg, rgba(248, 187, 217, 0.15) 0%, rgba(252, 228, 236, 0.2) 100%);
    border: 1px solid rgba(233, 30, 99, 0.2);
    border-radius: var(--border-radius);
    padding: 1.2rem;
    margin-top: 1rem;
    box-shadow: var(--shadow-soft);
}

/* 라벨 텍스트 */
label {
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
}

/* 정보 텍스트 */
.gr-info, .gr-textbox-info {
    color: var(--text-secondary) !important;
    font-size: 0.85rem !important;
    line-height: 1.4 !important;
}

/* 예시 마크다운 */
.example-region h3 {
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    margin-bottom: 1rem !important;
}

/* 폼 요소들 */
input[type="radio"], input[type="checkbox"] {
    accent-color: var(--color-primary) !important;
}

input[type="range"] {
    accent-color: var(--color-primary) !important;
}

/* 결과 이미지 컨테이너 */
.image-container {
    border-radius: var(--border-radius) !important;
    overflow: hidden !important;
    box-shadow: var(--shadow-medium) !important;
    background: rgba(255, 255, 255, 0.9) !important;
    border: 1px solid rgba(233, 30, 99, 0.1) !important;
}

/* 슬라이더 컨테이너 스타일링 */
.gr-slider {
    margin: 0.5rem 0 !important;
}

/* 아코디언 스타일 */
.gr-accordion {
    border: 1px solid rgba(233, 30, 99, 0.15) !important;
    border-radius: var(--border-radius) !important;
    background: rgba(255, 255, 255, 0.7) !important;
}

.gr-accordion-header {
    background: var(--background-accent) !important;
    color: var(--color-primary) !important;
    font-weight: 600 !important;
    border-radius: var(--border-radius) var(--border-radius) 0 0 !important;
}

/* 부드러운 애니메이션 */
.model-description, .input-container, .prompt-enhance-section, .style-preset-section {
    animation: fadeInUp 0.4s ease-out;
}

@keyframes fadeInUp {
    from {
        opacity: 0;
        transform: translateY(20px);
    }
    to {
        opacity: 1;
        transform: translateY(0);
    }
}

/* 전체적인 텍스트 가독성 향상 */
* {
    -webkit-font-smoothing: antialiased;
    -moz-osx-font-smoothing: grayscale;
}

/* 드롭다운 및 셀렉트 스타일 */
select, .gr-dropdown {
    background: rgba(255, 255, 255, 0.9) !important;
    border: 1px solid rgba(233, 30, 99, 0.2) !important;
    border-radius: 8px !important;
    color: var(--text-primary) !important;
}

/* 체크박스와 라디오 버튼 개선 */
.gr-checkbox, .gr-radio {
    background: transparent !important;
}

/* 전체 컨테이너 여백 조정 */
.gr-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 2rem 1rem !important;
}

/* 모바일 반응형 */
@media (max-width: 768px) {
    .title {
        font-size: 2.2rem !important;
    }
    
    .model-description, .input-container, .advanced-settings, .example-region {
        padding: 1rem !important;
        margin: 1rem 0 !important;
    }
    
    .large-prompt textarea {
        min-height: 100px !important;
        font-size: 14px !important;
    }
}
"""

# ===== Gradio UI =====
def create_interface():
    with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
        with gr.Group(elem_classes="model-description"):
            gr.HTML("""
            <p>
            <strong>Mr. CHO CS</strong><br>
            <small style="opacity: 0.8;">본 모델은 연구 목적으로 특정인의 얼굴과 외모를 LoRA 기술로 학습한 모델입니다.목적 외의 용도로 무단 사용하지 않도록 유의해 주세요. 프롬프트에 'cho'을 포함하여 주세요.</small><br><br>
            """)

        # ===== 메인 입력 =====
        with gr.Column():
            with gr.Row(elem_classes="input-container"):
                with gr.Column(scale=4):
                    user_prompt = gr.Text(
                        label="Prompt (프롬프트)", 
                        max_lines=5, 
                        value=examples[0],
                        elem_classes="large-prompt",
                        placeholder="Enter your image description here... (이미지 설명을 입력하세요...)"
                    )
                with gr.Column(scale=1):
                    run_button = gr.Button(
                        "Generate (생성)", 
                        variant="primary",
                        elem_classes="small-generate-btn"
                    )
            
            # 프롬프트 증강 옵션 (생성 버튼 아래)
            with gr.Group(elem_classes="prompt-enhance-section"):
                enhance_prompt_checkbox = gr.Checkbox(
                    label="🚀 Prompt Enhancement (프롬프트 증강)", 
                    value=False,
                    info="Automatically improve your prompt using OpenAI API for high-quality image generation (OpenAI API를 사용하여 고품질 이미지 생성을 위해 프롬프트를 자동으로 개선합니다)"
                )
            
            # 스타일 프리셋 섹션
            with gr.Group(elem_classes="style-preset-section"):
                style_select = gr.Radio(
                    label="🎨 Style Preset (스타일 프리셋)", 
                    choices=list(STYLE_PRESETS.keys()), 
                    value="None", 
                    interactive=True
                )

            result_image = gr.Image(label="Generated Image (생성된 이미지)")
            seed_output = gr.Number(label="Seed (시드값)")

            # ===== 고급 설정 =====
            with gr.Accordion("Advanced Settings (고급 설정)", open=False, elem_classes="advanced-settings"):
                seed = gr.Slider(label="Seed (시드값)", minimum=0, maximum=MAX_SEED, step=1, value=42)
                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=768)
                with gr.Row():
                    guidance_scale = gr.Slider(label="Guidance scale (가이던스 스케일)", minimum=0.0, maximum=10.0, step=0.1, value=3.5)
                    num_inference_steps = gr.Slider(label="Inference steps (추론 단계)", minimum=1, maximum=50, step=1, value=30)
                    lora_scale = gr.Slider(label="LoRA scale (LoRA 스케일)", minimum=0.0, maximum=1.0, step=0.1, value=1.0)

            # ===== 예시 영역 =====
            with gr.Group(elem_classes="example-region"):
                gr.Markdown("### Examples (예시)")
                gr.Examples(examples=examples, inputs=user_prompt, cache_examples=False)

        # ===== 이벤트 =====
        run_button.click(
            fn=generate_image,
            inputs=[
                user_prompt,
                style_select,
                enhance_prompt_checkbox,
                seed,
                randomize_seed,
                width,
                height,
                guidance_scale,
                num_inference_steps,
                lora_scale,
            ],
            outputs=[result_image, seed_output],
        )
    
    return demo

# ===== 애플리케이션 실행 =====
if __name__ == "__main__":
    demo = create_interface()
    demo.queue()
    demo.launch()