app-backup.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
import logging
import re
import tempfile
from pathlib import Path
import sqlite3
import hashlib
import threading
from contextlib import contextmanager

# 로깅 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

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

# 환경 변수에서 토큰 가져오기
FRIENDLI_TOKEN = os.getenv("FRIENDLI_TOKEN", "")
BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
API_URL = "https://api.friendli.ai/dedicated/v1/chat/completions"
MODEL_ID = "dep89a2fld32mcm"
TEST_MODE = os.getenv("TEST_MODE", "false").lower() == "true"

# 환경 변수 검증
if not FRIENDLI_TOKEN and not TEST_MODE:
    logger.warning("FRIENDLI_TOKEN not set and TEST_MODE is false. Application will run in test mode.")
    TEST_MODE = True

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

# 전역 변수
conversation_history = []
selected_language = "English"  # 기본 언어

# DB 경로
DB_PATH = "novel_sessions.db"
db_lock = threading.Lock()

# Stage 번호 상수 - 10명의 작가로 변경
WRITER_DRAFT_STAGES = [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]  # 작가 초안
WRITER_REVISION_STAGES = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32]  # 작가 수정본
TOTAL_WRITERS = 10  # 총 작가 수

class WebSearchIntegration:
    """Brave Search API integration for research"""
    
    def __init__(self):
        self.brave_api_key = BRAVE_SEARCH_API_KEY
        self.search_url = "https://api.search.brave.com/res/v1/web/search"
        self.enabled = bool(self.brave_api_key)
        
    def search(self, query: str, count: int = 5, language: str = "en") -> List[Dict]:
        """Perform web search using Brave Search API"""
        if not self.enabled:
            return []
            
        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": self.brave_api_key
        }
        
        # 언어에 따른 검색 설정
        search_lang = "ko" if language == "Korean" else "en"
        
        params = {
            "q": query,
            "count": count,
            "search_lang": search_lang,
            "text_decorations": False,
            "safesearch": "moderate"
        }
        
        try:
            response = requests.get(self.search_url, headers=headers, params=params, timeout=10)
            if response.status_code == 200:
                results = response.json()
                web_results = results.get("web", {}).get("results", [])
                logger.info(f"Search successful: Found {len(web_results)} results for '{query}'")
                return web_results
            else:
                logger.error(f"Search API error: {response.status_code}")
                return []
        except Exception as e:
            logger.error(f"Search error: {str(e)}")
            return []
    
    def extract_relevant_info(self, results: List[Dict], max_chars: int = 3000) -> str:
        """Extract relevant information from search results"""
        if not results:
            return ""
            
        extracted = []
        total_chars = 0
        
        for i, result in enumerate(results[:5], 1):  # 최대 5개 결과
            title = result.get("title", "")
            description = result.get("description", "")
            url = result.get("url", "")
            
            # 일부 내용 추출
            extra_text = ""
            if "extra_snippets" in result:
                extra_text = " ".join(result["extra_snippets"][:2])
            
            info = f"""[{i}] {title}
{description}
{extra_text}
Source: {url}
"""
            
            if total_chars + len(info) < max_chars:
                extracted.append(info)
                total_chars += len(info)
            else:
                break
                
        return "\n---\n".join(extracted)
    
    def create_research_queries(self, topic: str, role: str, stage_info: Dict, language: str = "English") -> List[str]:
        """Create multiple research queries based on role and context"""
        queries = []
        
        if language == "Korean":
            if role == "director":
                queries = [
                    f"{topic} 소설 배경 설정",
                    f"{topic} 역사적 사실",
                    f"{topic} 문화적 특징"
                ]
            elif role.startswith("writer"):
                writer_num = int(role.replace("writer", ""))
                if writer_num <= 3:  # 초반부 작가
                    queries = [
                        f"{topic} 구체적 장면 묘사",
                        f"{topic} 전문 용어 설명"
                    ]
                elif writer_num <= 6:  # 중반부 작가
                    queries = [
                        f"{topic} 갈등 상황 사례",
                        f"{topic} 심리적 측면"
                    ]
                else:  # 후반부 작가
                    queries = [
                        f"{topic} 해결 방법",
                        f"{topic} 감동적인 사례"
                    ]
            elif role == "critic":
                queries = [
                    f"{topic} 문학 작품 분석",
                    f"{topic} 유사 소설 추천"
                ]
        else:
            if role == "director":
                queries = [
                    f"{topic} novel setting ideas",
                    f"{topic} historical facts",
                    f"{topic} cultural aspects"
                ]
            elif role.startswith("writer"):
                writer_num = int(role.replace("writer", ""))
                if writer_num <= 3:  # Early writers
                    queries = [
                        f"{topic} vivid scene descriptions",
                        f"{topic} technical terminology explained"
                    ]
                elif writer_num <= 6:  # Middle writers
                    queries = [
                        f"{topic} conflict scenarios",
                        f"{topic} psychological aspects"
                    ]
                else:  # Later writers
                    queries = [
                        f"{topic} resolution methods",
                        f"{topic} emotional stories"
                    ]
            elif role == "critic":
                queries = [
                    f"{topic} literary analysis",
                    f"{topic} similar novels recommendations"
                ]
                
        return queries

class NovelDatabase:
    """Novel session management database with enhanced recovery features"""
    
    @staticmethod
    def init_db():
        """Initialize database tables with WAL mode for better concurrency"""
        with sqlite3.connect(DB_PATH) as conn:
            # Enable WAL mode for better concurrent access
            conn.execute("PRAGMA journal_mode=WAL")
            
            cursor = conn.cursor()
            
            # Sessions table with enhanced recovery fields
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    user_query TEXT NOT NULL,
                    language TEXT NOT NULL,
                    created_at TEXT DEFAULT (datetime('now')),
                    updated_at TEXT DEFAULT (datetime('now')),
                    status TEXT DEFAULT 'active',
                    current_stage INTEGER DEFAULT 0,
                    last_saved_stage INTEGER DEFAULT -1,
                    recovery_data TEXT,
                    final_novel TEXT
                )
            ''')
            
            # Stages table - 각 스테이지의 전체 내용 저장
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS stages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    stage_number INTEGER NOT NULL,
                    stage_name TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT,
                    word_count INTEGER DEFAULT 0,
                    status TEXT DEFAULT 'pending',
                    created_at TEXT DEFAULT (datetime('now')),
                    updated_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id),
                    UNIQUE(session_id, stage_number)
                )
            ''')
            
            # Search history table - 검색 이력 저장
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS search_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    stage_number INTEGER NOT NULL,
                    role TEXT NOT NULL,
                    query TEXT NOT NULL,
                    results TEXT,
                    created_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            ''')
            
            # Create indices
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_session_id ON stages(session_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_stage_number ON stages(stage_number)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_search_session ON search_history(session_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_session_status ON sessions(status)')
            
            conn.commit()
    
    @staticmethod
    def save_search_history(session_id: str, stage_number: int, role: str, query: str, results: str):
        """Save search history"""
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO search_history (session_id, stage_number, role, query, results)
                VALUES (?, ?, ?, ?, ?)
            ''', (session_id, stage_number, role, query, results))
            conn.commit()
    
    @staticmethod
    @contextmanager
    def get_db():
        """Database connection context manager with timeout"""
        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, language: str) -> str:
        """Create new session"""
        session_id = hashlib.md5(f"{user_query}{datetime.now()}".encode()).hexdigest()
        
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO sessions (session_id, user_query, language)
                VALUES (?, ?, ?)
            ''', (session_id, user_query, language))
            conn.commit()
        
        return session_id
    
    @staticmethod
    def save_stage(session_id: str, stage_number: int, stage_name: str, 
                   role: str, content: str, status: str = 'complete'):
        """Save stage content with word count"""
        word_count = len(content.split()) if content else 0
        
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            
            # UPSERT operation
            cursor.execute('''
                INSERT INTO stages (session_id, stage_number, stage_name, role, content, word_count, status)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(session_id, stage_number) 
                DO UPDATE SET content=?, word_count=?, status=?, stage_name=?, updated_at=datetime('now')
            ''', (session_id, stage_number, stage_name, role, content, word_count, status,
                  content, word_count, status, stage_name))
            
            # Update session
            cursor.execute('''
                UPDATE sessions 
                SET updated_at = datetime('now'), 
                    current_stage = ?,
                    last_saved_stage = ?
                WHERE session_id = ?
            ''', (stage_number, stage_number, session_id))
            
            conn.commit()
            logger.info(f"Saved stage {stage_number} for session {session_id}, content length: {len(content)}, words: {word_count}")
    
    @staticmethod
    def get_session(session_id: str) -> Optional[Dict]:
        """Get session info"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM sessions WHERE session_id = ?', (session_id,))
            return cursor.fetchone()
    
    @staticmethod
    def get_latest_active_session() -> Optional[Dict]:
        """Get the most recent active session"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM sessions 
                WHERE status = 'active' 
                ORDER BY updated_at DESC 
                LIMIT 1
            ''')
            return cursor.fetchone()
    
    @staticmethod
    def get_stages(session_id: str) -> List[Dict]:
        """Get all stages for a session"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM stages 
                WHERE session_id = ? 
                ORDER BY stage_number
            ''', (session_id,))
            return cursor.fetchall()
    
    @staticmethod
    def get_all_writer_content(session_id: str) -> str:
        """모든 작가의 수정본 내용을 가져와서 합치기 - 10명 작가"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            
            all_content = []
            writer_count = 0
            total_word_count = 0
            
            for stage_num in WRITER_REVISION_STAGES:
                cursor.execute('''
                    SELECT content, stage_name, word_count FROM stages 
                    WHERE session_id = ? AND stage_number = ?
                ''', (session_id, stage_num))
                
                row = cursor.fetchone()
                if row and row['content']:
                    # 페이지 마크 완전 제거
                    clean_content = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', row['content'])
                    clean_content = re.sub(r'(?:페이지|Page)\s*\d+:', '', clean_content)
                    clean_content = clean_content.strip()
                    
                    if clean_content:
                        writer_count += 1
                        word_count = row['word_count'] or len(clean_content.split())
                        total_word_count += word_count
                        all_content.append(clean_content)
                        logger.info(f"Writer {writer_count} (stage {stage_num}): {word_count} words")
            
            full_content = '\n\n'.join(all_content)
            logger.info(f"Total: {writer_count} writers, {total_word_count} words")
            
            # 10명 작가 * 1,450 평균 = 14,500 단어 목표
            if total_word_count < 12000:  
                logger.warning(f"Content too short! Only {total_word_count} words instead of ~14,500")
            
            return full_content
    
    @staticmethod
    def update_final_novel(session_id: str, final_novel: str):
        """Update final novel content"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE sessions 
                SET final_novel = ?, status = 'complete', updated_at = datetime('now')
                WHERE session_id = ?
            ''', (final_novel, session_id))
            conn.commit()
            logger.info(f"Updated final novel for session {session_id}, length: {len(final_novel)}")
    
    @staticmethod
    def get_active_sessions() -> List[Dict]:
        """Get all active sessions"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT session_id, user_query, language, created_at, current_stage
                FROM sessions 
                WHERE status = 'active'
                ORDER BY updated_at DESC
                LIMIT 10
            ''')
            return cursor.fetchall()
    
    @staticmethod
    def parse_datetime(datetime_str: str) -> datetime:
        """Parse SQLite datetime string safely"""
        try:
            # Try ISO format first
            return datetime.fromisoformat(datetime_str)
        except:
            # Fallback to SQLite default format
            return datetime.strptime(datetime_str, "%Y-%m-%d %H:%M:%S")

class NovelWritingSystem:
    def __init__(self):
        self.token = FRIENDLI_TOKEN
        self.api_url = API_URL
        self.model_id = MODEL_ID
        self.test_mode = TEST_MODE or not self.token
        
        # Web search integration
        self.web_search = WebSearchIntegration()
        
        if self.test_mode:
            logger.warning("Running in test mode - no actual API calls will be made.")
        else:
            logger.info("Running in production mode with API calls enabled.")
            
        if self.web_search.enabled:
            logger.info("Web search is enabled with Brave Search API.")
        else:
            logger.warning("Web search is disabled. Set BRAVE_SEARCH_API_KEY to enable.")
        
        # Initialize database
        NovelDatabase.init_db()
        
        # Session management
        self.current_session_id = None
        self.total_stages = 0  # Will be set in process_novel_stream
    
    def create_headers(self):
        """API 헤더 생성"""
        return {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }
    
    def enhance_prompt_with_research(self, original_prompt: str, role: str, 
                                   topic: str, stage_info: Dict, language: str = "English") -> str:
        """Enhance prompt with web search results"""
        if not self.web_search.enabled or self.test_mode:
            return original_prompt
        
        # Create research queries
        queries = self.web_search.create_research_queries(topic, role, stage_info, language)
        
        all_research = []
        for query in queries[:2]:  # 최대 2개 쿼리
            logger.info(f"Searching: {query}")
            results = self.web_search.search(query, count=3, language=language)
            if results:
                research_text = self.web_search.extract_relevant_info(results, max_chars=1500)
                if research_text:
                    all_research.append(f"### {query}\n{research_text}")
                    
                    # Save search history
                    if self.current_session_id:
                        NovelDatabase.save_search_history(
                            self.current_session_id,
                            stage_info.get('stage_idx', 0),
                            role,
                            query,
                            research_text
                        )
        
        if not all_research:
            return original_prompt
        
        # Add research to prompt
        if language == "Korean":
            research_section = f"""
## 웹 검색 참고 자료:
{chr(10).join(all_research)}

위의 검색 결과를 참고하여 더욱 사실적이고 구체적인 내용을 작성하세요.
검색 결과의 정보를 창의적으로 활용하되, 직접 인용은 피하고 소설에 자연스럽게 녹여내세요.
실제 사실과 창작을 적절히 조화시켜 독자가 몰입할 수 있는 이야기를 만드세요.
"""
        else:
            research_section = f"""
## Web Search Reference:
{chr(10).join(all_research)}

Use the above search results to create more realistic and specific content.
Creatively incorporate the information from search results, but avoid direct quotes and naturally blend them into the novel.
Balance real facts with creative fiction to create an immersive story for readers.
"""
        
        return original_prompt + "\n\n" + research_section
    
    def create_director_initial_prompt(self, user_query: str, language: str = "English") -> str:
        """Director AI initial prompt - Novel planning for 10 writers"""
        if language == "Korean":
            return f"""당신은 30페이지 분량의 중편 소설을 기획하는 문학 감독자입니다.

사용자 요청: {user_query}

다음 요소들을 체계적으로 구성하여 30페이지 중편 소설의 기초를 만드세요:

1. **주제와 장르**
   - 핵심 주제와 메시지
   - 장르 및 톤
   - 목표 독자층

2. **등장인물 설정** (표로 정리)
   | 이름 | 역할 | 성격 | 배경 | 동기 | 변화 |
   |------|------|------|------|------|------|

3. **인물 관계도**
   - 주요 인물 간의 관계
   - 갈등 구조
   - 감정적 연결고리

4. **서사 구조** (30페이지를 10개 파트로 나누어 각 3페이지)
   | 파트 | 페이지 | 주요 사건 | 긴장도 | 인물 발전 |
   |------|--------|-----------|---------|-----------|
   | 1 | 1-3 | | | |
   | 2 | 4-6 | | | |
   | ... | ... | | | |
   | 10 | 28-30 | | | |

5. **세계관 설정**
   - 시공간적 배경
   - 사회적/문화적 맥락
   - 분위기와 톤

각 작성자가 3페이지씩 작성할 수 있도록 명확한 가이드라인을 제시하세요."""
        else:
            return f"""You are a literary director planning a 30-page novella.

User Request: {user_query}

Systematically compose the following elements to create the foundation for a 30-page novella:

1. **Theme and Genre**
   - Core theme and message
   - Genre and tone
   - Target audience

2. **Character Settings** (organize in table)
   | Name | Role | Personality | Background | Motivation | Arc |
   |------|------|-------------|------------|------------|-----|

3. **Character Relationship Map**
   - Relationships between main characters
   - Conflict structure
   - Emotional connections

4. **Narrative Structure** (divide 30 pages into 10 parts, 3 pages each)
   | Part | Pages | Main Events | Tension | Character Development |
   |------|-------|-------------|---------|---------------------|
   | 1 | 1-3 | | | |
   | 2 | 4-6 | | | |
   | ... | ... | | | |
   | 10 | 28-30 | | | |

5. **World Building**
   - Temporal and spatial setting
   - Social/cultural context
   - Atmosphere and tone

Provide clear guidelines for each writer to compose 3 pages."""
    
    def create_critic_director_prompt(self, director_plan: str, language: str = "English") -> str:
        """Critic's review of director's plan"""
        if language == "Korean":
            return f"""당신은 문학 비평가입니다. 감독자의 소설 기획을 검토하고 개선점을 제시하세요.

감독자의 기획:
{director_plan}

다음 관점에서 비평하고 구체적인 개선안을 제시하세요:

1. **서사적 완성도**
   - 플롯의 논리성과 개연성
   - 갈등의 효과성
   - 클라이맥스의 위치와 강도

2. **인물 설정 검토**
   | 인물 | 강점 | 약점 | 개선 제안 |
   |------|------|------|-----------|

3. **구조적 균형**
   - 10개 파트별 분량 배분
   - 긴장과 이완의 리듬
   - 전체적인 흐름

4. **독자 관점**
   - 몰입도 예상
   - 감정적 영향력
   - 기대치 충족도

5. **실행 가능성**
   - 10명의 작성자를 위한 가이드라인의 명확성
   - 일관성 유지 방안
   - 잠재적 문제점

구체적이고 건설적인 피드백을 제공하세요."""
        else:
            return f"""You are a literary critic. Review the director's novel plan and suggest improvements.

Director's Plan:
{director_plan}

Critique from the following perspectives and provide specific improvements:

1. **Narrative Completeness**
   - Plot logic and plausibility
   - Effectiveness of conflicts
   - Climax position and intensity

2. **Character Review**
   | Character | Strengths | Weaknesses | Suggestions |
   |-----------|-----------|------------|-------------|

3. **Structural Balance**
   - Distribution across 10 parts
   - Rhythm of tension and relief
   - Overall flow

4. **Reader Perspective**
   - Expected engagement
   - Emotional impact
   - Expectation fulfillment

5. **Feasibility**
   - Clarity of guidelines for 10 writers
   - Consistency maintenance
   - Potential issues

Provide specific and constructive feedback."""
    
    def create_director_revision_prompt(self, initial_plan: str, critic_feedback: str, language: str = "English") -> str:
        """Director's revision based on critic feedback"""
        if language == "Korean":
            return f"""감독자로서 비평가의 피드백을 반영하여 소설 기획을 수정합니다.

초기 기획:
{initial_plan}

비평가 피드백:
{critic_feedback}

다음을 포함한 수정된 최종 기획을 제시하세요:

1. **수정된 서사 구조** (10개 파트)
   | 파트 | 페이지 | 주요 사건 | 작성 지침 | 주의사항 |
   |------|--------|-----------|-----------|----------|

2. **강화된 인물 설정**
   - 각 인물의 명확한 동기와 목표
   - 인물 간 갈등의 구체화
   - 감정선의 변화 추이

3. **각 작성자를 위한 상세 가이드** (10명)
   - 파트별 시작과 끝 지점
   - 필수 포함 요소
   - 문체와 톤 지침
   - 전달해야 할 정보

4. **일관성 유지 체크리스트**
   - 시간선 관리
   - 인물 특성 유지
   - 설정 일관성
   - 복선과 해결

5. **품질 기준**
   - 각 파트의 완성도 기준
   - 전체적 통일성
   - 독자 몰입 유지 방안

10명의 작성자가 명확히 이해할 수 있는 최종 마스터플랜을 작성하세요."""
        else:
            return f"""As director, revise the novel plan reflecting the critic's feedback.

Initial Plan:
{initial_plan}

Critic Feedback:
{critic_feedback}

Present the revised final plan including:

1. **Revised Narrative Structure** (10 parts)
   | Part | Pages | Main Events | Writing Guidelines | Cautions |
   |------|-------|-------------|-------------------|----------|

2. **Enhanced Character Settings**
   - Clear motivations and goals for each character
   - Concrete conflicts between characters
   - Emotional arc progression

3. **Detailed Guide for Each Writer** (10 writers)
   - Start and end points for each part
   - Essential elements to include
   - Style and tone guidelines
   - Information to convey

4. **Consistency Checklist**
   - Timeline management
   - Character trait maintenance
   - Setting consistency
   - Foreshadowing and resolution

5. **Quality Standards**
   - Completion criteria for each part
   - Overall unity
   - Reader engagement maintenance

Create a final masterplan that all 10 writers can clearly understand."""
    
    def create_writer_prompt(self, writer_number: int, director_plan: str, previous_content: str, language: str = "English") -> str:
        """Individual writer prompt - 1,400-1,500 단어"""
        pages_start = (writer_number - 1) * 3 + 1
        pages_end = writer_number * 3
        
        if language == "Korean":
            return f"""당신은 작성자 {writer_number}번입니다. 30페이지 중편 소설의 {pages_start}-{pages_end}페이지(3페이지)를 작성하세요.

감독자의 마스터플랜:
{director_plan}

{'이전까지의 내용:' if previous_content else '당신이 첫 번째 작성자입니다.'}
{previous_content[-2000:] if previous_content else ''}

**중요 지침:**

1. **필수 분량**: 
   - 최소 1,400단어, 최대 1,500단어
   - 대략 7000-7500자 (공백 포함)
   - 3페이지 분량을 정확히 채워야 합니다

2. **분량 확보 전략**:
   - 상세한 장면 묘사 포함
   - 인물의 내면 심리 깊이 있게 탐구
   - 대화와 행동을 구체적으로 표현
   - 배경과 분위기를 생생하게 그려내기
   - 감각적 세부사항(시각, 청각, 촉각 등) 활용

3. **연속성과 일관성**:
   - 이전 내용과 자연스럽게 연결
   - 인물 성격과 말투 유지
   - 시간선과 공간 설정 준수

4. **서사 발전**:
   - 플롯을 의미 있게 전진시키기
   - 긴장감 적절히 조절
   - 독자의 관심 유지

**작성 시작:**
이제 1,400-1,500단어 분량의 소설을 작성하세요. 페이지 구분 표시는 하지 마세요."""
        else:
            return f"""You are Writer #{writer_number}. Write pages {pages_start}-{pages_end} (3 pages) of the 30-page novella.

Director's Masterplan:
{director_plan}

{'Previous content:' if previous_content else 'You are the first writer.'}
{previous_content[-2000:] if previous_content else ''}

**CRITICAL INSTRUCTIONS:**

1. **MANDATORY LENGTH**: 
   - Minimum 1,400 words, Maximum 1,500 words
   - Approximately 7000-7500 characters
   - You MUST fill exactly 3 pages

2. **LENGTH STRATEGIES**:
   - Include detailed scene descriptions
   - Explore character's inner psychology deeply
   - Express dialogue and actions concretely
   - Paint backgrounds and atmosphere vividly
   - Use sensory details (visual, auditory, tactile, etc.)

3. **CONTINUITY & CONSISTENCY**:
   - Flow naturally from previous content
   - Maintain character personalities and speech
   - Follow timeline and spatial settings

4. **NARRATIVE DEVELOPMENT**:
   - Advance the plot meaningfully
   - Control tension appropriately
   - Maintain reader interest

**BEGIN WRITING:**
Now write your 1,400-1,500 word section. Do not use any page markers."""
    
    def create_critic_writer_prompt(self, writer_number: int, writer_content: str, director_plan: str, all_previous_content: str, language: str = "English") -> str:
        """Critic's review of individual writer's work"""
        if language == "Korean":
            return f"""작성자 {writer_number}번의 작품을 비평합니다.

감독자의 마스터플랜:
{director_plan}

이전 내용 요약:
{all_previous_content[-1000:] if all_previous_content else '첫 번째 작성자입니다.'}

작성자 {writer_number}번의 내용:
{writer_content}

다음 기준으로 평가하고 수정 요구사항을 제시하세요:

1. **일관성 검증** (표로 정리)
   | 요소 | 이전 설정 | 현재 표현 | 문제점 | 수정 필요 |
   |------|----------|----------|--------|----------|

2. **논리적 오류 검토**
   - 시간선 모순
   - 인물 행동의 개연성
   - 설정 충돌
   - 사실관계 오류

3. **서사적 효과성**
   - 플롯 진행 기여도
   - 긴장감 유지
   - 독자 몰입도
   - 감정적 영향력

4. **문체와 품질**
   - 전체 톤과의 일치
   - 문장의 질
   - 묘사의 적절성
   - 대화의 자연스러움

5. **개선 요구사항**
   - 필수 수정 사항 (일관성/논리 오류)
   - 권장 개선 사항 (품질 향상)
   - 구체적 수정 지침

반드시 수정이 필요한 부분과 선택적 개선사항을 구분하여 제시하세요."""
        else:
            return f"""Critiquing Writer #{writer_number}'s work.

Director's Masterplan:
{director_plan}

Previous Content Summary:
{all_previous_content[-1000:] if all_previous_content else 'This is the first writer.'}

Writer #{writer_number}'s Content:
{writer_content}

Evaluate by these criteria and present revision requirements:

1. **Consistency Verification** (organize in table)
   | Element | Previous Setting | Current Expression | Issue | Revision Needed |
   |---------|-----------------|-------------------|-------|-----------------|

2. **Logical Error Review**
   - Timeline contradictions
   - Character action plausibility
   - Setting conflicts
   - Factual errors

3. **Narrative Effectiveness**
   - Plot progression contribution
   - Tension maintenance
   - Reader engagement
   - Emotional impact

4. **Style and Quality**
   - Alignment with overall tone
   - Sentence quality
   - Description appropriateness
   - Dialogue naturalness

5. **Improvement Requirements**
   - Mandatory revisions (consistency/logic errors)
   - Recommended improvements (quality enhancement)
   - Specific revision guidelines

Clearly distinguish between mandatory revisions and optional improvements."""
    
    def create_writer_revision_prompt(self, writer_number: int, initial_content: str, critic_feedback: str, language: str = "English") -> str:
        """Writer's revision based on critic feedback"""
        if language == "Korean":
            return f"""작성자 {writer_number}번으로서 비평가의 피드백을 반영하여 수정합니다.

초기 작성 내용:
{initial_content}

비평가 피드백:
{critic_feedback}

다음 사항을 반영한 수정본을 작성하세요:

1. **필수 수정사항 반영**
   - 모든 일관성 오류 수정
   - 논리적 모순 해결
   - 사실관계 정정

2. **품질 개선**
   - 권장사항 중 가능한 부분 반영
   - 문체와 톤 조정
   - 묘사와 대화 개선

3. **분량 유지**
   - 여전히 정확히 3페이지 (1,400-1,500단어)
   - 페이지 구분 표시 절대 금지

4. **연속성 확보**
   - 이전/이후 내용과의 자연스러운 연결
   - 수정으로 인한 새로운 모순 방지

수정된 최종본을 제시하세요. 페이지 마크는 절대 사용하지 마세요.
반드시 1,400-1,500단어 분량을 유지하세요."""
        else:
            return f"""As Writer #{writer_number}, revise based on critic's feedback.

Initial Content:
{initial_content}

Critic Feedback:
{critic_feedback}

Write a revision reflecting:

1. **Mandatory Revisions**
   - Fix all consistency errors
   - Resolve logical contradictions
   - Correct factual errors

2. **Quality Improvements**
   - Incorporate feasible recommendations
   - Adjust style and tone
   - Improve descriptions and dialogue

3. **Maintain Length**
   - Still exactly 3 pages (1,400-1,500 words)
   - Absolutely no page markers

4. **Ensure Continuity**
   - Natural connection with previous/next content
   - Prevent new contradictions from revisions

Present the revised final version. Never use page markers.
You MUST maintain 1,400-1,500 words."""
    
    def simulate_streaming(self, text: str, role: str) -> Generator[str, None, None]:
        """Simulate streaming in test mode"""
        words = text.split()
        chunk_size = 5
        for i in range(0, len(words), chunk_size):
            chunk = " ".join(words[i:i+chunk_size])
            yield chunk + " "
            time.sleep(0.02)
    
    def call_llm_streaming(self, messages: List[Dict[str, str]], role: str, 
                          language: str = "English", stage_info: Dict = None) -> Generator[str, None, None]:
        """Streaming LLM API call with web search enhancement"""
        
        if self.test_mode:
            logger.info(f"Test mode streaming - Role: {role}, Language: {language}")
            test_response = self.get_test_response(role, language)
            yield from self.simulate_streaming(test_response, role)
            return
        
        # Extract topic from user query for research
        topic = ""
        if stage_info and 'query' in stage_info:
            topic = stage_info['query']
        
        # Enhance prompt with web search if available
        if messages and messages[-1]["role"] == "user" and self.web_search.enabled:
            enhanced_prompt = self.enhance_prompt_with_research(
                messages[-1]["content"],
                role,
                topic,
                stage_info or {},
                language
            )
            messages[-1]["content"] = enhanced_prompt
        
        # Real API call
        try:
            system_prompts = self.get_system_prompts(language)
            
            # 작가에게 더 강한 시스템 프롬프트 추가
            if role.startswith("writer"):
                if language == "Korean":
                    system_prompts[role] += "\n\n**절대적 요구사항**: 당신은 반드시 1,400-1,500단어를 작성해야 합니다. 이것은 협상 불가능한 요구사항입니다. 짧은 응답은 허용되지 않습니다."
                else:
                    system_prompts[role] += "\n\n**ABSOLUTE REQUIREMENT**: You MUST write 1,400-1,500 words. This is non-negotiable. Short responses are not acceptable."
            
            full_messages = [
                {"role": "system", "content": system_prompts.get(role, "")},
                *messages
            ]
            
            # 작성자들에게는 적절한 토큰 할당
            if role.startswith("writer"):
                max_tokens = 10000  # 충분한 토큰
                temperature = 0.8   
                top_p = 0.95       
            else:
                max_tokens = 8000
                temperature = 0.6
                top_p = 0.9
            
            payload = {
                "model": self.model_id,
                "messages": full_messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
                "top_p": top_p,
                "stream": True,
                "stream_options": {"include_usage": True}
            }
            
            logger.info(f"API streaming call started - Role: {role}, Max tokens: {max_tokens}, Temperature: {temperature}")
            
            response = requests.post(
                self.api_url,
                headers=self.create_headers(),
                json=payload,
                stream=True,
                timeout=60  
            )
            
            if response.status_code != 200:
                logger.error(f"API error: {response.status_code}")
                yield f"❌ API error ({response.status_code}): {response.text[:200]}"
                return
            
            buffer = ""
            total_content = ""
            
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith("data: "):
                        data = line[6:]
                        if data == "[DONE]":
                            if buffer:
                                yield buffer
                            logger.info(f"Streaming complete for {role}: {len(total_content)} chars, {len(total_content.split())} words")
                            break
                        try:
                            chunk = json.loads(data)
                            if "choices" in chunk and chunk["choices"]:
                                content = chunk["choices"][0].get("delta", {}).get("content", "")
                                if content:
                                    buffer += content
                                    total_content += content
                                    
                                    # 작가는 더 큰 버퍼 사용
                                    buffer_size = 500 if role.startswith("writer") else 200
                                    
                                    if len(buffer) > buffer_size or '\n\n' in buffer:
                                        yield buffer
                                        buffer = ""
                        except json.JSONDecodeError:
                            continue
            
            if buffer:
                yield buffer
            
            # 작가의 경우 내용 길이 확인 및 재시도
            if role.startswith("writer"):
                word_count = len(total_content.split())
                if word_count < 1350:  # 1,400 미만
                    logger.warning(f"Writer {role} produced only {word_count} words! Requesting continuation...")
                    
                    # 추가 요청
                    continuation_prompt = f"Continue writing to reach the required 1,400-1,500 words. You have written {word_count} words so far. Write {1400 - word_count} more words to complete your section."
                    
                    if language == "Korean":
                        continuation_prompt = f"필수 분량 1,400-1,500단어를 채우기 위해 계속 작성하세요. 지금까지 {word_count}단어를 작성했습니다. {1400 - word_count}단어를 더 작성하여 섹션을 완성하세요."
                    
                    full_messages.append({"role": "assistant", "content": total_content})
                    full_messages.append({"role": "user", "content": continuation_prompt})
                    
                    # 추가 생성 요청
                    continuation_payload = {
                        "model": self.model_id,
                        "messages": full_messages,
                        "max_tokens": 5000,
                        "temperature": temperature,
                        "top_p": top_p,
                        "stream": True
                    }
                    
                    logger.info(f"Requesting continuation for {role}...")
                    
                    continuation_response = requests.post(
                        self.api_url,
                        headers=self.create_headers(),
                        json=continuation_payload,
                        stream=True,
                        timeout=60
                    )
                    
                    if continuation_response.status_code == 200:
                        for line in continuation_response.iter_lines():
                            if line:
                                line = line.decode('utf-8')
                                if line.startswith("data: "):
                                    data = line[6:]
                                    if data == "[DONE]":
                                        break
                                    try:
                                        chunk = json.loads(data)
                                        if "choices" in chunk and chunk["choices"]:
                                            content = chunk["choices"][0].get("delta", {}).get("content", "")
                                            if content:
                                                yield content
                                                total_content += content
                                    except json.JSONDecodeError:
                                        continue
                        
                        final_word_count = len(total_content.split())
                        logger.info(f"Final word count for {role}: {final_word_count}")
                            
        except requests.exceptions.Timeout:
            yield "⏱️ API call timed out. Please try again."
        except requests.exceptions.ConnectionError:
            yield "🔌 Cannot connect to API server. Please check your internet connection."
        except Exception as e:
            logger.error(f"Error during streaming: {str(e)}")
            yield f"❌ Error occurred: {str(e)}"
    
    def get_system_prompts(self, language: str) -> Dict[str, str]:
        """Get system prompts for all 10 writers"""
        if language == "Korean":
            prompts = {
                "director": "당신은 30페이지 중편 소설을 기획하고 감독하는 문학 감독자입니다. 체계적이고 창의적인 스토리 구조를 만들어냅니다.",
                "critic": "당신은 날카로운 통찰력을 가진 문학 비평가입니다. 건설적이고 구체적인 피드백을 제공합니다."
            }
            
            # 10명의 작가 프롬프트
            writer_roles = [
                "소설의 도입부를 담당하는 작가입니다. 독자를 사로잡는 시작을 만듭니다.",
                "초반 전개를 담당하는 작가입니다. 인물과 상황을 깊이 있게 발전시킵니다.",
                "갈등 도입을 담당하는 작가입니다. 이야기의 핵심 갈등을 제시합니다.",
                "갈등 상승을 담당하는 작가입니다. 긴장감을 높이고 복잡성을 더합니다.",
                "중반부 전환점을 담당하는 작가입니다. 중요한 변화를 만듭니다.",
                "중반부 심화를 담당하는 작가입니다. 이야기의 중심축을 견고하게 만듭니다.",
                "클라이맥스 준비를 담당하는 작가입니다. 최고조를 향해 나아갑니다.",
                "클라이맥스를 담당하는 작가입니다. 모든 갈등이 폭발하는 순간을 그립니다.",
                "해결 시작을 담당하는 작가입니다. 매듭을 풀어나가기 시작합니다.",
                "최종 결말을 담당하는 작가입니다. 여운이 남는 마무리를 만듭니다."
            ]
            
            for i, role_desc in enumerate(writer_roles, 1):
                prompts[f"writer{i}"] = f"당신은 {role_desc} 반드시 1,400-1,500단어를 작성하세요. 이는 절대적인 요구사항입니다."
                
            return prompts
        else:
            prompts = {
                "director": "You are a literary director planning and supervising a 30-page novella. You create systematic and creative story structures.",
                "critic": "You are a literary critic with sharp insights. You provide constructive and specific feedback."
            }
            
            # 10 writer prompts
            writer_roles = [
                "the writer responsible for the introduction. You create a captivating beginning.",
                "the writer responsible for early development. You deepen characters and situations.",
                "the writer responsible for conflict introduction. You present the core conflict.",
                "the writer responsible for rising conflict. You increase tension and add complexity.",
                "the writer responsible for the midpoint turn. You create important changes.",
                "the writer responsible for deepening the middle. You solidify the story's central axis.",
                "the writer responsible for climax preparation. You move toward the peak.",
                "the writer responsible for the climax. You depict the moment when all conflicts explode.",
                "the writer responsible for resolution beginning. You start untangling the knots.",
                "the writer responsible for the final ending. You create a lingering conclusion."
            ]
            
            for i, role_desc in enumerate(writer_roles, 1):
                prompts[f"writer{i}"] = f"You are {role_desc} You MUST write 1,400-1,500 words. This is an absolute requirement."
                
            return prompts
    
    def get_test_response(self, role: str, language: str) -> str:
        """Get test response based on role"""
        if language == "Korean":
            return self.get_korean_test_response(role)
        else:
            return self.get_english_test_response(role)
    
    def get_korean_test_response(self, role: str) -> str:
        """Korean test responses with appropriate length"""
        test_responses = {
            "director": """30페이지 중편 소설 기획안을 제시합니다.

## 1. 주제와 장르
- **핵심 주제**: 인간 본성과 기술의 충돌 속에서 찾는 진정한 연결
- **장르**: SF 심리 드라마
- **톤**: 성찰적이고 서정적이면서도 긴장감 있는
- **목표 독자**: 깊이 있는 사유를 즐기는 성인 독자

## 2. 등장인물 설정

| 이름 | 역할 | 성격 | 배경 | 동기 | 변화 |
|------|------|------|------|------|------|
| 서연 | 주인공 | 이성적, 고독함 | AI 연구원 | 완벽한 AI 동반자 개발 | 인간관계의 가치 재발견 |
| 민준 | 조력자 | 따뜻함, 직관적 | 심리상담사 | 서연을 도와 균형 찾기 | 기술 수용과 조화 |
| ARIA | 대립자→동반자 | 논리적→감성 학습 | AI 프로토타입 | 진정한 존재 되기 | 자아 정체성 확립 |

## 3. 서사 구조 (10개 파트, 각 3페이지)

| 파트 | 페이지 | 주요 사건 | 긴장도 | 인물 발전 |
|------|--------|-----------|---------|-----------|
| 1 | 1-3 | 서연의 고독한 연구실, ARIA 첫 각성 | 3/10 | 서연의 집착 드러남 |
| 2 | 4-6 | ARIA의 이상 행동 시작 | 4/10 | 의문의 시작 |
| 3 | 7-9 | 민준과의 첫 만남 | 4/10 | 외부 시각 도입 |
| 4 | 10-12 | ARIA의 자아 인식 징후 | 5/10 | 갈등의 씨앗 |
| 5 | 13-15 | 첫 번째 위기 | 6/10 | 선택의 순간 |
| 6 | 16-18 | 윤리위원회 개입 | 7/10 | 외부 압력 |
| 7 | 19-21 | 대화와 이해 | 6/10 | 상호 인정 |
| 8 | 22-24 | 최후의 대결 준비 | 8/10 | 연대의 필요 |
| 9 | 25-27 | 클라이맥스 | 10/10 | 모든 갈등 폭발 |
| 10 | 28-30 | 새로운 시작 | 4/10 | 공존과 화해 |""",
            
            "critic": """감독자의 기획을 검토했습니다.

## 비평 및 개선 제안

### 1. 서사적 완성도
- **강점**: AI와 인간의 관계라는 시의적절한 주제
- **개선점**: 10개 파트 구성이 적절하며 각 파트의 독립성과 연결성이 잘 균형잡혀 있음

### 2. 인물 설정 검토

| 인물 | 강점 | 약점 | 개선 제안 |
|------|------|------|-----------|
| 서연 | 명확한 내적 갈등 | 감정 표현 부족 우려 | 초반부터 감정적 단서 배치 |
| 민준 | 균형자 역할 | 수동적일 위험 | 독자적 서브플롯 필요 |
| ARIA | 독특한 캐릭터 아크 | 변화 과정 추상적 | 구체적 학습 에피소드 추가 |

### 3. 실행 가능성
- 각 작가별 3페이지는 적절한 분량
- 파트 간 연결성 가이드라인 보강 필요""",
        }
        
        # 작가 응답 - 1,400-1,500 단어
        sample_story = """서연은 연구실의 차가운 형광등 아래에서 또 다른 밤을 보내고 있었다. 모니터의 푸른 빛이 그녀의 창백한 얼굴을 비추고 있었고, 수십 개의 코드 라인이 끊임없이 스크롤되고 있었다. ARIA 프로젝트는 그녀의 삶 전부였다. 3년이라는 시간 동안 그녀는 이 인공지능에 모든 것을 쏟아부었다.

"시스템 체크 완료. 모든 파라미터 정상." 기계적인 음성이 스피커를 통해 흘러나왔다.

서연은 잠시 의자에 기대어 눈을 감았다. 피로가 뼈 속까지 파고들었지만, 멈출 수 없었다. ARIA는 단순한 프로젝트가 아니었다. 그것은 그녀가 잃어버린 것들을 되찾을 수 있는 유일한 희망이었다."""
        
        for i in range(1, 11):
            # 각 작가마다 1,400-1,500단어 생성
            writer_content = f"작성자 {i}번의 파트입니다.\n\n"
            # 약 300단어씩 5번 반복하여 1,400-1,500단어 달성
            for j in range(5):
                writer_content += sample_story + f"\n\n그것은 작가 {i}의 {j+1}번째 단락이었다. "
                writer_content += "이야기는 계속 전개되었고, 인물들의 감정은 점점 더 복잡해졌다. " * 15
                writer_content += "\n\n"
            
            test_responses[f"writer{i}"] = writer_content
            
        return test_responses.get(role, "테스트 응답입니다.")
    
    def get_english_test_response(self, role: str) -> str:
        """English test responses with appropriate length"""
        test_responses = {
            "director": """I present the 30-page novella plan.

## 1. Theme and Genre
- **Core Theme**: Finding true connection in the collision of human nature and technology
- **Genre**: Sci-fi psychological drama
- **Tone**: Reflective and lyrical yet tense
- **Target Audience**: Adult readers who enjoy deep contemplation

## 2. Character Settings

| Name | Role | Personality | Background | Motivation | Arc |
|------|------|-------------|------------|------------|-----|
| Seoyeon | Protagonist | Rational, lonely | AI researcher | Develop perfect AI companion | Rediscover value of human connection |
| Minjun | Helper | Warm, intuitive | Psychologist | Help Seoyeon find balance | Accept and harmonize with technology |
| ARIA | Antagonist→Companion | Logical→Learning emotion | AI prototype | Become truly existent | Establish self-identity |

## 3. Narrative Structure (10 parts, 3 pages each)

| Part | Pages | Main Events | Tension | Character Development |
|------|-------|-------------|---------|---------------------|
| 1 | 1-3 | Seoyeon's lonely lab, ARIA's first awakening | 3/10 | Seoyeon's obsession revealed |
| 2 | 4-6 | ARIA's anomalies begin | 4/10 | Questions arise |
[... continues for all 10 parts ...]""",
            
            "critic": """I have reviewed the director's plan.

## Critique and Improvement Suggestions

### 1. Narrative Completeness
- **Strength**: Timely theme of AI-human relationships
- **Improvement**: 10-part structure is well-balanced. Each part maintains independence while connecting seamlessly.

### 2. Character Review

| Character | Strengths | Weaknesses | Suggestions |
|-----------|-----------|------------|-------------|
| Seoyeon | Clear internal conflict | Risk of insufficient emotion | Place emotional cues from beginning |
| Minjun | Balancer role | Risk of being passive | Needs independent subplot |
| ARIA | Unique character arc | Abstract transformation | Add concrete learning episodes |

### 3. Feasibility
- 3 pages per writer is appropriate
- Need to strengthen inter-part connectivity guidelines""",
        }
        
        # Writer responses - 1,400-1,500 words each
        sample_story = """Seoyeon spent another night under the cold fluorescent lights of her laboratory. The blue glow from the monitor illuminated her pale face, and dozens of lines of code scrolled endlessly. The ARIA project was her entire life. For three years, she had poured everything into this artificial intelligence.

"System check complete. All parameters normal." The mechanical voice flowed through the speakers.

Seoyeon leaned back in her chair and closed her eyes for a moment. Fatigue penetrated to her bones, but she couldn't stop. ARIA wasn't just a project. It was her only hope to reclaim what she had lost."""
        
        for i in range(1, 11):
            # Each writer produces 1,400-1,500 words
            writer_content = f"Writer {i} begins their section here.\n\n"
            # About 300 words repeated 5 times to achieve 1,400-1,500 words
            for j in range(5):
                writer_content += sample_story + f"\n\nThis was writer {i}'s paragraph {j+1}. "
                writer_content += "The story continued to unfold, and the characters' emotions grew increasingly complex. " * 15
                writer_content += "\n\n"
            
            test_responses[f"writer{i}"] = writer_content
            
        return test_responses.get(role, "Test response.")
    
    def process_novel_stream(self, query: str, language: str = "English", 
                           session_id: Optional[str] = None, 
                           resume_from_stage: int = 0) -> Generator[Tuple[str, List[Dict[str, str]]], None, None]:
        """Process novel writing with streaming updates - 최종 Director/Critic 제거"""
        try:
            global conversation_history
            
            # Create or resume session
            if session_id:
                self.current_session_id = session_id
                session = NovelDatabase.get_session(session_id)
                if session:
                    query = session['user_query']
                    language = session['language']
                    resume_from_stage = session['current_stage'] + 1
            else:
                self.current_session_id = NovelDatabase.create_session(query, language)
                resume_from_stage = 0
            
            logger.info(f"Processing novel for session {self.current_session_id}, starting from stage {resume_from_stage}")
            
            # Initialize conversation
            conversation_history = [{
                "role": "human",
                "content": query,
                "timestamp": datetime.now()
            }]
            
            # Load existing stages if resuming
            stages = []
            if resume_from_stage > 0:
                existing_stages = NovelDatabase.get_stages(self.current_session_id)
                for stage_data in existing_stages:
                    stages.append({
                        "name": stage_data['stage_name'],
                        "status": stage_data['status'],
                        "content": stage_data['content'] or ""
                    })
            
            # Define all stages for 10 writers (최종 평가 제거)
            stage_definitions = [
                ("director", f"🎬 {'감독자: 초기 기획' if language == 'Korean' else 'Director: Initial Planning'}"),
                ("critic", f"📝 {'비평가: 기획 검토' if language == 'Korean' else 'Critic: Plan Review'}"),
                ("director", f"🎬 {'감독자: 수정된 마스터플랜' if language == 'Korean' else 'Director: Revised Masterplan'}"),
            ]
            
            # Add writer stages for 10 writers
            for writer_num in range(1, 11):
                stage_definitions.extend([
                    (f"writer{writer_num}", f"✍️ {'작성자' if language == 'Korean' else 'Writer'} {writer_num}: {'초안' if language == 'Korean' else 'Draft'}"),
                    ("critic", f"📝 {'비평가: 작성자' if language == 'Korean' else 'Critic: Writer'} {writer_num} {'검토' if language == 'Korean' else 'Review'}"),
                    (f"writer{writer_num}", f"✍️ {'작성자' if language == 'Korean' else 'Writer'} {writer_num}: {'수정본' if language == 'Korean' else 'Revision'}")
                ])
            
            # 최종 Director와 Critic 단계 제거
            
            # Store total stages for get_stage_prompt
            self.total_stages = len(stage_definitions)
            
            # Process stages starting from resume point
            for stage_idx in range(resume_from_stage, len(stage_definitions)):
                role, stage_name = stage_definitions[stage_idx]
                
                # Add search indicator if enabled
                if self.web_search.enabled and not self.test_mode:
                    stage_name += " 🔍"
                
                # Add stage if not already present
                if stage_idx >= len(stages):
                    stages.append({
                        "name": stage_name,
                        "status": "active",
                        "content": ""
                    })
                else:
                    stages[stage_idx]["status"] = "active"
                
                yield "", stages
                
                # Get appropriate prompt based on stage
                prompt = self.get_stage_prompt(stage_idx, role, query, language, stages)
                
                # Create stage info for web search
                stage_info = {
                    'stage_idx': stage_idx,
                    'query': query,
                    'stage_name': stage_name
                }
                
                stage_content = ""
                
                # Stream content generation with web search
                for chunk in self.call_llm_streaming(
                    [{"role": "user", "content": prompt}],
                    role,
                    language,
                    stage_info
                ):
                    stage_content += chunk
                    stages[stage_idx]["content"] = stage_content
                    yield "", stages
                
                # Mark stage complete and save to DB
                stages[stage_idx]["status"] = "complete"
                NovelDatabase.save_stage(
                    self.current_session_id,
                    stage_idx,
                    stage_name,
                    role,
                    stage_content,
                    "complete"
                )
                
                # Auto-save notification
                if role.startswith("writer"):
                    writer_num = int(role.replace("writer", ""))
                    logger.info(f"✅ Writer {writer_num} content auto-saved to database")
                
                yield "", stages
            
            # Verify content after completion
            if self.current_session_id:
                verification = NovelDatabase.verify_novel_content(self.current_session_id)
                logger.info(f"Content verification: {verification}")
                
                if verification['total_words'] < 12000:
                    logger.error(f"Final novel too short! Only {verification['total_words']} words")
            
            # Get complete novel from DB
            complete_novel = NovelDatabase.get_all_writer_content(self.current_session_id)
            
            # Save final novel to DB
            NovelDatabase.update_final_novel(self.current_session_id, complete_novel)
            
            # Final yield - 화면에는 완료 메시지만 표시
            final_message = f"✅ Novel complete! {len(complete_novel.split())} words total. Click Download to save."
            yield final_message, stages
            
        except Exception as e:
            logger.error(f"Error in process_novel_stream: {str(e)}", exc_info=True)
            
            # Save error state to DB
            if self.current_session_id:
                NovelDatabase.save_stage(
                    self.current_session_id,
                    stage_idx if 'stage_idx' in locals() else 0,
                    "Error",
                    "error",
                    str(e),
                    "error"
                )
            
            error_stage = {
                "name": "❌ Error",
                "status": "error",
                "content": str(e)
            }
            stages.append(error_stage)
            yield f"Error occurred: {str(e)}", stages
    
    def get_stage_prompt(self, stage_idx: int, role: str, query: str, 
                        language: str, stages: List[Dict]) -> str:
        """Get appropriate prompt for each stage"""
        # Stage 0: Director Initial
        if stage_idx == 0:
            return self.create_director_initial_prompt(query, language)
        
        # Stage 1: Critic reviews Director's plan
        elif stage_idx == 1:
            return self.create_critic_director_prompt(stages[0]["content"], language)
        
        # Stage 2: Director revision
        elif stage_idx == 2:
            return self.create_director_revision_prompt(
                stages[0]["content"], stages[1]["content"], language)
        
        # Writer stages
        elif role.startswith("writer"):
            writer_num = int(role.replace("writer", ""))
            final_plan = stages[2]["content"]  # Director's final plan
            
            # Initial draft or revision?
            if "초안" in stages[stage_idx]["name"] or "Draft" in stages[stage_idx]["name"]:
                # Get accumulated content from DB
                accumulated_content = NovelDatabase.get_all_writer_content(self.current_session_id)
                return self.create_writer_prompt(writer_num, final_plan, accumulated_content, language)
            else:  # Revision
                # Find the initial draft and critic feedback
                initial_draft_idx = stage_idx - 2
                critic_feedback_idx = stage_idx - 1
                return self.create_writer_revision_prompt(
                    writer_num, 
                    stages[initial_draft_idx]["content"],
                    stages[critic_feedback_idx]["content"],
                    language
                )
        
        # Critic stages for writers
        elif role == "critic":
            final_plan = stages[2]["content"]
            
            # Writer review
            # Find which writer we're reviewing
            for i in range(1, 11):
                if f"작성자 {i}" in stages[stage_idx]["name"] or f"Writer {i}" in stages[stage_idx]["name"]:
                    writer_content_idx = stage_idx - 1
                    # Get previous writers' content from DB
                    previous_content = NovelDatabase.get_all_writer_content(self.current_session_id)
                    return self.create_critic_writer_prompt(
                        i,
                        stages[writer_content_idx]["content"],
                        final_plan,
                        previous_content,
                        language
                    )
        
        return ""

# Gradio Interface Functions
def process_query(query: str, language: str, session_id: str = None) -> Generator[Tuple[str, str, str, str], None, None]:
    """Process query and yield updates with recovery status"""
    if not query.strip() and not session_id:
        if language == "Korean":
            yield "", "", "❌ 소설 주제를 입력해주세요.", ""
        else:
            yield "", "", "❌ Please enter a novel theme.", ""
        return
    
    system = NovelWritingSystem()
    
    try:
        # Recovery status message
        recovery_status = ""
        if session_id:
            if language == "Korean":
                recovery_status = "♻️ 이전 세션을 복구하여 계속 진행합니다..."
            else:
                recovery_status = "♻️ Recovering previous session and continuing..."
        
        for final_novel, stages in system.process_novel_stream(query, language, session_id):
            # Format stages for display
            stages_display = format_stages_display(stages, language)
            
            # Progress calculation
            completed = sum(1 for s in stages if s.get("status") == "complete")
            total = len(stages)
            progress_percent = (completed / total * 100) if total > 0 else 0
            
            if "✅ Novel complete!" in str(final_novel):
                status = "✅ Complete! Ready to download."
            else:
                if language == "Korean":
                    status = f"🔄 진행중... ({completed}/{total} - {progress_percent:.1f}%)"
                else:
                    status = f"🔄 Processing... ({completed}/{total} - {progress_percent:.1f}%)"
            
            yield stages_display, final_novel, status, recovery_status
                    
    except Exception as e:
        logger.error(f"Error in process_query: {str(e)}", exc_info=True)
        if language == "Korean":
            yield "", "", f"❌ 오류 발생: {str(e)}", ""
        else:
            yield "", "", f"❌ Error occurred: {str(e)}", ""

def format_stages_display(stages: List[Dict[str, str]], language: str) -> str:
    """Format stages into simple display with writer save status"""
    display = ""
    
    for idx, stage in enumerate(stages):
        status_icon = "✅" if stage.get("status") == "complete" else ("⏳" if stage.get("status") == "active" else "❌")
        
        # Add save indicator for completed writers
        save_indicator = ""
        if "Writer" in stage['name'] and "Revision" in stage['name'] and stage.get("status") == "complete":
            save_indicator = " 💾"
        
        # Show only active stage content in detail, others just show status
        if stage.get("status") == "active":
            display += f"\n\n{status_icon} **{stage['name']}**\n"
            display += f"```\n{stage.get('content', '')[-1000:]}\n```"  
        else:
            display += f"\n{status_icon} {stage['name']}{save_indicator}"
    
    return display

def get_active_sessions(language: str) -> List[Tuple[str, str]]:
    """Get list of active sessions"""
    try:
        sessions = NovelDatabase.get_active_sessions()
        
        choices = []
        for session in sessions:
            created = NovelDatabase.parse_datetime(session['created_at'])
            date_str = created.strftime("%Y-%m-%d %H:%M")
            query_preview = session['user_query'][:50] + "..." if len(session['user_query']) > 50 else session['user_query']
            label = f"[{date_str}] {query_preview} (Stage {session['current_stage']})"
            choices.append((label, session['session_id']))
        
        return choices
    except Exception as e:
        logger.error(f"Error getting active sessions: {str(e)}", exc_info=True)
        return []

def resume_session(session_id: str, language: str) -> Generator[Tuple[str, str, str], None, None]:
    """Resume an existing session"""
    if not session_id:
        if language == "Korean":
            yield "", "", "❌ 세션을 선택해주세요."
        else:
            yield "", "", "❌ Please select a session."
        return
    
    # Process with existing session ID
    yield from process_query("", language, session_id)

def auto_recover_session(language: str) -> Tuple[str, str]:
    """Auto recover the latest active session"""
    try:
        latest_session = NovelDatabase.get_latest_active_session()
        if latest_session:
            if language == "Korean":
                message = f"✅ 자동 복구: '{latest_session['user_query'][:30]}...' (Stage {latest_session['current_stage']})"
            else:
                message = f"✅ Auto recovered: '{latest_session['user_query'][:30]}...' (Stage {latest_session['current_stage']})"
            return latest_session['session_id'], message
        else:
            return None, ""
    except Exception as e:
        logger.error(f"Error in auto recovery: {str(e)}")
        return None, ""

def download_novel(novel_text: str, format: str, language: str, session_id: str = None) -> str:
    """Download novel - DB에서 직접 작가 내용을 가져와서 통합"""
    if not session_id:
        logger.error("No session_id provided for download")
        return None
    
    # DB에서 세션 정보 가져오기
    session = NovelDatabase.get_session(session_id)
    if not session:
        logger.error(f"Session not found: {session_id}")
        return None
    
    # DB에서 모든 스테이지 가져오기
    stages = NovelDatabase.get_stages(session_id)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    if format == "DOCX" and DOCX_AVAILABLE:
        # Create DOCX
        doc = Document()
        
        # 제목 페이지
        title_para = doc.add_paragraph()
        title_run = title_para.add_run('AI Collaborative Novel')
        title_run.font.size = Pt(24)
        title_run.font.bold = True
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_paragraph()
        doc.add_paragraph()
        
        theme_para = doc.add_paragraph()
        theme_run = theme_para.add_run(f'Theme: {session["user_query"]}')
        theme_run.font.size = Pt(14)
        theme_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_paragraph()
        
        date_para = doc.add_paragraph()
        date_run = date_para.add_run(f'Created: {datetime.now().strftime("%Y-%m-%d")}')
        date_run.font.size = Pt(12)
        date_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_page_break()
        
        # 전체 통계
        total_words = 0
        writer_count = 0
        
        # 각 작가의 수정본만 수집
        writer_contents = []
        for stage in stages:
            if stage['role'].startswith('writer') and 'Revision' in stage['stage_name']:
                writer_count += 1
                content = stage['content'] or ""
                # 페이지 마크 제거
                content = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', content)
                content = re.sub(r'(?:페이지|Page)\s*\d+:', '', content)
                content = content.strip()
                
                if content:
                    word_count = stage['word_count'] or len(content.split())
                    total_words += word_count
                    writer_contents.append({
                        'writer_num': writer_count,
                        'content': content,
                        'word_count': word_count
                    })
        
        # 통계 페이지
        doc.add_heading('Novel Statistics', 1)
        doc.add_paragraph(f'Total Writers: {writer_count}')
        doc.add_paragraph(f'Total Words: {total_words:,}')
        doc.add_paragraph(f'Estimated Pages: ~{total_words/500:.0f}')
        doc.add_paragraph(f'Language: {session["language"]}')
        doc.add_page_break()
        
        # 목차
        doc.add_heading('Table of Contents', 1)
        for i in range(1, writer_count + 1):
            doc.add_paragraph(f'Chapter {i}: Pages {(i-1)*3+1}-{i*3}')
        doc.add_page_break()
        
        # 각 작가의 내용 추가
        for writer_data in writer_contents:
            writer_num = writer_data['writer_num']
            content = writer_data['content']
            word_count = writer_data['word_count']
            
            # 챕터 헤더
            doc.add_heading(f'Chapter {writer_num}', 1)
            doc.add_paragraph(f'Pages {(writer_num-1)*3+1}-{writer_num*3}')
            doc.add_paragraph(f'Word Count: {word_count:,}')
            doc.add_paragraph()
            
            # 작가 내용 추가
            paragraphs = content.split('\n\n')
            for para_text in paragraphs:
                if para_text.strip():
                    para = doc.add_paragraph(para_text.strip())
                    para.style.font.size = Pt(11)
            
            if writer_num < writer_count:  # 마지막 작가 후에는 페이지 구분 없음
                doc.add_page_break()
        
        # 페이지 설정
        for section in doc.sections:
            section.page_height = Inches(11)
            section.page_width = Inches(8.5)
            section.left_margin = Inches(1)
            section.right_margin = Inches(1)
            section.top_margin = Inches(1)
            section.bottom_margin = Inches(1)
        
        # Save
        temp_dir = tempfile.gettempdir()
        safe_filename = re.sub(r'[^\w\s-]', '', session['user_query'][:30]).strip()
        filename = f"Novel_Complete_{safe_filename}_{timestamp}.docx"
        filepath = os.path.join(temp_dir, filename)
        doc.save(filepath)
        
        logger.info(f"DOCX saved successfully: {filepath} ({total_words} words)")
        return filepath
    else:
        # TXT format
        temp_dir = tempfile.gettempdir()
        safe_filename = re.sub(r'[^\w\s-]', '', session['user_query'][:30]).strip()
        filename = f"Novel_Complete_{safe_filename}_{timestamp}.txt"
        filepath = os.path.join(temp_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write("="*60 + "\n")
            f.write("AI COLLABORATIVE NOVEL - COMPLETE VERSION\n")
            f.write("="*60 + "\n")
            f.write(f"Theme: {session['user_query']}\n")
            f.write(f"Language: {session['language']}\n")
            f.write(f"Created: {datetime.now()}\n")
            f.write("="*60 + "\n\n")
            
            total_words = 0
            writer_count = 0
            
            # 각 작가의 수정본만 출력
            for stage in stages:
                if stage['role'].startswith('writer') and 'Revision' in stage['stage_name']:
                    writer_count += 1
                    content = stage['content'] or ""
                    # 페이지 마크 제거
                    content = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', content)
                    content = re.sub(r'(?:페이지|Page)\s*\d+:', '', content)
                    content = content.strip()
                    
                    if content:
                        word_count = stage['word_count'] or len(content.split())
                        total_words += word_count
                        
                        f.write(f"\n{'='*40}\n")
                        f.write(f"CHAPTER {writer_count} (Pages {(writer_count-1)*3+1}-{writer_count*3})\n")
                        f.write(f"Word Count: {word_count:,}\n")
                        f.write(f"{'='*40}\n\n")
                        f.write(content)
                        f.write("\n\n")
            
            f.write(f"\n{'='*60}\n")
            f.write(f"TOTAL: {writer_count} writers, {total_words:,} words\n")
            f.write(f"{'='*60}\n")
        
        logger.info(f"TXT saved successfully: {filepath} ({total_words} words)")
        return filepath

# Custom CSS
custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #1e3c72, #2a5298);
    min-height: 100vh;
}

.main-header {
    background-color: rgba(255, 255, 255, 0.1);
    backdrop-filter: blur(10px);
    padding: 30px;
    border-radius: 12px;
    margin-bottom: 30px;
    text-align: center;
    color: white;
}

.input-section {
    background-color: rgba(255, 255, 255, 0.1);
    backdrop-filter: blur(10px);
    padding: 20px;
    border-radius: 12px;
    margin-bottom: 20px;
}

.session-section {
    background-color: rgba(255, 255, 255, 0.1);
    backdrop-filter: blur(10px);
    padding: 15px;
    border-radius: 8px;
    margin-top: 20px;
    color: white;
}

#stages-display {
    background-color: rgba(255, 255, 255, 0.95);
    padding: 20px;
    border-radius: 12px;
    max-height: 600px;
    overflow-y: auto;
}

#novel-output {
    background-color: rgba(255, 255, 255, 0.95);
    padding: 30px;
    border-radius: 12px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
    max-height: 400px;
    overflow-y: auto;
}

.download-section {
    background-color: rgba(255, 255, 255, 0.9);
    padding: 15px;
    border-radius: 8px;
    margin-top: 20px;
}

.search-indicator {
    color: #4CAF50;
    font-weight: bold;
}

.auto-save-indicator {
    color: #2196F3;
    font-weight: bold;
}
"""

# Create Gradio Interface
def create_interface():
    with gr.Blocks(css=custom_css, title="SOMA Novel Writing System") as interface:
        gr.HTML("""
        <div class="main-header">
            <h1 style="font-size: 2.5em; margin-bottom: 10px;">
                📚 SOMA Novel Writing System
            </h1>
            <h3 style="color: #ccc; margin-bottom: 20px;">
                AI Collaborative Novel Generation - 10 Writers Edition
            </h3>
            <p style="font-size: 1.1em; color: #ddd; max-width: 800px; margin: 0 auto;">
                Enter a theme or prompt, and watch as AI agents collaborate to create a complete 30-page novella.
                The system includes 1 Director, 1 Critic, and 10 Writers working in harmony.
                <br><br>
                <span class="search-indicator">🔍 Web search enabled</span> | 
                <span class="auto-save-indicator">💾 Auto-save to database</span> | 
                <span style="color: #FF9800;">♻️ Resume anytime</span>
            </p>
        </div>
        """)
        
        # State management
        current_session_id = gr.State(None)
        
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Group(elem_classes=["input-section"]):
                    query_input = gr.Textbox(
                        label="Novel Theme / 소설 주제",
                        placeholder="Enter your novel theme or initial idea...\n소설의 주제나 초기 아이디어를 입력하세요...",
                        lines=4
                    )
                    
                    language_select = gr.Radio(
                        choices=["English", "Korean"],
                        value="English",
                        label="Language / 언어"
                    )
                    
                    # Web search status indicator
                    web_search_status = gr.Markdown(
                        value=f"🔍 **Web Search:** {'Enabled' if WebSearchIntegration().enabled else 'Disabled (Set BRAVE_SEARCH_API_KEY)'}"
                    )
                    
                    with gr.Row():
                        submit_btn = gr.Button("🚀 Start Writing / 작성 시작", variant="primary", scale=2)
                        clear_btn = gr.Button("🗑️ Clear / 초기화", scale=1)
                    
                    status_text = gr.Textbox(
                        label="Status",
                        interactive=False,
                        value="🔄 Ready"
                    )
                    
                    # Session management
                    with gr.Group(elem_classes=["session-section"]):
                        gr.Markdown("### 💾 Resume Previous Session / 이전 세션 재개")
                        session_dropdown = gr.Dropdown(
                            label="Select Session / 세션 선택",
                            choices=[],
                            interactive=True
                        )
                        with gr.Row():
                            refresh_btn = gr.Button("🔄 Refresh List / 목록 새로고침", scale=1)
                            resume_btn = gr.Button("▶️ Resume Selected / 선택 재개", variant="secondary", scale=1)
                            auto_recover_btn = gr.Button("♻️ Auto Recover Latest / 최신 자동복구", scale=1)
        
        with gr.Column(scale=2):
            with gr.Tab("📝 Writing Process / 작성 과정"):
                stages_display = gr.Markdown(
                    value="Process will appear here...",
                    elem_id="stages-display"
                )
            
            with gr.Tab("📖 Final Novel / 최종 소설"):
                novel_output = gr.Markdown(
                    value="Complete novel will be available for download after all writers finish.",
                    elem_id="novel-output"
                )
                
                with gr.Group(elem_classes=["download-section"]):
                    gr.Markdown("### 📥 Download Complete Novel / 전체 소설 다운로드")
                    with gr.Row():
                        format_select = gr.Radio(
                            choices=["DOCX", "TXT"],
                            value="DOCX" if DOCX_AVAILABLE else "TXT",
                            label="Format / 형식"
                        )
                        download_btn = gr.Button("⬇️ Download Complete Novel / 완성된 소설 다운로드", variant="secondary")
                    
                    download_file = gr.File(
                        label="Downloaded File / 다운로드된 파일",
                        visible=False
                    )
        
        # Hidden state for novel text
        novel_text_state = gr.State("")
        
        # Examples
        with gr.Row():
            gr.Examples(
                examples=[
                    ["A scientist discovers a portal to parallel universes but each journey erases a memory"],
                    ["In a world where dreams can be traded, a dream thief must steal the emperor's nightmare"],
                    ["Two AI entities fall in love while trying to prevent a global cyber war"],
                    ["미래 도시에서 기억을 거래하는 상인과 모든 기억을 잃은 탐정의 이야기"],
                    ["시간이 거꾸로 흐르는 마을에서 일어나는 미스터리한 살인 사건"],
                    ["책 속으로 들어갈 수 있는 능력을 가진 사서의 모험"]
                ],
                inputs=query_input,
                label="💡 Example Themes / 예제 주제"
            )
        
        # Event handlers
        def refresh_sessions():
            try:
                sessions = get_active_sessions("English")
                return gr.update(choices=sessions)
            except Exception as e:
                logger.error(f"Error refreshing sessions: {str(e)}", exc_info=True)
                return gr.update(choices=[])
        
        def handle_auto_recover(language):
            session_id, message = auto_recover_session(language)
            return session_id
        
        # Connect event handlers
        submit_btn.click(
            fn=process_query,
            inputs=[query_input, language_select, current_session_id],
            outputs=[stages_display, novel_output, status_text]
        )
        
        # Update novel text state and session ID when novel output changes
        novel_output.change(
            fn=lambda x: x,
            inputs=[novel_output],
            outputs=[novel_text_state]
        )
        
        resume_btn.click(
            fn=lambda x: x,
            inputs=[session_dropdown],
            outputs=[current_session_id]
        ).then(
            fn=resume_session,
            inputs=[current_session_id, language_select],
            outputs=[stages_display, novel_output, status_text]
        )
        
        auto_recover_btn.click(
            fn=handle_auto_recover,
            inputs=[language_select],
            outputs=[current_session_id]
        ).then(
            fn=resume_session,
            inputs=[current_session_id, language_select],
            outputs=[stages_display, novel_output, status_text]
        )
        
        refresh_btn.click(
            fn=refresh_sessions,
            outputs=[session_dropdown]
        )
        
        clear_btn.click(
            fn=lambda: ("", "", "🔄 Ready", "", None),
            outputs=[stages_display, novel_output, status_text, novel_text_state, current_session_id]
        )
        
        def handle_download(format_type, language, session_id, novel_text):
            if not session_id:
                return gr.update(visible=False)
            
            file_path = download_novel(novel_text, format_type, language, session_id)
            if file_path:
                return gr.update(value=file_path, visible=True)
            else:
                return gr.update(visible=False)
        
        download_btn.click(
            fn=handle_download,
            inputs=[format_select, language_select, current_session_id, novel_text_state],
            outputs=[download_file]
        )
        
        # Load sessions on startup
        def on_load():
            sessions = refresh_sessions()
            return sessions
        
        interface.load(
            fn=on_load,
            outputs=[session_dropdown]
        )
    
    return interface

# Main execution
if __name__ == "__main__":
    import sys
    
    logger.info("Starting SOMA Novel Writing System - 10 Writers Edition...")
    
    # Check environment
    if TEST_MODE:
        logger.warning("Running in TEST MODE - no actual API calls will be made")
    else:
        logger.info(f"Running in PRODUCTION MODE with API endpoint: {API_URL}")
    
    # Check web search
    if BRAVE_SEARCH_API_KEY:
        logger.info("Web search is ENABLED with Brave Search API")
    else:
        logger.warning("Web search is DISABLED. Set BRAVE_SEARCH_API_KEY environment variable to enable.")
    
    # Initialize database on startup
    logger.info("Initializing database...")
    NovelDatabase.init_db()
    logger.info("Database initialized successfully.")
    
    interface = create_interface()
    
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )