app.py

import gradio as gr
import os
import json
import requests
from datetime import datetime
import time
from typing import List, Dict, Any, Generator, Tuple, Optional, Set
import logging
import re
import tempfile
from pathlib import Path
import sqlite3
import hashlib
import threading
from contextlib import contextmanager
from dataclasses import dataclass, field, asdict
from collections import defaultdict
import random
from huggingface_hub import HfApi, upload_file, hf_hub_download

# --- Logging setup ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Document export imports ---
try:
    from docx import Document
    from docx.shared import Inches, Pt, RGBColor, Mm
    from docx.enum.text import WD_ALIGN_PARAGRAPH
    from docx.enum.style import WD_STYLE_TYPE
    from docx.oxml.ns import qn
    from docx.oxml import OxmlElement
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False
    logger.warning("python-docx not installed. DOCX export will be disabled.")

import io  # Add io import for DOCX export

# --- Environment variables and constants ---
FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "")
BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
MODEL_ID = "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507"
DB_PATH = "webtoon_sessions_v1.db"

# Target settings for webtoon - UPDATED FOR WEBTOON
TARGET_EPISODES = 40  # 40화 완결
PANELS_PER_EPISODE = 30  # 각 화당 30개 패널
TARGET_PANELS = TARGET_EPISODES * PANELS_PER_EPISODE  # 총 1200 패널

# Webtoon genres
WEBTOON_GENRES = {
    "로맨스": "Romance",
    "로판": "Romance Fantasy", 
    "판타지": "Fantasy",
    "현판": "Modern Fantasy",
    "무협": "Martial Arts",
    "스릴러": "Thriller",
    "일상": "Slice of Life",
    "개그": "Comedy",
    "스포츠": "Sports"
}

# --- Environment validation ---
if not FIREWORKS_API_KEY:
    logger.error("FIREWORKS_API_KEY not set. Application will not work properly.")
    FIREWORKS_API_KEY = "dummy_token_for_testing"

if not BRAVE_SEARCH_API_KEY:
    logger.warning("BRAVE_SEARCH_API_KEY not set. Web search features will be disabled.")

# --- Global variables ---
db_lock = threading.Lock()

# --- Data classes ---
@dataclass
class WebtoonBible:
    """Webtoon story bible for maintaining consistency"""
    genre: str = ""
    title: str = ""
    characters: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    settings: Dict[str, str] = field(default_factory=dict)
    plot_points: List[Dict[str, Any]] = field(default_factory=list)
    episode_hooks: Dict[int, str] = field(default_factory=dict)
    genre_elements: Dict[str, Any] = field(default_factory=dict)
    visual_style: Dict[str, Any] = field(default_factory=dict)
    panel_compositions: List[str] = field(default_factory=list)
    
@dataclass
class StoryboardPanel:
    """Individual storyboard panel"""
    panel_number: int
    scene_type: str  # wide, close-up, medium, establishing
    image_prompt: str  # Image generation prompt
    dialogue: List[str] = field(default_factory=list)
    narration: str = ""
    sound_effects: List[str] = field(default_factory=list)
    emotion_notes: str = ""
    camera_angle: str = ""
    background: str = ""

@dataclass
class EpisodeStoryboard:
    """Complete storyboard for one episode"""
    episode_number: int
    title: str
    panels: List[StoryboardPanel] = field(default_factory=list)
    total_panels: int = 30
    cliffhanger: str = ""

# --- Genre-specific prompts and elements ---
GENRE_ELEMENTS = {
    "로맨스": {
        "key_elements": ["감정선", "오해와 화해", "달콤한 순간", "질투", "고백"],
        "visual_styles": ["소프트 톤", "파스텔", "꽃 배경", "빛망울 효과", "분홍빛 필터"],
        "panel_types": ["클로즈업 감정샷", "투샷", "손 클로즈업", "눈빛 교환", "백허그"],
        "typical_scenes": ["카페 데이트", "우산 씬", "불꽃놀이", "옥상 고백", "공항 이별"]
    },
    "로판": {
        "key_elements": ["회귀/빙의", "드레스", "무도회", "마법", "신분 상승"],
        "visual_styles": ["화려한 의상", "유럽풍 배경", "반짝이 효과", "마법진", "성 배경"],
        "panel_types": ["전신샷", "드레스 디테일", "마법 이펙트", "회상씬", "충격 리액션"],
        "typical_scenes": ["무도회장", "정원 산책", "서재", "마법 수업", "알현실"]
    },
    "판타지": {
        "key_elements": ["마법체계", "레벨업", "던전", "길드", "모험"],
        "visual_styles": ["다이나믹 액션", "이펙트 강조", "몬스터 디자인", "판타지 배경", "빛 효과"],
        "panel_types": ["액션씬", "풀샷 전투", "스킬 발동", "몬스터 등장", "파워업"],
        "typical_scenes": ["던전 입구", "보스전", "길드 회관", "수련장", "아이템 획득"]
    },
    "현판": {
        "key_elements": ["게이트", "헌터", "각성", "현대 도시", "능력"],
        "visual_styles": ["도시 배경", "네온 효과", "현대적 액션", "특수 효과", "어반 판타지"],
        "panel_types": ["도시 전경", "능력 발현", "게이트 출현", "전투 액션", "일상 대비"],
        "typical_scenes": ["게이트 현장", "헌터 협회", "훈련장", "병원", "학교"]
    },
    "무협": {
        "key_elements": ["무공", "문파", "강호", "복수", "의협"],
        "visual_styles": ["동양풍", "먹 효과", "기 표현", "중국풍 의상", "산수화 배경"],
        "panel_types": ["검술 동작", "경공술", "기공 수련", "대결 구도", "폭발 이펙트"],
        "typical_scenes": ["무림맹", "객잔", "절벽", "폭포 수련", "비무대회"]
    },
    "스릴러": {
        "key_elements": ["서스펜스", "공포", "추격", "심리전", "반전"],
        "visual_styles": ["어두운 톤", "그림자 강조", "대비 효과", "불안한 구도", "붉은색 강조"],
        "panel_types": ["극단 클로즈업", "Dutch angle", "실루엣", "충격 컷", "공포 연출"],
        "typical_scenes": ["어두운 골목", "폐건물", "지하실", "추격씬", "대치 상황"]
    },
    "일상": {
        "key_elements": ["일상", "공감", "소소한 재미", "관계", "성장"],
        "visual_styles": ["따뜻한 색감", "부드러운 선", "일상 배경", "캐주얼", "편안한 구도"],
        "panel_types": ["일상 컷", "리액션", "대화씬", "배경 묘사", "감정 표현"],
        "typical_scenes": ["집", "학교", "회사", "동네", "편의점"]
    },
    "개그": {
        "key_elements": ["개그", "패러디", "과장", "반전", "슬랩스틱"],
        "visual_styles": ["과장된 표정", "데포르메", "효과선", "말풍선 연출", "파격 구도"],
        "panel_types": ["과장 리액션", "개그 컷", "패러디", "충격 표정", "망가짐"],
        "typical_scenes": ["개그 상황", "일상 붕괴", "오해 상황", "추격전", "단체 개그"]
    },
    "스포츠": {
        "key_elements": ["경기", "훈련", "팀워크", "라이벌", "성장"],
        "visual_styles": ["다이나믹", "스피드선", "땀 표현", "근육 묘사", "경기장"],
        "panel_types": ["액션 컷", "결정적 순간", "전신 동작", "표정 클로즈업", "경기 전경"],
        "typical_scenes": ["경기장", "훈련장", "라커룸", "벤치", "시상대"]
    }
}

# Panel composition types
PANEL_COMPOSITIONS = {
    "establishing": "전체 배경과 분위기를 보여주는 원거리 샷",
    "wide": "캐릭터와 배경이 모두 보이는 와이드 샷",
    "medium": "캐릭터 상반신이 보이는 미디엄 샷",
    "close_up": "얼굴이나 표정에 집중하는 클로즈업",
    "extreme_close_up": "눈, 입 등 특정 부위 극단 클로즈업",
    "over_shoulder": "어깨 너머로 보는 오버숄더 샷",
    "dutch_angle": "긴장감을 위한 기울어진 앵글",
    "bird_eye": "위에서 내려다보는 부감 샷",
    "worm_eye": "아래서 올려다보는 앙각 샷",
    "action": "움직임을 강조하는 액션 샷"
}

# --- Core logic classes ---
class WebtoonTracker:
    """Webtoon narrative and storyboard tracker"""
    def __init__(self):
        self.story_bible = WebtoonBible()
        self.episode_storyboards: Dict[int, EpisodeStoryboard] = {}
        self.episodes: Dict[int, str] = {}
        self.total_panel_count = 0
        
    def set_genre(self, genre: str):
        """Set the webtoon genre"""
        self.story_bible.genre = genre
        self.story_bible.genre_elements = GENRE_ELEMENTS.get(genre, {})
        
    def add_storyboard(self, episode_num: int, storyboard: EpisodeStoryboard):
        """Add episode storyboard"""
        self.episode_storyboards[episode_num] = storyboard
        self.total_panel_count += len(storyboard.panels)

class WebtoonDatabase:
    """Database management for webtoon system"""
    @staticmethod
    def init_db():
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            cursor = conn.cursor()
            
            # Sessions table with genre
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    user_query TEXT NOT NULL,
                    genre TEXT NOT NULL,
                    language TEXT NOT NULL,
                    title TEXT,
                    created_at TEXT DEFAULT (datetime('now')),
                    updated_at TEXT DEFAULT (datetime('now')),
                    status TEXT DEFAULT 'active',
                    current_episode INTEGER DEFAULT 0,
                    total_episodes INTEGER DEFAULT 40,
                    planning_doc TEXT,
                    story_bible TEXT,
                    visual_style TEXT
                )
            ''')
            
            # Storyboards table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS storyboards (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    episode_number INTEGER NOT NULL,
                    title TEXT,
                    storyboard_data TEXT,
                    panel_count INTEGER DEFAULT 30,
                    status TEXT DEFAULT 'pending',
                    created_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id),
                    UNIQUE(session_id, episode_number)
                )
            ''')
            
            # Panels table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS panels (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    episode_number INTEGER NOT NULL,
                    panel_number INTEGER NOT NULL,
                    scene_type TEXT,
                    image_prompt TEXT,
                    dialogue TEXT,
                    narration TEXT,
                    sound_effects TEXT,
                    created_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            ''')
            
            conn.commit()

    @staticmethod
    @contextmanager
    def get_db():
        with db_lock:
            conn = sqlite3.connect(DB_PATH, timeout=30.0)
            conn.row_factory = sqlite3.Row
            try:
                yield conn
            finally:
                conn.close()

    @staticmethod
    def create_session(user_query: str, genre: str, language: str) -> str:
        session_id = hashlib.md5(f"{user_query}{genre}{datetime.now()}".encode()).hexdigest()
        with WebtoonDatabase.get_db() as conn:
            conn.cursor().execute(
                '''INSERT INTO sessions (session_id, user_query, genre, language) 
                   VALUES (?, ?, ?, ?)''',
                (session_id, user_query, genre, language)
            )
            conn.commit()
        return session_id

    @staticmethod
    def save_storyboard(session_id: str, episode_num: int, storyboard: EpisodeStoryboard):
        with WebtoonDatabase.get_db() as conn:
            cursor = conn.cursor()
            
            # Save storyboard
            cursor.execute('''
                INSERT INTO storyboards (session_id, episode_number, title, 
                                        storyboard_data, panel_count, status)
                VALUES (?, ?, ?, ?, ?, 'complete')
                ON CONFLICT(session_id, episode_number)
                DO UPDATE SET title=?, storyboard_data=?, panel_count=?, status='complete'
            ''', (session_id, episode_num, storyboard.title, 
                  json.dumps(asdict(storyboard)), len(storyboard.panels),
                  storyboard.title, json.dumps(asdict(storyboard)), len(storyboard.panels)))
            
            # Save individual panels
            for panel in storyboard.panels:
                cursor.execute('''
                    INSERT INTO panels (session_id, episode_number, panel_number,
                                      scene_type, image_prompt, dialogue, narration, sound_effects)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ''', (session_id, episode_num, panel.panel_number,
                      panel.scene_type, panel.image_prompt,
                      json.dumps(panel.dialogue), panel.narration,
                      json.dumps(panel.sound_effects)))
            
            conn.commit()

# --- LLM Integration ---
class WebtoonSystem:
    """Webtoon planning and storyboard generation system"""
    def __init__(self):
        self.api_key = FIREWORKS_API_KEY
        self.api_url = API_URL
        self.model_id = MODEL_ID
        self.tracker = WebtoonTracker()
        self.current_session_id = None
        WebtoonDatabase.init_db()

    def create_headers(self):
        return {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

    # --- Prompt generation functions ---
    def create_planning_prompt(self, query: str, genre: str, language: str) -> str:
        """Create initial planning prompt for webtoon"""
        genre_info = GENRE_ELEMENTS.get(genre, {})
        
        lang_prompts = {
            "Korean": f"""한국 웹툰 시장을 겨냥한 {genre} 장르 웹툰을 기획하세요.

**[핵심 스토리 설정 - 반드시 이 내용을 중심으로 전개하세요]**
{query}

**장르:** {genre}
**목표:** 40화 완결 웹툰

⚠️ **중요**: 위에 제시된 스토리 설정을 반드시 기반으로 하여 플롯을 구성하세요.

**장르 필수 요소:**
- 핵심 요소: {', '.join(genre_info.get('key_elements', []))}
- 비주얼 스타일: {', '.join(genre_info.get('visual_styles', []))}
- 주요 씬: {', '.join(genre_info.get('typical_scenes', []))}

**전체 구성 (40화 기준):**
1. **1-5화**: 세계관 소개, 주인공 등장, 핵심 갈등 제시
2. **6-15화**: 갈등 심화, 캐릭터 관계 발전, 서브플롯 전개
3. **16-25화**: 중간 클라이맥스, 반전, 새로운 위기
4. **26-35화**: 최종 갈등 고조, 결전 준비
5. **36-40화**: 클라이맥스, 해결, 에필로그

**웹툰 특화 요소:**
- 매 화 마지막 강력한 클리프행어
- 세로 스크롤에 적합한 연출
- 임팩트 있는 대사와 액션
- 독자 댓글을 유발하는 전개

다음 형식으로 작성하세요:

📚 **작품 제목:** [임팩트 있는 제목]

🎨 **비주얼 컨셉:**
- 그림체: [작품에 어울리는 그림체]
- 색감: [주요 색상 톤]
- 캐릭터 디자인 특징: [주인공들의 비주얼 특징]

👥 **주요 캐릭터:**
- 주인공: [이름] - [외모 특징, 성격, 목표]
- 캐릭터2: [이름] - [역할, 특징]
- 캐릭터3: [이름] - [역할, 특징]

📖 **시놉시스:**
[3-4줄로 전체 스토리 요약]

📝 **40화 전체 구성안:**
각 화별로 핵심 사건과 클리프행어를 포함하여 작성하세요.

1화: [제목] - [핵심 사건] - 클리프행어: [충격적인 마무리]
2화: [제목] - [핵심 사건] - 클리프행어: [충격적인 마무리]
...
40화: [제목] - [핵심 사건] - [대단원의 마무리]""",

            "English": f"""Plan a Korean-style webtoon for {genre} genre.

**[Core Story Setting - MUST base the story on this]**
{query}

**Genre:** {genre}
**Goal:** 40 episodes webtoon

⚠️ **IMPORTANT**: You MUST base the plot on the story setting provided above.

**Genre Requirements:**
- Key elements: {', '.join(genre_info.get('key_elements', []))}
- Visual styles: {', '.join(genre_info.get('visual_styles', []))}
- Typical scenes: {', '.join(genre_info.get('typical_scenes', []))}

**Overall Structure (40 episodes):**
1. **Episodes 1-5**: World introduction, protagonist debut, core conflict
2. **Episodes 6-15**: Deepening conflict, character development, subplots
3. **Episodes 16-25**: Mid-climax, plot twist, new crisis
4. **Episodes 26-35**: Final conflict escalation, preparation for showdown
5. **Episodes 36-40**: Climax, resolution, epilogue

**Webtoon-specific elements:**
- Strong cliffhanger at end of each episode
- Vertical scroll-optimized directing
- Impactful dialogue and action
- Comment-inducing development

Format as follows:

📚 **Title:** [Impactful title]

🎨 **Visual Concept:**
- Art style: [Suitable art style]
- Color tone: [Main color palette]
- Character design: [Visual characteristics]

👥 **Main Characters:**
- Protagonist: [Name] - [Appearance, personality, goal]
- Character2: [Name] - [Role, traits]
- Character3: [Name] - [Role, traits]

📖 **Synopsis:**
[3-4 line story summary]

📝 **40 Episode Structure:**
Include key events and cliffhangers for each episode.

Episode 1: [Title] - [Key event] - Cliffhanger: [Shocking ending]
Episode 2: [Title] - [Key event] - Cliffhanger: [Shocking ending]
...
Episode 40: [Title] - [Key event] - [Grand finale]"""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_storyboard_prompt(self, episode_num: int, plot_outline: str, 
                                genre: str, language: str) -> str:
        """Create prompt for episode 1 storyboard with 30 panels"""
        genre_info = GENRE_ELEMENTS.get(genre, {})
        
        lang_prompts = {
            "Korean": f"""웹툰 {episode_num}화 스토리보드를 30개 패널로 작성하세요.

**장르:** {genre}
**1화 내용:** {self._extract_episode_plan(plot_outline, episode_num)}

**패널 구성 지침:**
- 총 30개 패널로 구성
- 다양한 샷 사이즈 활용 (와이드샷, 클로즈업, 미디엄샷 등)
- 장르 특성에 맞는 연출: {', '.join(genre_info.get('panel_types', []))}
- 세로 스크롤 웹툰에 최적화된 구성

**각 패널별로 다음을 포함하여 작성:**

패널 1:
- 샷 타입: [establishing/wide/medium/close_up 등]
- 이미지 프롬프트: [상세한 한글 이미지 생성 프롬프트]
- 대사: [캐릭터 대사가 있다면]
- 나레이션: [해설이 있다면]
- 효과음: [필요한 효과음]
- 배경: [배경 설명]

...이런 식으로 30개 패널 모두 작성

**중요 연출 포인트:**
1. 첫 패널은 임팩트 있는 establishing shot
2. 감정선은 클로즈업으로 강조
3. 액션은 다이나믹한 앵글 활용
4. 마지막 패널은 강력한 클리프행어
5. 대사와 이미지가 조화롭게 구성

반드시 30개 패널을 모두 작성하세요.""",

            "English": f"""Create Episode {episode_num} storyboard with 30 panels.

**Genre:** {genre}
**Episode 1 content:** {self._extract_episode_plan(plot_outline, episode_num)}

**Panel composition guidelines:**
- Total 30 panels
- Various shot sizes (wide, close-up, medium, etc.)
- Genre-appropriate directing: {', '.join(genre_info.get('panel_types', []))}
- Optimized for vertical scroll webtoon

**For each panel include:**

Panel 1:
- Shot type: [establishing/wide/medium/close_up etc]
- Image prompt: [Detailed Korean image generation prompt]
- Dialogue: [Character dialogue if any]
- Narration: [Narration if any]
- Sound effects: [Required sound effects]
- Background: [Background description]

...continue for all 30 panels

**Key directing points:**
1. First panel: impactful establishing shot
2. Emphasize emotions with close-ups
3. Dynamic angles for action
4. Last panel: powerful cliffhanger
5. Harmonious dialogue and image composition

Must write all 30 panels."""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def _extract_episode_plan(self, plot_outline: str, episode_num: int) -> str:
        """Extract specific episode plan from outline"""
        lines = plot_outline.split('\n')
        episode_section = []
        capturing = False
        
        patterns = [
            f"{episode_num}화:", f"Episode {episode_num}:", 
            f"제{episode_num}화:", f"EP{episode_num}:",
            f"{episode_num}.", f"[{episode_num}]"
        ]
        
        next_patterns = [
            f"{episode_num+1}화:", f"Episode {episode_num+1}:", 
            f"제{episode_num+1}화:", f"EP{episode_num+1}:",
            f"{episode_num+1}.", f"[{episode_num+1}]"
        ]
        
        for line in lines:
            if any(pattern in line for pattern in patterns):
                capturing = True
                episode_section.append(line)
            elif capturing and any(pattern in line for pattern in next_patterns):
                break
            elif capturing:
                episode_section.append(line)
        
        if episode_section:
            return '\n'.join(episode_section)
        
        return f"에피소드 {episode_num} 내용을 플롯에서 참조하여 작성하세요."

    # --- LLM call functions ---
    def call_llm_sync(self, messages: List[Dict[str, str]], role: str, language: str) -> str:
        full_content = ""
        for chunk in self.call_llm_streaming(messages, role, language):
            full_content += chunk
        if full_content.startswith("❌"):
            raise Exception(f"LLM Call Failed: {full_content}")
        return full_content

    def call_llm_streaming(self, messages: List[Dict[str, str]], role: str, 
                          language: str) -> Generator[str, None, None]:
        try:
            system_prompts = self.get_system_prompts(language)
            full_messages = [{"role": "system", "content": system_prompts.get(role, "")}, *messages]
            
            max_tokens = 15000 if role == "storyboarder" else 10000
            
            payload = {
                "model": self.model_id,
                "messages": full_messages,
                "max_tokens": max_tokens,
                "temperature": 0.7,
                "top_p": 1,
                "top_k": 40,
                "presence_penalty": 0,
                "frequency_penalty": 0,
                "stream": True
            }
            
            response = requests.post(
                self.api_url, 
                headers=self.create_headers(), 
                json=payload, 
                stream=True, 
                timeout=180
            )
            
            if response.status_code != 200:
                yield f"❌ API Error (Status Code: {response.status_code})"
                return
            
            buffer = ""
            for line in response.iter_lines():
                if not line:
                    continue
                
                try:
                    line_str = line.decode('utf-8').strip()
                    if not line_str.startswith("data: "):
                        continue
                    
                    data_str = line_str[6:]
                    if data_str == "[DONE]":
                        break
                    
                    data = json.loads(data_str)
                    choices = data.get("choices", [])
                    if choices and choices[0].get("delta", {}).get("content"):
                        content = choices[0]["delta"]["content"]
                        buffer += content
                        
                        if len(buffer) >= 50 or '\n' in buffer:
                            yield buffer
                            buffer = ""
                            time.sleep(0.01)
                    
                except Exception as e:
                    logger.error(f"Chunk processing error: {str(e)}")
                    continue
            
            if buffer:
                yield buffer
                
        except Exception as e:
            logger.error(f"Streaming error: {type(e).__name__}: {str(e)}")
            yield f"❌ Error occurred: {str(e)}"

    def get_system_prompts(self, language: str) -> Dict[str, str]:
        """System prompts for webtoon roles"""
        base_prompts = {
            "Korean": {
                "planner": """당신은 한국 웹툰 시장을 완벽히 이해하는 웹툰 기획자입니다.
독자를 사로잡는 스토리와 비주얼 연출을 기획합니다.
40화 완결 구조로 완벽한 기승전결을 설계합니다.
각 화마다 강력한 클리프행어로 다음 화를 기대하게 만듭니다.
장르별 특성과 독자층을 정확히 파악합니다.

⚠️ 가장 중요한 원칙: 사용자가 제공한 스토리 설정을 절대적으로 우선시하고, 이를 중심으로 모든 플롯을 구성합니다.""",
                
                "storyboarder": """당신은 웹툰 스토리보드 전문가입니다.
30개 패널로 한 화를 완벽하게 구성합니다.
세로 스크롤에 최적화된 연출을 합니다.
각 패널마다 상세한 이미지 프롬프트를 한글로 작성합니다.
대사, 나레이션, 효과음을 적절히 배치합니다.
다양한 카메라 앵글과 샷 사이즈를 활용합니다.
감정선과 액션을 시각적으로 극대화합니다.

⚠️ 가장 중요한 원칙: 반드시 30개 패널을 모두 작성하고, 각 패널마다 이미지 생성이 가능한 구체적인 한글 프롬프트를 제공합니다."""
            },
            "English": {
                "planner": """You perfectly understand the Korean webtoon market.
Design stories and visual direction that captivate readers.
Create perfect story structure in 40 episodes.
Make readers anticipate next episode with strong cliffhangers.
Accurately understand genre characteristics and readership.

⚠️ Most important principle: Absolutely prioritize the user's story setting and build all plots around it.""",
                
                "storyboarder": """You are a webtoon storyboard specialist.
Perfectly compose one episode with 30 panels.
Optimize direction for vertical scrolling.
Write detailed image prompts in Korean for each panel.
Properly place dialogue, narration, and sound effects.
Use various camera angles and shot sizes.
Visually maximize emotions and action.

⚠️ Most important principle: Must write all 30 panels and provide specific Korean prompts for image generation for each panel."""
            }
        }
        
        return base_prompts.get(language, base_prompts["Korean"])

    # --- Main process ---
    def process_webtoon_stream(self, query: str, genre: str, language: str, 
                              session_id: Optional[str] = None) -> Generator[Tuple[str, str, str, str], None, None]:
        """Webtoon planning and storyboard generation process"""
        try:
            if not session_id:
                self.current_session_id = WebtoonDatabase.create_session(query, genre, language)
                self.tracker.set_genre(genre)
                logger.info(f"Created new session: {self.current_session_id}")
                self.original_query = query
            else:
                self.current_session_id = session_id

            # Phase 1: Generate planning document (40 episodes structure)
            yield "🎬 웹툰 기획안 작성 중...", "", f"장르: {genre}", self.current_session_id
            
            planning_prompt = self.create_planning_prompt(query, genre, language)
            planning_doc = self.call_llm_sync(
                [{"role": "user", "content": planning_prompt}], 
                "planner", language
            )
            
            self.planning_doc = planning_doc
            
            yield "✅ 기획안 완성!", planning_doc, "40화 구성 완료", self.current_session_id

            # Phase 2: Generate Episode 1 Storyboard (30 panels)
            yield "🎨 1화 스토리보드 작성 중...", planning_doc, "30개 패널 구성 중", self.current_session_id
            
            storyboard_prompt = self.create_storyboard_prompt(1, planning_doc, genre, language)
            storyboard_content = self.call_llm_sync(
                [{"role": "user", "content": storyboard_prompt}],
                "storyboarder", language
            )
            
            # Parse storyboard into structured format
            storyboard = self.parse_storyboard(storyboard_content, 1)
            
            # Save to database
            WebtoonDatabase.save_storyboard(self.current_session_id, 1, storyboard)
            
            # Format final output
            final_output = f"{planning_doc}\n\n{'='*50}\n\n## 📋 1화 스토리보드 (30 패널)\n\n{storyboard_content}"
            
            yield "🎉 완성!", final_output, "기획안 + 1화 스토리보드 완료", self.current_session_id

        except Exception as e:
            logger.error(f"Webtoon generation error: {e}", exc_info=True)
            yield f"❌ 오류 발생: {e}", "", "오류", self.current_session_id

    def parse_storyboard(self, content: str, episode_num: int) -> EpisodeStoryboard:
        """Parse storyboard text into structured format"""
        storyboard = EpisodeStoryboard(episode_number=episode_num, title=f"{episode_num}화")
        
        # Simple parsing - can be enhanced
        panels = []
        current_panel = None
        panel_number = 0
        
        lines = content.split('\n')
        for line in lines:
            if '패널' in line or 'Panel' in line:
                if current_panel:
                    panels.append(current_panel)
                panel_number += 1
                # Initialize with empty image_prompt as required argument
                current_panel = StoryboardPanel(
                    panel_number=panel_number, 
                    scene_type="medium",
                    image_prompt=""  # Initialize with empty string
                )
            elif current_panel:
                if '이미지 프롬프트:' in line or 'Image prompt:' in line:
                    current_panel.image_prompt = line.split(':', 1)[1].strip()
                elif '대사:' in line or 'Dialogue:' in line:
                    dialogue = line.split(':', 1)[1].strip()
                    if dialogue:
                        current_panel.dialogue.append(dialogue)
                elif '나레이션:' in line or 'Narration:' in line:
                    current_panel.narration = line.split(':', 1)[1].strip()
                elif '효과음:' in line or 'Sound effects:' in line:
                    effects = line.split(':', 1)[1].strip()
                    if effects:
                        current_panel.sound_effects.append(effects)
                elif '샷 타입:' in line or 'Shot type:' in line:
                    current_panel.scene_type = line.split(':', 1)[1].strip()
        
        if current_panel:
            panels.append(current_panel)
        
        storyboard.panels = panels[:30]  # Limit to 30 panels
        return storyboard

# --- Export functions ---
def export_to_txt(planning_doc: str, storyboard: str, genre: str, title: str = "") -> str:
    """Export webtoon planning and storyboard to TXT format"""
    content = f"{'=' * 50}\n"
    content += f"{title if title else genre + ' 웹툰'}\n"
    content += f"{'=' * 50}\n\n"
    content += f"장르: {genre}\n"
    content += f"총 40화 기획\n"
    content += f"{'=' * 50}\n\n"
    
    content += "📚 기획안\n\n"
    content += planning_doc
    content += f"\n\n{'=' * 50}\n\n"
    
    content += "🎨 1화 스토리보드 (30 패널)\n\n"
    content += storyboard
    
    return content

def export_to_docx(planning_doc: str, storyboard: str, genre: str, title: str = "") -> bytes:
    """Export webtoon planning and storyboard to DOCX format"""
    if not DOCX_AVAILABLE:
        raise Exception("python-docx is not installed")
    
    doc = Document()
    
    # Title
    doc.add_heading(title if title else f"{genre} 웹툰", 0)
    
    # Info
    doc.add_paragraph(f"장르: {genre}")
    doc.add_paragraph("총 40화 기획")
    doc.add_page_break()
    
    # Planning document
    doc.add_heading("📚 기획안", 1)
    for line in planning_doc.split('\n'):
        if line.strip():
            if line.startswith('#'):
                doc.add_heading(line.replace('#', '').strip(), 2)
            else:
                doc.add_paragraph(line)
    
    doc.add_page_break()
    
    # Storyboard
    doc.add_heading("🎨 1화 스토리보드 (30 패널)", 1)
    for line in storyboard.split('\n'):
        if line.strip():
            if '패널' in line or 'Panel' in line:
                doc.add_heading(line, 2)
            else:
                doc.add_paragraph(line)
    
    # Save to bytes
    bytes_io = io.BytesIO()
    doc.save(bytes_io)
    bytes_io.seek(0)
    return bytes_io.getvalue()

def generate_random_webtoon_theme(genre: str, language: str) -> str:
    """Generate random webtoon theme"""
    templates = {
        "로맨스": [
            "재벌 3세 상사와 신입사원의 비밀 계약연애",
            "고등학교 때 첫사랑과 10년 만의 재회",
            "냉혈 검사와 열혈 변호사의 법정 로맨스"
        ],
        "로판": [
            "악녀로 빙의했는데 1년 후 처형 예정",
            "회귀한 황녀, 버려진 왕자와 손잡다",
            "계약결혼한 북부 공작이 집착남이 되었다"
        ],
        "판타지": [
            "F급 헌터가 SSS급 네크로맨서로 각성",
            "100층 탑을 역주행하는 회귀자",
            "버그로 최강이 된 게임 속 NPC"
        ],
        "현판": [
            "무능력자인 줄 알았는데 SSS급 생산직",
            "게이트 속에서 10년, 돌아온 최강자",
            "헌터 고등학교의 숨겨진 랭킹 1위"
        ],
        "무협": [
            "천하제일문 막내가 마교 교주 제자가 되다",
            "100년 전으로 회귀한 화산파 장문인",
            "폐급 무공으로 천하를 제패하다"
        ],
        "스릴러": [
            "폐교에 갇힌 동창회, 한 명씩 사라진다",
            "타임루프 속 연쇄살인범 찾기",
            "내 남편이 사이코패스였다"
        ],
        "일상": [
            "편의점 알바생의 소소한 일상",
            "30대 직장인의 퇴사 준비 일기",
            "우리 동네 고양이들의 비밀 회의"
        ],
        "개그": [
            "이세계 용사인데 스탯이 이상하다",
            "우리 학교 선생님은 전직 마왕",
            "좀비 아포칼립스인데 나만 개그 캐릭터"
        ],
        "스포츠": [
            "벤치 멤버에서 에이스가 되기까지",
            "여자 야구부 창설기",
            "은퇴 선수의 코치 도전기"
        ]
    }
    
    genre_themes = templates.get(genre, templates["로맨스"])
    return random.choice(genre_themes)

# --- UI functions ---
def format_planning_display(planning_doc: str) -> str:
    """Format planning document for display"""
    return planning_doc

def format_storyboard_display(storyboard: str) -> str:
    """Format storyboard for display"""
    return storyboard

# --- Gradio interface ---
def create_interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="K-Webtoon Storyboard Generator") as interface:
        gr.HTML("""
        <style>
            .main-header {
                text-align: center;
                margin-bottom: 2rem;
                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                padding: 2rem;
                border-radius: 15px;
                color: white;
            }
            
            .header-title {
                font-size: 3rem;
                margin-bottom: 1rem;
                text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
            }
            
            .header-subtitle {
                font-size: 1.2rem;
                margin-bottom: 0.5rem;
            }
            
            .header-description {
                margin-bottom: 1.5rem;
            }
        </style>
        
        <div class="main-header">
            <h1 class="header-title">🎨 K-Webtoon Storyboard Generator</h1>
            <p class="header-subtitle">한국형 웹툰 기획 및 스토리보드 자동 생성 시스템</p>
            <p class="header-description">40화 전체 기획안과 1화 스토리보드(30패널)를 생성합니다</p>
        </div>
        """)
        
        # State
        current_session_id = gr.State(None)
        
        with gr.Tab("🎨 웹툰 기획하기"):
            with gr.Group():
                gr.Markdown("### 🎯 웹툰 설정")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        genre_select = gr.Radio(
                            choices=list(WEBTOON_GENRES.keys()),
                            value="로맨스",
                            label="장르 선택",
                            info="원하는 장르를 선택하세요"
                        )
                        
                        query_input = gr.Textbox(
                            label="스토리 콘셉트",
                            placeholder="웹툰의 기본 설정이나 주제를 입력하세요...",
                            lines=3
                        )
                        
                        with gr.Row():
                            random_btn = gr.Button("🎲 랜덤 테마", variant="secondary")
                            submit_btn = gr.Button("📝 기획 시작", variant="primary", size="lg")
                    
                    with gr.Column(scale=1):
                        language_select = gr.Radio(
                            choices=["Korean", "English"],
                            value="Korean",
                            label="언어"
                        )
                        
                        gr.Markdown("""
                        **장르별 특징:**
                        - 로맨스: 감정선, 관계 중심
                        - 로판: 회귀/빙의 판타지
                        - 판타지: 모험과 성장
                        - 현판: 현대 배경 능력물
                        - 무협: 무공과 복수
                        - 스릴러: 서스펜스와 반전
                        - 일상: 공감과 힐링
                        - 개그: 웃음과 패러디
                        - 스포츠: 열정과 성장
                        """)
                
                status_text = gr.Textbox(
                    label="진행 상황",
                    interactive=False,
                    value="장르를 선택하고 콘셉트를 입력하세요"
                )
            
            # Output
            with gr.Row():
                with gr.Column():
                    planning_display = gr.Markdown("*기획안이 여기에 표시됩니다*")
                
                with gr.Column():
                    storyboard_display = gr.Markdown("*1화 스토리보드가 여기에 표시됩니다*")
                    
                    with gr.Row():
                        download_format = gr.Radio(
                            choices=["TXT", "DOCX"],
                            value="TXT",
                            label="다운로드 형식"
                        )
                        download_btn = gr.Button("📥 다운로드", variant="secondary")
                    
                    download_file = gr.File(visible=False)
        
        with gr.Tab("📚 콘셉트 예시"):
            gr.Markdown("""
            ### 장르별 인기 콘셉트
            
            **로맨스**
            - 계약결혼한 재벌과 평범한 회사원
            - 첫사랑과의 운명적 재회
            - 라이벌에서 연인이 되기까지
            
            **판타지**
            - F급에서 SSS급으로 각성
            - 회귀자의 두 번째 기회
            - 버그 캐릭터의 세계 정복
            
            **현판**
            - 숨겨진 최강자의 일상
            - 게이트와 헌터의 세계
            - 각성자 학교의 비밀
            
            **무협**
            - 천재의 몰락과 부활
            - 정파와 마교의 경계
            - 전설의 비급을 얻다
            """)
        
        # Event handlers
        def process_query(query, genre, language, session_id):
            system = WebtoonSystem()
            planning = ""
            storyboard = ""
            
            for planning_content, storyboard_content, status, new_session_id in system.process_webtoon_stream(query, genre, language, session_id):
                # Split the combined output if it contains both planning and storyboard
                if '1화 스토리보드' in storyboard_content:
                    parts = storyboard_content.split('=' * 50)
                    if len(parts) >= 2:
                        planning = parts[0]
                        storyboard = ('=' * 50) + ('=' * 50).join(parts[1:])
                    else:
                        planning = planning_content
                        storyboard = storyboard_content
                else:
                    planning = planning_content
                    storyboard = storyboard_content
                    
                yield planning, storyboard, status, new_session_id
        
        def handle_random_theme(genre, language):
            return generate_random_webtoon_theme(genre, language)
        
        def handle_download(download_format, session_id, planning, storyboard, genre):
            """Handle download request"""
            try:
                title = f"{genre} 웹툰"
                
                if download_format == "TXT":
                    content = export_to_txt(planning, storyboard, genre, title)
                    
                    with tempfile.NamedTemporaryFile(mode='w', encoding='utf-8', 
                                                   suffix='.txt', delete=False) as f:
                        f.write(content)
                        return f.name
                
                elif download_format == "DOCX":
                    if not DOCX_AVAILABLE:
                        gr.Warning("DOCX export requires python-docx library")
                        return None
                    
                    content = export_to_docx(planning, storyboard, genre, title)
                    
                    with tempfile.NamedTemporaryFile(mode='wb', suffix='.docx', 
                                                   delete=False) as f:
                        f.write(content)
                        return f.name
                        
            except Exception as e:
                logger.error(f"Download error: {e}")
                gr.Warning(f"다운로드 중 오류 발생: {str(e)}")
                return None
        
        # Connect events
        submit_btn.click(
            fn=process_query,
            inputs=[query_input, genre_select, language_select, current_session_id],
            outputs=[planning_display, storyboard_display, status_text, current_session_id]
        )
        
        random_btn.click(
            fn=handle_random_theme,
            inputs=[genre_select, language_select],
            outputs=[query_input]
        )
        
        download_btn.click(
            fn=handle_download,
            inputs=[download_format, current_session_id, planning_display, storyboard_display, genre_select],
            outputs=[download_file]
        ).then(
            fn=lambda x: gr.update(visible=True) if x else gr.update(visible=False),
            inputs=[download_file],
            outputs=[download_file]
        )
        
        # Examples
        gr.Examples(
            examples=[
                ["재벌 3세 상사와 신입사원의 비밀 계약연애", "로맨스"],
                ["악녀로 빙의했는데 1년 후 처형 예정", "로판"],
                ["F급 헌터가 SSS급 네크로맨서로 각성", "판타지"],
                ["무능력자인 줄 알았는데 SSS급 생산직", "현판"],
                ["천하제일문 막내가 마교 교주 제자가 되다", "무협"],
                ["폐교에 갇힌 동창회, 한 명씩 사라진다", "스릴러"],
                ["편의점 알바생의 소소한 일상", "일상"],
                ["이세계 용사인데 스탯이 이상하다", "개그"],
                ["벤치 멤버에서 에이스가 되기까지", "스포츠"]
            ],
            inputs=[query_input, genre_select]
        )
    
    return interface

# Main
if __name__ == "__main__":
    logger.info("K-Webtoon Storyboard Generator Starting...")
    logger.info("=" * 60)
    
    # Environment check
    logger.info(f"API Endpoint: {API_URL}")
    logger.info(f"Model: {MODEL_ID}")
    logger.info(f"Target: {TARGET_EPISODES} episodes, {PANELS_PER_EPISODE} panels per episode")
    logger.info("Genres: " + ", ".join(WEBTOON_GENRES.keys()))
    
    logger.info("=" * 60)
    
    # Initialize database
    logger.info("Initializing database...")
    WebtoonDatabase.init_db()
    logger.info("Database ready.")
    
    # Launch interface
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )