import gradio as gr
import replicate
import requests
import os
import json
import asyncio
import concurrent.futures
from io import BytesIO
from PIL import Image
from typing import List, Tuple, Dict

# 환경 변수에서 토큰 가져오기
REPLICATE_API_TOKEN = os.getenv("RAPI_TOKEN")
FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN")

# 스타일 정의
STYLE_TEMPLATES = {
    "3D Style (Pixar-like)": {
        "name": "3D Style",
        "description": "Pixar-esque 3D render with volumetric lighting",
        "example": "A fluffy ginger cat wearing a tiny spacesuit, floating amidst a vibrant nebula in a 3D render. The cat is gazing curiously at a swirling planet with rings made of candy. Background is filled with sparkling stars and colorful gas clouds, lit with soft, volumetric lighting. Style: Pixar-esque, highly detailed, playful. Colors: Deep blues, purples, oranges, and pinks. Rendered in Octane, 8k resolution."
    },
    "Elegant SWOT Quadrant": {
        "name": "SWOT Analysis",
        "description": "Flat-design 4-grid layout with minimal shadows",
        "example": "Elegant SWOT quadrant: flat-design 4-grid on matte-white backdrop, thin pastel separators, top-left 'Strengths' panel shows glowing shield icon and subtle motif, top-right 'Weaknesses' panel with cracked chain icon in soft crimson, bottom-left 'Opportunities' panel with sunrise-over-horizon icon in optimistic teal, bottom-right 'Threats' panel with storm-cloud & lightning icon in deep indigo, minimal shadows, no text, no watermark, 16:9, 4K"
    },
    "Colorful Mind Map": {
        "name": "Mind Map",
        "description": "Hand-drawn educational style with vibrant colors",
        "example": "A handrawn colorful mind map diagram: educational style, vibrant colors, clear hierarchy, golden ratio layout. Central concept with branching sub-topics, each branch with unique color coding, organic flowing connections, doodle-style icons for each node"
    },
    "Business Workflow": {
        "name": "Business Process",
        "description": "End-to-end business workflow with clear phases",
        "example": "A detailed hand-drawn diagram illustrating an end-to-end business workflow with Market Analysis, Strategy Development, Product Design, Implementation, and Post-Launch Review phases. Clear directional arrows, iconography for each component, vibrant educational yet professional style"
    },
    "Industrial Design": {
        "name": "Product Design",
        "description": "Sleek industrial design concept sketch",
        "example": "A sleek industrial design concept: Curved metallic body with minimal bezel, Touchscreen panel for settings, Modern matte black finish, Hand-drawn concept sketch style with annotations and dimension lines"
    },
    "3D Bubble Chart": {
        "name": "Bubble Chart",
        "description": "Clean 3D bubble visualization",
        "example": "3-D bubble chart on clean white 2×2 grid, quadrant titles hidden, four translucent spheres in lime, azure, amber, magenta, gentle depth-of-field, modern consulting aesthetic, no text, 4K"
    },
    "Timeline Ribbon": {
        "name": "Timeline",
        "description": "Horizontal ribbon timeline with cyber-futuristic vibe",
        "example": "Horizontal ribbon timeline, milestone pins glowing hot pink on charcoal, year markers as circles, faint motion streaks, cyber-futuristic vibe, no text, 1920×1080"
    },
    "Risk Heat Map": {
        "name": "Heat Map",
        "description": "Risk assessment heat map with gradient colors",
        "example": "Risk Heat Map: square grid, smooth gradient from mint to fire-red, cells beveled, simple legend strip hidden, long subtle shadow, sterile white frame, no text"
    },
    "Pyramid/Funnel": {
        "name": "Funnel Chart",
        "description": "Multi-layer gradient funnel visualization",
        "example": "Pyramid / Funnel: 5-layer gradient funnel narrowing downwards, top vivid sky-blue, mid mint-green, bottom sunset-orange, glass reflection, minimal background, no text"
    },
    "KPI Dashboard": {
        "name": "Dashboard",
        "description": "Dark-mode analytics dashboard with sci-fi interface",
        "example": "KPI Dashboard: Dark-mode analytic dashboard, three glass speedometers glowing neon lime, two sparkline charts under, black glass background, sci-fi interface, no text, 4K"
    },
    "Value Chain": {
        "name": "Value Chain",
        "description": "Horizontal value chain with industrial look",
        "example": "Value Chain Diagram: Horizontal value chain blocks, steel-blue gradient bars with subtle bevel, small gear icons above each segment, sleek industrial look, shadow cast, no text"
    },
    "Gantt Chart": {
        "name": "Gantt Chart",
        "description": "Hand-drawn style Gantt chart with playful colors",
        "example": "Gantt Chart: Hand-drawn style Gantt bars sketched with vibrant markers on dotted grid notebook page, sticky-note color palette, playful yet organized, perspective tilt, no text"
    },
    "Mobile App Mockup": {
        "name": "App Mockup",
        "description": "Clean wireframe for mobile app design",
        "example": "MOCKUP DESIGN: A clean hand-drawn style wireframe for a mobile app with Title screen, Login screen, Dashboard with sections, Bottom navigation bar, minimalist design with annotations"
    },
    "Flowchart": {
        "name": "Flowchart",
        "description": "Vibrant flowchart with minimalistic icons",
        "example": "FLOWCHART DESIGN: A hand-drawn style flowchart, vibrant colors, minimalistic icons showing process flow from START to END with decision points, branches, and clear directional arrows"
    }
}

def generate_prompt_with_llm(topic: str, style_example: str = None) -> str:
    """주제와 스타일 예제를 받아서 LLM을 사용해 이미지 프롬프트를 생성"""
    url = "https://api.friendli.ai/dedicated/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {FRIENDLI_TOKEN}",
        "Content-Type": "application/json"
    }
    
    # 강화된 시스템 프롬프트
    system_prompt = """You are an expert image prompt engineer specializing in creating detailed, visually rich prompts for AI image generation. 

Your task is to create prompts that:
1. Are highly specific and visual, describing exact details, colors, lighting, and composition
2. Include style references (e.g., "3D render", "hand-drawn", "flat design", "industrial sketch")
3. Specify technical details like resolution, aspect ratio, or rendering style when appropriate
4. Use descriptive adjectives for materials, textures, and atmospheres
5. Avoid abstract concepts - focus on concrete visual elements

Important guidelines:
- If given a style example, adapt the topic to match that specific visual style
- Maintain the technical vocabulary and visual descriptors from the style example
- Always output ONLY the prompt without any explanation or additional text
- Make prompts between 50-150 words for optimal results"""
    
    user_message = f"Topic: {topic}"
    if style_example:
        user_message += f"\n\nStyle reference to follow:\n{style_example}\n\nAdapt the topic to match this exact visual style."
    
    payload = {
        "model": "dep89a2fld32mcm",
        "messages": [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": user_message
            }
        ],
        "max_tokens": 300,
        "top_p": 0.8,
        "temperature": 0.7,
        "stream": False
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        if response.status_code == 200:
            result = response.json()
            return result['choices'][0]['message']['content'].strip()
        else:
            return f"프롬프트 생성 실패: {response.status_code}"
    except Exception as e:
        return f"프롬프트 생성 중 오류 발생: {str(e)}"

def translate_to_english(text: str) -> str:
    """한글 텍스트를 영어로 번역 (LLM 사용)"""
    if not any(ord('가') <= ord(char) <= ord('힣') for char in text):
        return text
    
    url = "https://api.friendli.ai/dedicated/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {FRIENDLI_TOKEN}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "model": "dep89a2fld32mcm",
        "messages": [
            {
                "role": "system",
                "content": "You are a translator. Translate the given Korean text to English. Only return the translation without any explanation."
            },
            {
                "role": "user",
                "content": text
            }
        ],
        "max_tokens": 500,
        "top_p": 0.8,
        "stream": False
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        if response.status_code == 200:
            result = response.json()
            return result['choices'][0]['message']['content'].strip()
        else:
            return text
    except Exception as e:
        return text

def generate_image(prompt: str, seed: int = 10) -> Tuple[Image.Image, str]:
    """Replicate API를 사용해 이미지 생성"""
    try:
        english_prompt = translate_to_english(prompt)
        
        if not REPLICATE_API_TOKEN:
            return None, "RAPI_TOKEN 환경변수가 설정되지 않았습니다."
        
        client = replicate.Client(api_token=REPLICATE_API_TOKEN)
        
        input_params = {
            "seed": seed,
            "prompt": english_prompt,
            "speed_mode": "Extra Juiced 🚀 (even more speed)",
            "output_quality": 80
        }
        
        output = client.run(
            "prunaai/hidream-l1-fast:17c237d753218fed0ed477cb553902b6b75735f48c128537ab829096ef3d3645",
            input=input_params
        )
        
        if output:
            if isinstance(output, str) and output.startswith('http'):
                response = requests.get(output)
                img = Image.open(BytesIO(response.content))
                return img, english_prompt
            else:
                img = Image.open(BytesIO(output.read()))
                return img, english_prompt
        else:
            return None, "이미지 생성 실패"
            
    except Exception as e:
        return None, f"오류: {str(e)}"

def generate_images_parallel(topic: str, selected_styles: List[str], seed: int) -> List[Dict]:
    """선택된 스타일들에 대해 병렬로 이미지 생성"""
    results = []
    
    # 각 스타일에 대한 프롬프트 생성
    prompts = []
    for style_name in selected_styles:
        if style_name in STYLE_TEMPLATES:
            style_info = STYLE_TEMPLATES[style_name]
            prompt = generate_prompt_with_llm(topic, style_info["example"])
            prompts.append({
                "style": style_name,
                "prompt": prompt,
                "style_info": style_info
            })
    
    # 병렬로 이미지 생성
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        future_to_style = {}
        for prompt_data in prompts:
            future = executor.submit(generate_image, prompt_data["prompt"], seed)
            future_to_style[future] = prompt_data
        
        for future in concurrent.futures.as_completed(future_to_style):
            prompt_data = future_to_style[future]
            try:
                img, used_prompt = future.result()
                results.append({
                    "style": prompt_data["style"],
                    "image": img,
                    "prompt": prompt_data["prompt"],
                    "used_prompt": used_prompt,
                    "success": img is not None
                })
            except Exception as e:
                results.append({
                    "style": prompt_data["style"],
                    "image": None,
                    "prompt": prompt_data["prompt"],
                    "used_prompt": str(e),
                    "success": False
                })
    
    return results

def process_multiple_styles(topic: str, selected_styles: List[str], seed: int):
    """여러 스타일로 이미지 생성 처리"""
    if not topic.strip():
        return [], "주제를 입력해주세요."
    
    if not selected_styles:
        return [], "최소 하나의 스타일을 선택해주세요."
    
    status = f"선택된 {len(selected_styles)}개 스타일로 이미지 생성 중..."
    
    # 병렬로 이미지 생성
    results = generate_images_parallel(topic, selected_styles, seed)
    
    # 결과 정리
    images = []
    prompts_info = []
    
    for result in results:
        if result["success"]:
            images.append((result["image"], result["style"]))
            prompts_info.append(f"**{result['style']}**\n프롬프트: {result['prompt']}\n")
        else:
            prompts_info.append(f"**{result['style']}** - 생성 실패: {result['used_prompt']}\n")
    
    final_status = f"총 {len(images)}개 이미지 생성 완료\n\n" + "\n".join(prompts_info)
    
    return images, final_status

# Gradio 인터페이스 생성
with gr.Blocks(title="AI 멀티스타일 이미지 생성기", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 AI 멀티스타일 이미지 생성기
    
    주제를 입력하고 원하는 스타일을 선택하면, 각 스타일에 맞는 이미지를 동시에 생성합니다.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            topic_input = gr.Textbox(
                label="주제 입력",
                placeholder="예: 우주를 여행하는 고양이, 미래의 도시, 혁신적인 제품 디자인",
                lines=2
            )
            
            gr.Markdown("### 스타일 선택 (복수 선택 가능)")
            
            # 스타일 체크박스 그룹
            style_checkboxes = gr.CheckboxGroup(
                choices=list(STYLE_TEMPLATES.keys()),
                label="생성할 스타일",
                value=["3D Style (Pixar-like)"],
                info="각 스타일은 고유한 시각적 특성을 가지고 있습니다"
            )
            
            seed_input = gr.Slider(
                minimum=1,
                maximum=100,
                value=10,
                step=1,
                label="시드 값 (동일한 시드는 동일한 이미지 생성)"
            )
            
            generate_btn = gr.Button("🚀 선택한 스타일로 이미지 생성", variant="primary", size="lg")
            
        with gr.Column(scale=2):
            # 갤러리로 여러 이미지 표시
            output_gallery = gr.Gallery(
                label="생성된 이미지들",
                show_label=True,
                elem_id="gallery",
                columns=2,
                rows=3,
                object_fit="contain",
                height="auto"
            )
            
            status_text = gr.Markdown(
                value="생성 상태 및 프롬프트 정보가 여기에 표시됩니다."
            )
    
    # 스타일 설명 섹션
    with gr.Accordion("📚 스타일 가이드", open=False):
        style_guide_text = "### 사용 가능한 스타일:\n\n"
        for style_name, style_info in STYLE_TEMPLATES.items():
            style_guide_text += f"**{style_name}**: {style_info['description']}\n\n"
        gr.Markdown(style_guide_text)
    
    # 이미지 생성 이벤트
    generate_btn.click(
        fn=process_multiple_styles,
        inputs=[topic_input, style_checkboxes, seed_input],
        outputs=[output_gallery, status_text]
    )
    
    gr.Markdown("""
    ---
    ### 💡 사용 팁:
    - **주제**: 구체적일수록 좋은 결과를 얻을 수 있습니다
    - **스타일 조합**: 여러 스타일을 선택하면 다양한 시각적 표현을 비교할 수 있습니다
    - **병렬 처리**: 선택한 모든 스타일의 이미지가 동시에 생성됩니다
    - **시드 값**: 동일한 시드로 재현 가능한 결과를 얻을 수 있습니다
    
    각 스타일은 전문적으로 큐레이션된 예제를 기반으로 프롬프트를 생성합니다.
    """)

# 앱 실행
if __name__ == "__main__":
    # 환경 변수 확인
    if not REPLICATE_API_TOKEN:
        print("경고: RAPI_TOKEN 환경 변수가 설정되지 않았습니다.")
    if not FRIENDLI_TOKEN:
        print("경고: FRIENDLI_TOKEN 환경 변수가 설정되지 않았습니다.")
    
    demo.launch()