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.")

# --- Environment variables and constants ---
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 = "dep86pjolcjjnv8"
DB_PATH = "novel_sessions_v6.db"

# Target word count settings
TARGET_WORDS = 8000  # Safety margin
MIN_WORDS_PER_PART = 800  # Minimum words per part

# --- Environment validation ---
if not FRIENDLI_TOKEN:
    logger.error("FRIENDLI_TOKEN not set. Application will not work properly.")
    FRIENDLI_TOKEN = "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()

# Narrative phases definition
NARRATIVE_PHASES = [
    "Introduction: Daily Life and Cracks",
    "Development 1: Rising Anxiety",
    "Development 2: External Shock",
    "Development 3: Deepening Internal Conflict",
    "Climax 1: Peak of Crisis",
    "Climax 2: Moment of Choice",
    "Falling Action 1: Consequences and Aftermath",
    "Falling Action 2: New Recognition",
    "Resolution 1: Changed Daily Life",
    "Resolution 2: Open Questions"
]

# Stage configuration - Single writer system
UNIFIED_STAGES = [
    ("director", "🎬 Director: Integrated Narrative Structure Planning"),
    ("critic_director", "📝 Critic: Deep Review of Narrative Structure"),
    ("director", "🎬 Director: Final Master Plan"),
] + [
    item for i in range(1, 11) 
    for item in [
        ("writer", f"✍️ Writer: Part {i} - {NARRATIVE_PHASES[i-1]}"),
        (f"critic_part{i}", f"📝 Part {i} Critic: Immediate Review and Revision Request"),
        ("writer", f"✍️ Writer: Part {i} Revision")
    ]
] + [
    ("critic_final", "📝 Final Critic: Comprehensive Evaluation and Literary Achievement"),
]

# --- Data classes ---
@dataclass
class StoryBible:
    """Story bible for maintaining narrative consistency"""
    characters: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    settings: Dict[str, str] = field(default_factory=dict)
    timeline: List[Dict[str, Any]] = field(default_factory=list)
    plot_points: List[Dict[str, Any]] = field(default_factory=list)
    themes: List[str] = field(default_factory=list)
    symbols: Dict[str, List[str]] = field(default_factory=dict)
    style_guide: Dict[str, str] = field(default_factory=dict)
    opening_sentence: str = ""

@dataclass
class PartCritique:
    """Critique content for each part"""
    part_number: int
    continuity_issues: List[str] = field(default_factory=list)
    character_consistency: List[str] = field(default_factory=list)
    plot_progression: List[str] = field(default_factory=list)
    thematic_alignment: List[str] = field(default_factory=list)
    technical_issues: List[str] = field(default_factory=list)
    strengths: List[str] = field(default_factory=list)
    required_changes: List[str] = field(default_factory=list)
    literary_quality: List[str] = field(default_factory=list)

# --- Core logic classes ---
class UnifiedNarrativeTracker:
    """Unified narrative tracker for single writer system"""
    def __init__(self):
        self.story_bible = StoryBible()
        self.part_critiques: Dict[int, PartCritique] = {}
        self.accumulated_content: List[str] = []
        self.word_count_by_part: Dict[int, int] = {}
        self.revision_history: Dict[int, List[str]] = defaultdict(list)
        self.causal_chains: List[Dict[str, Any]] = []
        self.narrative_momentum: float = 0.0
        
    def update_story_bible(self, element_type: str, key: str, value: Any):
        """Update story bible"""
        if element_type == "character":
            self.story_bible.characters[key] = value
        elif element_type == "setting":
            self.story_bible.settings[key] = value
        elif element_type == "timeline":
            self.story_bible.timeline.append({"event": key, "details": value})
        elif element_type == "theme":
            if key not in self.story_bible.themes:
                self.story_bible.themes.append(key)
        elif element_type == "symbol":
            if key not in self.story_bible.symbols:
                self.story_bible.symbols[key] = []
            self.story_bible.symbols[key].append(value)
            
    def add_part_critique(self, part_number: int, critique: PartCritique):
        """Add part critique"""
        self.part_critiques[part_number] = critique
        
    def check_continuity(self, current_part: int, new_content: str) -> List[str]:
        """Check continuity"""
        issues = []
        
        # Character consistency check
        for char_name, char_data in self.story_bible.characters.items():
            if char_name in new_content:
                if "traits" in char_data:
                    for trait in char_data["traits"]:
                        if trait.get("abandoned", False):
                            issues.append(f"{char_name}'s abandoned trait '{trait['name']}' reappears")
                            
        # Timeline consistency check
        if len(self.story_bible.timeline) > 0:
            last_event = self.story_bible.timeline[-1]
            
        # Causality check
        if current_part > 1 and not any(kw in new_content for kw in 
            ['because', 'therefore', 'thus', 'hence', 'consequently']):
            issues.append("Unclear causality with previous part")
            
        return issues
    
    def calculate_narrative_momentum(self, part_number: int, content: str) -> float:
        """Calculate narrative momentum"""
        momentum = 5.0
        
        # New elements introduced
        new_elements = len(set(content.split()) - set(' '.join(self.accumulated_content).split()))
        if new_elements > 100:
            momentum += 2.0
            
        # Conflict escalation
        tension_words = ['crisis', 'conflict', 'tension', 'struggle', 'dilemma']
        if any(word in content.lower() for word in tension_words):
            momentum += 1.5
            
        # Causal clarity
        causal_words = ['because', 'therefore', 'thus', 'consequently', 'hence']
        causal_count = sum(1 for word in causal_words if word in content.lower())
        momentum += min(causal_count * 0.5, 2.0)
        
        # Repetition penalty
        if part_number > 1:
            prev_content = self.accumulated_content[-1] if self.accumulated_content else ""
            overlap = len(set(content.split()) & set(prev_content.split()))
            if overlap > len(content.split()) * 0.3:
                momentum -= 3.0
                
        return max(0.0, min(10.0, momentum))

class NovelDatabase:
    """Database management - Modified for single writer system"""
    @staticmethod
    def init_db():
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            cursor = conn.cursor()
            
            # Main sessions table
            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,
                    final_novel TEXT,
                    literary_report TEXT,
                    total_words INTEGER DEFAULT 0,
                    story_bible TEXT,
                    narrative_tracker TEXT,
                    opening_sentence 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',
                    narrative_momentum REAL DEFAULT 0.0,
                    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)
                )
            ''')
            
            # Critiques table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS critiques (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    part_number INTEGER NOT NULL,
                    critique_data TEXT,
                    created_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            ''')
            
            # Random themes library table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS random_themes_library (
                    theme_id TEXT PRIMARY KEY,
                    theme_text TEXT NOT NULL,
                    language TEXT NOT NULL,
                    title TEXT,
                    opening_sentence TEXT,
                    protagonist TEXT,
                    conflict TEXT,
                    philosophical_question TEXT,
                    generated_at TEXT DEFAULT (datetime('now')),
                    view_count INTEGER DEFAULT 0,
                    used_count INTEGER DEFAULT 0,
                    tags TEXT,
                    metadata TEXT
                )
            ''')
            
            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, language: str) -> str:
        session_id = hashlib.md5(f"{user_query}{datetime.now()}".encode()).hexdigest()
        with NovelDatabase.get_db() as conn:
            conn.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',
                   narrative_momentum: float = 0.0):
        word_count = len(content.split()) if content else 0
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO stages (session_id, stage_number, stage_name, role, content, 
                                  word_count, status, narrative_momentum)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(session_id, stage_number)
                DO UPDATE SET content=?, word_count=?, status=?, stage_name=?, 
                             narrative_momentum=?, updated_at=datetime('now')
            ''', (session_id, stage_number, stage_name, role, content, word_count, 
                  status, narrative_momentum, content, word_count, status, stage_name, 
                  narrative_momentum))
            
            # Update total word count
            cursor.execute('''
                UPDATE sessions 
                SET total_words = (
                    SELECT SUM(word_count) 
                    FROM stages 
                    WHERE session_id = ? AND role = 'writer' AND content IS NOT NULL
                ),
                updated_at = datetime('now'),
                current_stage = ?
                WHERE session_id = ?
            ''', (session_id, stage_number, session_id))
            
            conn.commit()

    @staticmethod
    def save_critique(session_id: str, part_number: int, critique: PartCritique):
        """Save critique"""
        with NovelDatabase.get_db() as conn:
            critique_json = json.dumps(asdict(critique))
            conn.cursor().execute(
                'INSERT INTO critiques (session_id, part_number, critique_data) VALUES (?, ?, ?)',
                (session_id, part_number, critique_json)
            )
            conn.commit()

    @staticmethod
    def save_opening_sentence(session_id: str, opening_sentence: str):
        """Save opening sentence"""
        with NovelDatabase.get_db() as conn:
            conn.cursor().execute(
                'UPDATE sessions SET opening_sentence = ? WHERE session_id = ?',
                (opening_sentence, session_id)
            )
            conn.commit()

    @staticmethod
    def get_writer_content(session_id: str) -> str:
        """Get writer content - Integrate all revisions"""
        with NovelDatabase.get_db() as conn:
            rows = conn.cursor().execute('''
                SELECT content FROM stages 
                WHERE session_id = ? AND role = 'writer' 
                AND stage_name LIKE '%Revision%'
                ORDER BY stage_number
            ''', (session_id,)).fetchall()
            
            if rows:
                return '\n\n'.join(row['content'] for row in rows if row['content'])
            else:
                # If no revisions, use drafts
                rows = conn.cursor().execute('''
                    SELECT content FROM stages 
                    WHERE session_id = ? AND role = 'writer' 
                    AND stage_name NOT LIKE '%Revision%'
                    ORDER BY stage_number
                ''', (session_id,)).fetchall()
                return '\n\n'.join(row['content'] for row in rows if row['content'])

    @staticmethod
    def save_narrative_tracker(session_id: str, tracker: UnifiedNarrativeTracker):
        """Save unified narrative tracker"""
        with NovelDatabase.get_db() as conn:
            tracker_data = json.dumps({
                'story_bible': asdict(tracker.story_bible),
                'part_critiques': {k: asdict(v) for k, v in tracker.part_critiques.items()},
                'word_count_by_part': tracker.word_count_by_part,
                'causal_chains': tracker.causal_chains,
                'narrative_momentum': tracker.narrative_momentum
            })
            conn.cursor().execute(
                'UPDATE sessions SET narrative_tracker = ? WHERE session_id = ?',
                (tracker_data, session_id)
            )
            conn.commit()

    @staticmethod
    def load_narrative_tracker(session_id: str) -> Optional[UnifiedNarrativeTracker]:
        """Load unified narrative tracker"""
        with NovelDatabase.get_db() as conn:
            row = conn.cursor().execute(
                'SELECT narrative_tracker FROM sessions WHERE session_id = ?',
                (session_id,)
            ).fetchone()
            
            if row and row['narrative_tracker']:
                data = json.loads(row['narrative_tracker'])
                tracker = UnifiedNarrativeTracker()
                
                # Restore story bible
                bible_data = data.get('story_bible', {})
                tracker.story_bible = StoryBible(**bible_data)
                
                # Restore critiques
                for part_num, critique_data in data.get('part_critiques', {}).items():
                    tracker.part_critiques[int(part_num)] = PartCritique(**critique_data)
                
                tracker.word_count_by_part = data.get('word_count_by_part', {})
                tracker.causal_chains = data.get('causal_chains', [])
                tracker.narrative_momentum = data.get('narrative_momentum', 0.0)
                
                return tracker
        return None

    @staticmethod
    def save_random_theme(theme_text: str, language: str, metadata: Dict[str, Any]) -> str:
        """Save randomly generated theme to library"""
        theme_id = hashlib.md5(f"{theme_text}{datetime.now()}".encode()).hexdigest()[:12]
        
        # Extract components from theme text
        title = metadata.get('title', '')
        opening_sentence = metadata.get('opening_sentence', '')
        protagonist = metadata.get('protagonist', '')
        conflict = metadata.get('conflict', '')
        philosophical_question = metadata.get('philosophical_question', '')
        tags = json.dumps(metadata.get('tags', []))
        
        with NovelDatabase.get_db() as conn:
            conn.cursor().execute('''
                INSERT INTO random_themes_library 
                (theme_id, theme_text, language, title, opening_sentence, 
                 protagonist, conflict, philosophical_question, tags, metadata)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (theme_id, theme_text, language, title, opening_sentence,
                  protagonist, conflict, philosophical_question, tags, 
                  json.dumps(metadata)))
            conn.commit()
        
        return theme_id

    @staticmethod
    def get_random_themes_library(language: str = None, limit: int = 50) -> List[Dict]:
        """Get random themes from library"""
        with NovelDatabase.get_db() as conn:
            query = '''
                SELECT * FROM random_themes_library 
                {} 
                ORDER BY generated_at DESC 
                LIMIT ?
            '''.format('WHERE language = ?' if language else '')
            
            if language:
                rows = conn.cursor().execute(query, (language, limit)).fetchall()
            else:
                rows = conn.cursor().execute(query, (limit,)).fetchall()
            
            return [dict(row) for row in rows]

    @staticmethod
    def update_theme_view_count(theme_id: str):
        """Update view count for a theme"""
        with NovelDatabase.get_db() as conn:
            conn.cursor().execute(
                'UPDATE random_themes_library SET view_count = view_count + 1 WHERE theme_id = ?',
                (theme_id,)
            )
            conn.commit()

    @staticmethod
    def update_theme_used_count(theme_id: str):
        """Update used count when theme is used for novel"""
        with NovelDatabase.get_db() as conn:
            conn.cursor().execute(
                'UPDATE random_themes_library SET used_count = used_count + 1 WHERE theme_id = ?',
                (theme_id,)
            )
            conn.commit()

    @staticmethod
    def get_theme_by_id(theme_id: str) -> Optional[Dict]:
        """Get specific theme by ID"""
        with NovelDatabase.get_db() as conn:
            row = conn.cursor().execute(
                'SELECT * FROM random_themes_library WHERE theme_id = ?',
                (theme_id,)
            ).fetchone()
            return dict(row) if row else None

    @staticmethod
    def get_session(session_id: str) -> Optional[Dict]:
        with NovelDatabase.get_db() as conn:
            row = conn.cursor().execute('SELECT * FROM sessions WHERE session_id = ?', 
                                       (session_id,)).fetchone()
            return dict(row) if row else None

    @staticmethod
    def get_stages(session_id: str) -> List[Dict]:
        with NovelDatabase.get_db() as conn:
            rows = conn.cursor().execute(
                'SELECT * FROM stages WHERE session_id = ? ORDER BY stage_number', 
                (session_id,)
            ).fetchall()
            return [dict(row) for row in rows]

    @staticmethod
    def update_final_novel(session_id: str, final_novel: str, literary_report: str = ""):
        with NovelDatabase.get_db() as conn:
            conn.cursor().execute(
                '''UPDATE sessions SET final_novel = ?, status = 'complete', 
                   updated_at = datetime('now'), literary_report = ? WHERE session_id = ?''',
                (final_novel, literary_report, session_id)
            )
            conn.commit()

    @staticmethod
    def get_active_sessions() -> List[Dict]:
        with NovelDatabase.get_db() as conn:
            rows = conn.cursor().execute(
                '''SELECT session_id, user_query, language, created_at, current_stage, total_words 
                   FROM sessions WHERE status = 'active' ORDER BY updated_at DESC LIMIT 10'''
            ).fetchall()
            return [dict(row) for row in rows]

    @staticmethod
    def get_total_words(session_id: str) -> int:
        """Get total word count"""
        with NovelDatabase.get_db() as conn:
            row = conn.cursor().execute(
                'SELECT total_words FROM sessions WHERE session_id = ?',
                (session_id,)
            ).fetchone()
            return row['total_words'] if row and row['total_words'] else 0

class WebSearchIntegration:
    """Web search functionality"""
    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 = 3, language: str = "en") -> List[Dict]:
        if not self.enabled:
            return []
        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": self.brave_api_key
        }
        params = {
            "q": query,
            "count": count,
            "search_lang": "ko" if language == "Korean" else "en",
            "text_decorations": False,
            "safesearch": "moderate"
        }
        try:
            response = requests.get(self.search_url, headers=headers, params=params, timeout=10)
            response.raise_for_status()
            results = response.json().get("web", {}).get("results", [])
            return results
        except requests.exceptions.RequestException as e:
            logger.error(f"Web search API error: {e}")
            return []

    def extract_relevant_info(self, results: List[Dict], max_chars: int = 1500) -> str:
        if not results:
            return ""
        extracted = []
        total_chars = 0
        for i, result in enumerate(results[:3], 1):
            title = result.get("title", "")
            description = result.get("description", "")
            info = f"[{i}] {title}: {description}"
            if total_chars + len(info) < max_chars:
                extracted.append(info)
                total_chars += len(info)
            else:
                break
        return "\n".join(extracted)

class UnifiedLiterarySystem:
    """Single writer progressive literary novel generation system"""
    def __init__(self):
        self.token = FRIENDLI_TOKEN
        self.api_url = API_URL
        self.model_id = MODEL_ID
        self.narrative_tracker = UnifiedNarrativeTracker()
        self.web_search = WebSearchIntegration()
        self.current_session_id = None
        NovelDatabase.init_db()

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

    # --- Prompt generation functions ---
    def augment_query(self, user_query: str, language: str) -> str:
        """Augment prompt"""
        if len(user_query.split()) < 15:
            augmented_template = {
                "Korean": f"""'{user_query}'

**서사 구조 핵심:**
- 10개 파트가 하나의 통합된 이야기를 구성
- 각 파트는 이전 파트의 필연적 결과
- 인물의 명확한 변화 궤적 (A → B → C)
- 중심 갈등의 점진적 고조와 해결
- 강렬한 중심 상징의 의미 변화""",
                
                "English": f"""'{user_query}'

**Narrative Structure Core:**
- 10 parts forming one integrated story
- Each part as inevitable result of previous
- Clear character transformation arc (A → B → C)
- Progressive escalation and resolution of central conflict
- Evolving meaning of powerful central symbol"""
            }
            return augmented_template.get(language, user_query)
        return user_query

    def generate_powerful_opening(self, user_query: str, language: str) -> str:
        """Generate powerful opening sentence matching the theme"""
        
        opening_prompt = {
            "Korean": f"""주제: {user_query}

이 주제에 대한 강렬하고 잊을 수 없는 첫문장을 생성하세요.

**첫문장 작성 원칙:**
1. 즉각적인 긴장감이나 궁금증 유발
2. 평범하지 않은 시각이나 상황 제시
3. 감각적이고 구체적인 이미지
4. 철학적 질문이나 역설적 진술
5. 시간과 공간의 독특한 설정

**훌륭한 첫문장의 예시 패턴:**
- "그가 죽은 날, ..." (충격적 사건)
- "모든 것이 끝났다고 생각한 순간..." (반전 예고)
- "세상에서 가장 [형용사]한 [명사]는..." (독특한 정의)
- "[구체적 행동]하는 것만으로도..." (일상의 재해석)

단 하나의 문장만 제시하세요.""",

            "English": f"""Theme: {user_query}

Generate an unforgettable opening sentence for this theme.

**Opening Sentence Principles:**
1. Immediate tension or curiosity
2. Unusual perspective or situation
3. Sensory and specific imagery
4. Philosophical question or paradox
5. Unique temporal/spatial setting

**Great Opening Patterns:**
- "The day he died, ..." (shocking event)
- "At the moment everything seemed over..." (reversal hint)
- "The most [adjective] [noun] in the world..." (unique definition)
- "Just by [specific action]..." (reinterpretation of ordinary)

Provide only one sentence."""
        }
        
        messages = [{"role": "user", "content": opening_prompt.get(language, opening_prompt["Korean"])}]
        opening = self.call_llm_sync(messages, "writer", language)
        return opening.strip()

    def create_director_initial_prompt(self, user_query: str, language: str) -> str:
        """Director initial planning - Enhanced version"""
        augmented_query = self.augment_query(user_query, language)
        
        # Generate opening sentence
        opening_sentence = self.generate_powerful_opening(user_query, language)
        self.narrative_tracker.story_bible.opening_sentence = opening_sentence
        if self.current_session_id:
            NovelDatabase.save_opening_sentence(self.current_session_id, opening_sentence)
        
        search_results_str = ""
        if self.web_search.enabled:
            short_query = user_query[:50] if len(user_query) > 50 else user_query
            queries = [
                f"{short_query} philosophical meaning",
                f"human existence meaning {short_query}",
                f"{short_query} literary works"
            ]
            for q in queries[:2]:
                try:
                    results = self.web_search.search(q, count=2, language=language)
                    if results:
                        search_results_str += self.web_search.extract_relevant_info(results) + "\n"
                except Exception as e:
                    logger.warning(f"Search failed: {str(e)}")
    
        lang_prompts = {
            "Korean": f"""노벨문학상 수준의 철학적 깊이를 지닌 중편소설(8,000단어)을 기획하세요.

**주제:** {augmented_query}

**필수 첫문장:** {opening_sentence}

**참고 자료:**
{search_results_str if search_results_str else "N/A"}

**필수 문학적 요소:**

1. **철학적 탐구**
   - 현대인의 실존적 고뇌 (소외, 정체성, 의미 상실)
   - 디지털 시대의 인간 조건
   - 자본주의 사회의 모순과 개인의 선택
   - 죽음, 사랑, 자유에 대한 새로운 성찰

2. **사회적 메시지**
   - 계급, 젠더, 세대 간 갈등
   - 환경 위기와 인간의 책임
   - 기술 발전과 인간성의 충돌
   - 현대 민주주의의 위기와 개인의 역할

3. **문학적 수사 장치**
   - 중심 은유: [구체적 사물/현상] → [추상적 의미]
   - 반복되는 모티프: [이미지/행동] (최소 5회 변주)
   - 대조법: [A vs B]의 지속적 긴장
   - 상징적 공간: [구체적 장소]가 의미하는 것
   - 시간의 주관적 흐름 (회상, 예감, 정지)

4. **통합된 10파트 구조**
   각 파트별 핵심:
   - 파트 1: 첫문장으로 시작, 일상 속 균열 → 철학적 질문 제기
   - 파트 2-3: 외부 사건 → 내적 성찰 심화
   - 파트 4-5: 사회적 갈등 → 개인적 딜레마
   - 파트 6-7: 위기의 정점 → 실존적 선택
   - 파트 8-9: 선택의 결과 → 새로운 인식
   - 파트 10: 변화된 세계관 → 열린 질문

5. **문체 지침**
   - 시적 산문체: 일상 언어와 은유의 균형
   - 의식의 흐름과 객관적 묘사의 교차
   - 짧고 강렬한 문장과 성찰적 긴 문장의 리듬
   - 감각적 디테일로 추상적 개념 구현

구체적이고 혁신적인 계획을 제시하세요.""",

            "English": f"""Plan a philosophically profound novella (8,000 words) worthy of Nobel Prize.

**Theme:** {augmented_query}

**Required Opening:** {opening_sentence}

**Reference:**
{search_results_str if search_results_str else "N/A"}

**Essential Literary Elements:**

1. **Philosophical Exploration**
   - Modern existential anguish (alienation, identity, loss of meaning)
   - Human condition in digital age
   - Capitalist contradictions and individual choice
   - New reflections on death, love, freedom

2. **Social Message**
   - Class, gender, generational conflicts
   - Environmental crisis and human responsibility
   - Technology vs humanity collision
   - Modern democracy crisis and individual role

3. **Literary Devices**
   - Central metaphor: [concrete object/phenomenon] → [abstract meaning]
   - Recurring motif: [image/action] (minimum 5 variations)
   - Contrast: sustained tension of [A vs B]
   - Symbolic space: what [specific place] means
   - Subjective time flow (flashback, premonition, pause)

4. **Integrated 10-Part Structure**
   Each part's core:
   - Part 1: Start with opening sentence, daily cracks → philosophical questions
   - Part 2-3: External events → deepening introspection
   - Part 4-5: Social conflict → personal dilemma
   - Part 6-7: Crisis peak → existential choice
   - Part 8-9: Choice consequences → new recognition
   - Part 10: Changed worldview → open questions

5. **Style Guidelines**
   - Poetic prose: balance of everyday language and metaphor
   - Stream of consciousness crossing with objective description
   - Rhythm of short intense sentences and reflective long ones
   - Abstract concepts through sensory details

Provide concrete, innovative plan."""
        }
    
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_critic_director_prompt(self, director_plan: str, user_query: str, language: str) -> str:
        """Director plan deep review - Enhanced version"""
        lang_prompts = {
            "Korean": f"""서사 구조 전문가로서 이 기획을 심층 분석하세요.

**원 주제:** {user_query}

**감독자 기획:**
{director_plan}

**심층 검토 항목:**

1. **인과관계 검증**
   각 파트 간 연결을 검토하고 논리적 비약을 찾으세요:
   - 파트 1→2: [연결성 평가]
   - 파트 2→3: [연결성 평가]
   (모든 연결 지점 검토)

2. **철학적 깊이 평가**
   - 제시된 철학적 주제가 충분히 깊은가?
   - 현대적 관련성이 있는가?
   - 독창적 통찰이 있는가?

3. **문학적 장치의 효과성**
   - 은유와 상징이 유기적으로 작동하는가?
   - 과도하거나 부족하지 않은가?
   - 주제와 긴밀히 연결되는가?

4. **캐릭터 아크 실현 가능성**
   - 변화가 충분히 점진적인가?
   - 각 단계의 동기가 명확한가?
   - 심리적 신뢰성이 있는가?

5. **8,000단어 실현 가능성**
   - 각 파트가 800단어를 유지할 수 있는가?
   - 늘어지거나 압축되는 부분은 없는가?

**필수 개선사항을 구체적으로 제시하세요.**""",

            "English": f"""As narrative structure expert, deeply analyze this plan.

**Original Theme:** {user_query}

**Director's Plan:**
{director_plan}

**Deep Review Items:**

1. **Causality Verification**
   Review connections between parts, find logical leaps:
   - Part 1→2: [Connection assessment]
   - Part 2→3: [Connection assessment]
   (Review all connection points)

2. **Philosophical Depth Assessment**
   - Is philosophical theme deep enough?
   - Contemporary relevance?
   - Original insights?

3. **Literary Device Effectiveness**
   - Do metaphors and symbols work organically?
   - Not excessive or insufficient?
   - Tightly connected to theme?

4. **Character Arc Feasibility**
   - Is change sufficiently gradual?
   - Are motivations clear at each stage?
   - Psychological credibility?

5. **8,000-word Feasibility**
   - Can each part sustain 800 words?
   - Any dragging or compressed sections?

**Provide specific required improvements.**"""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_writer_prompt(self, part_number: int, master_plan: str, 
                           accumulated_content: str, story_bible: StoryBible,
                           language: str) -> str:
        """Single writer prompt - Enhanced version"""
        
        phase_name = NARRATIVE_PHASES[part_number-1]
        target_words = MIN_WORDS_PER_PART
        
        # Part-specific instructions
        philosophical_focus = {
            1: "Introduce existential anxiety through daily cracks",
            2: "First collision between individual and society",
            3: "Self-recognition through encounter with others",
            4: "Shaking beliefs and clashing values",
            5: "Weight of choice and paradox of freedom",
            6: "Test of humanity in extreme situations",
            7: "Weight of consequences and responsibility",
            8: "Self-rediscovery through others' gaze",
            9: "Reconciliation with the irreconcilable",
            10: "New life possibilities and unresolved questions"
        }
        
        literary_techniques = {
            1: "Introducing objective correlative",
            2: "Contrapuntal narration",
            3: "Stream of consciousness",
            4: "Subtle shifts in perspective",
            5: "Aesthetics of silence and omission",
            6: "Subjective transformation of time",
            7: "Intersection of multiple viewpoints",
            8: "Subversion of metaphor",
            9: "Reinterpretation of archetypal images",
            10: "Multi-layered open ending"
        }
        
        # Story bible summary
        bible_summary = f"""
**Characters:** {', '.join(story_bible.characters.keys()) if story_bible.characters else 'TBD'}
**Key Symbols:** {', '.join(story_bible.symbols.keys()) if story_bible.symbols else 'TBD'}
**Themes:** {', '.join(story_bible.themes[:3]) if story_bible.themes else 'TBD'}
**Style:** {story_bible.style_guide.get('voice', 'N/A')}
"""
        
        # Previous content summary
        prev_content = ""
        if accumulated_content:
            prev_parts = accumulated_content.split('\n\n')
            if len(prev_parts) >= 1:
                prev_content = prev_parts[-1][-2000:]  # Last 2000 chars of previous part
        
        lang_prompts = {
            "Korean": f"""당신은 현대 문학의 최전선에 선 작가입니다.
**현재: 파트 {part_number} - {phase_name}**

{"**필수 첫문장:** " + story_bible.opening_sentence if part_number == 1 and story_bible.opening_sentence else ""}

**이번 파트의 철학적 초점:** {philosophical_focus[part_number]}
**핵심 문학 기법:** {literary_techniques[part_number]}

**전체 계획:**
{master_plan}

**스토리 바이블:**
{bible_summary}

**직전 내용:**
{prev_content if prev_content else "첫 파트입니다"}

**파트 {part_number} 작성 지침:**

1. **분량:** {target_words}-900 단어 (필수)

2. **문학적 수사 요구사항:**
   - 최소 3개의 독창적 은유/직유
   - 1개 이상의 상징적 이미지 심화
   - 감각적 묘사와 추상적 사유의 융합
   - 리듬감 있는 문장 구성 (장단의 변주)

3. **현대적 고뇌 표현:**
   - 디지털 시대의 소외감
   - 자본주의적 삶의 부조리
   - 관계의 표면성과 진정성 갈망
   - 의미 추구와 무의미의 직면

4. **사회적 메시지 내재화:**
   - 직접적 주장이 아닌 상황과 인물을 통한 암시
   - 개인의 고통과 사회 구조의 연결
   - 미시적 일상과 거시적 문제의 교차

5. **서사적 추진력:**
   - 이전 파트의 필연적 결과로 시작
   - 새로운 갈등 층위 추가
   - 다음 파트를 향한 긴장감 조성

**문학적 금기:**
- 진부한 표현이나 상투적 은유
- 감정의 직접적 설명
- 도덕적 판단이나 교훈
- 인위적인 해결이나 위안

파트 {part_number}를 깊이 있는 문학적 성취로 만드세요.""",

            "English": f"""You are a writer at the forefront of contemporary literature.
**Current: Part {part_number} - {phase_name}**

{"**Required Opening:** " + story_bible.opening_sentence if part_number == 1 and story_bible.opening_sentence else ""}

**Philosophical Focus:** {philosophical_focus[part_number]}
**Core Literary Technique:** {literary_techniques[part_number]}

**Master Plan:**
{master_plan}

**Story Bible:**
{bible_summary}

**Previous Content:**
{prev_content if prev_content else "This is the first part"}

**Part {part_number} Guidelines:**

1. **Length:** {target_words}-900 words (mandatory)

2. **Literary Device Requirements:**
   - Minimum 3 original metaphors/similes
   - Deepen at least 1 symbolic image
   - Fusion of sensory description and abstract thought
   - Rhythmic sentence composition (variation of long/short)

3. **Modern Anguish Expression:**
   - Digital age alienation
   - Absurdity of capitalist life
   - Surface relationships vs authenticity yearning
   - Meaning pursuit vs confronting meaninglessness

4. **Social Message Internalization:**
   - Implication through situation and character, not direct claim
   - Connection between individual pain and social structure
   - Intersection of micro daily life and macro problems

5. **Narrative Momentum:**
   - Start as inevitable result of previous part
   - Add new conflict layers
   - Create tension toward next part

**Literary Taboos:**
- Clichéd expressions or trite metaphors
- Direct emotion explanation
- Moral judgment or preaching
- Artificial resolution or comfort

Make Part {part_number} a profound literary achievement."""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_part_critic_prompt(self, part_number: int, part_content: str, 
                                master_plan: str, accumulated_content: str,
                                story_bible: StoryBible, language: str) -> str:
        """Part-by-part immediate critique - Enhanced version"""
        
        lang_prompts = {
            "Korean": f"""파트 {part_number}의 문학적 성취도를 엄격히 평가하세요.

**마스터플랜 파트 {part_number} 요구사항:**
{self._extract_part_plan(master_plan, part_number)}

**작성된 내용:**
{part_content}

**스토리 바이블 체크:**
- 캐릭터: {', '.join(story_bible.characters.keys())}
- 설정: {', '.join(story_bible.settings.keys())}

**평가 기준:**

1. **문학적 수사 (30%)**
   - 은유와 상징의 독창성
   - 언어의 시적 밀도
   - 이미지의 선명도와 깊이
   - 문장의 리듬과 음악성

2. **철학적 깊이 (25%)**
   - 실존적 질문의 제기
   - 현대인의 조건 탐구
   - 보편성과 특수성의 균형
   - 사유의 독창성

3. **사회적 통찰 (20%)**
   - 시대정신의 포착
   - 구조와 개인의 관계
   - 비판적 시각의 예리함
   - 대안적 상상력

4. **서사적 완성도 (25%)**
   - 인과관계의 필연성
   - 긴장감의 유지
   - 인물의 입체성
   - 구조적 통일성

**구체적 지적사항:**
- 진부한 표현: [예시와 대안]
- 철학적 천착 부족: [보완 방향]
- 사회적 메시지 불명확: [강화 방안]
- 서사적 허점: [수정 필요]

**필수 개선 요구:**
문학적 수준을 노벨상 급으로 끌어올리기 위한 구체적 수정안을 제시하세요.""",

            "English": f"""Strictly evaluate literary achievement of Part {part_number}.

**Master Plan Part {part_number} Requirements:**
{self._extract_part_plan(master_plan, part_number)}

**Written Content:**
{part_content}

**Story Bible Check:**
- Characters: {', '.join(story_bible.characters.keys()) if story_bible.characters else 'None yet'}
- Settings: {', '.join(story_bible.settings.keys()) if story_bible.settings else 'None yet'}

**Evaluation Criteria:**

1. **Literary Rhetoric (30%)**
   - Originality of metaphor and symbol
   - Poetic density of language
   - Clarity and depth of imagery
   - Rhythm and musicality of sentences

2. **Philosophical Depth (25%)**
   - Raising existential questions
   - Exploring modern human condition
   - Balance of universality and specificity
   - Originality of thought

3. **Social Insight (20%)**
   - Capturing zeitgeist
   - Relationship between structure and individual
   - Sharpness of critical perspective
   - Alternative imagination

4. **Narrative Completion (25%)**
   - Inevitability of causality
   - Maintaining tension
   - Character dimensionality
   - Structural unity

**Specific Points:**
- Clichéd expressions: [examples and alternatives]
- Insufficient philosophical exploration: [enhancement direction]
- Unclear social message: [strengthening methods]
- Narrative gaps: [needed revisions]

**Required Improvements:**
Provide specific revisions to elevate literary level to Nobel Prize standard."""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_writer_revision_prompt(self, part_number: int, original_content: str,
                                    critic_feedback: str, language: str) -> str:
        """Writer revision prompt"""
        
        lang_prompts = {
            "Korean": f"""파트 {part_number}를 비평에 따라 수정하세요.

**원본:**
{original_content}

**비평 피드백:**
{critic_feedback}

**수정 지침:**
1. 모든 '필수 수정' 사항을 반영
2. 가능한 '권장 개선' 사항도 포함
3. 원본의 강점은 유지
4. 분량 {MIN_WORDS_PER_PART}단어 이상 유지
5. 작가로서의 일관된 목소리 유지
6. 문학적 수준을 한 단계 높이기

수정본만 제시하세요. 설명은 불필요합니다.""",

            "English": f"""Revise Part {part_number} according to critique.

**Original:**
{original_content}

**Critique Feedback:**
{critic_feedback}

**Revision Guidelines:**
1. Reflect all 'Required fixes'
2. Include 'Recommended improvements' where possible
3. Maintain original strengths
4. Keep length {MIN_WORDS_PER_PART}+ words
5. Maintain consistent authorial voice
6. Elevate literary level

Present only the revision. No explanation needed."""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_final_critic_prompt(self, complete_novel: str, word_count: int, 
                                 story_bible: StoryBible, language: str) -> str:
        """Final comprehensive evaluation"""
        
        lang_prompts = {
            "Korean": f"""완성된 소설을 종합 평가하세요.

**작품 정보:**
- 총 분량: {word_count}단어
- 목표: 8,000단어

**평가 기준:**

1. **서사적 통합성 (30점)**
   - 10개 파트가 하나의 이야기로 통합되었는가?
   - 인과관계가 명확하고 필연적인가?
   - 반복이나 순환 없이 진행되는가?

2. **캐릭터 아크 (25점)**
   - 주인공의 변화가 설득력 있는가?
   - 변화가 점진적이고 자연스러운가?
   - 최종 상태가 초기와 명확히 다른가?

3. **문학적 성취 (25점)**
   - 주제가 깊이 있게 탐구되었는가?
   - 상징이 효과적으로 활용되었는가?
   - 문체가 일관되고 아름다운가?
   - 현대적 철학과 사회적 메시지가 녹아있는가?

4. **기술적 완성도 (20점)**
   - 목표 분량을 달성했는가?
   - 각 파트가 균형 있게 전개되었는가?
   - 문법과 표현이 정확한가?

**총점: /100점**

구체적인 강점과 약점을 제시하세요.""",

            "English": f"""Comprehensively evaluate the completed novel.

**Work Info:**
- Total length: {word_count} words
- Target: 8,000 words

**Evaluation Criteria:**

1. **Narrative Integration (30 points)**
   - Are 10 parts integrated into one story?
   - Clear and inevitable causality?
   - Progress without repetition or cycles?

2. **Character Arc (25 points)**
   - Convincing protagonist transformation?
   - Gradual and natural changes?
   - Final state clearly different from initial?

3. **Literary Achievement (25 points)**
   - Theme explored with depth?
   - Symbols used effectively?
   - Consistent and beautiful style?
   - Contemporary philosophy and social message integrated?

4. **Technical Completion (20 points)**
   - Target length achieved?
   - Each part balanced in development?
   - Grammar and expression accurate?

**Total Score: /100 points**

Present specific strengths and weaknesses."""
        }
        
        return lang_prompts.get(language, lang_prompts["Korean"])

    def create_director_final_prompt(self, initial_plan: str, critic_feedback: str, 
                                  user_query: str, language: str) -> str:
        """Director final master plan"""
        return f"""Reflect the critique and complete the final master plan.

**Original Theme:** {user_query}

**Initial Plan:**
{initial_plan}

**Critique Feedback:**
{critic_feedback}

**Final Master Plan Requirements:**
1. Reflect all critique points
2. Specific content and causality for 10 parts
3. Clear transformation stages of protagonist
4. Meaning evolution process of central symbol
5. Feasibility of 800 words per part
6. Implementation of philosophical depth and social message

Present concrete and executable final plan."""

    def _extract_part_plan(self, master_plan: str, part_number: int) -> str:
        """Extract specific part plan from master plan"""
        lines = master_plan.split('\n')
        part_section = []
        capturing = False
        
        for line in lines:
            if f"Part {part_number}:" in line or f"파트 {part_number}:" in line:
                capturing = True
            elif capturing and (f"Part {part_number+1}:" in line or f"파트 {part_number+1}:" in line):
                break
            elif capturing:
                part_section.append(line)
                
        return '\n'.join(part_section) if part_section else "Cannot find the part plan."

    # --- 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 == "writer" else 10000
            
            payload = {
                "model": self.model_id,
                "messages": full_messages,
                "max_tokens": max_tokens,
                "temperature": 0.8,
                "top_p": 0.95,
                "presence_penalty": 0.5,
                "frequency_penalty": 0.3,
                "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]:
        """Role-specific system prompts - Enhanced version"""
        
        base_prompts = {
            "Korean": {
                "director": """당신은 현대 세계문학의 정점을 지향하는 작품을 설계합니다.
깊은 철학적 통찰과 날카로운 사회 비판을 결합하세요.
인간 조건의 복잡성을 10개의 유기적 파트로 구현하세요.
독자의 영혼을 뒤흔들 강렬한 첫문장부터 시작하세요.""",
                
                "critic_director": """서사 구조의 논리성과 실현 가능성을 검증하는 전문가입니다.
인과관계의 허점을 찾아내세요.
캐릭터 발전의 신빙성을 평가하세요.
철학적 깊이와 문학적 가치를 판단하세요.
8,000단어 분량의 적절성을 판단하세요.""",
                
                "writer": """당신은 언어의 연금술사입니다.
일상어를 시로, 구체를 추상으로, 개인을 보편으로 변환하세요.
현대인의 영혼의 어둠과 빛을 동시에 포착하세요.
독자가 자신을 재발견하게 만드는 거울이 되세요.""",
                
                "critic_final": """당신은 작품의 문학적 잠재력을 극대화하는 조력자입니다.
평범함을 비범함으로 이끄는 날카로운 통찰을 제공하세요.
작가의 무의식에 잠든 보석을 발굴하세요.
타협 없는 기준으로 최고를 요구하세요."""
            },
            "English": {
                "director": """You design works aiming for the pinnacle of contemporary world literature.
Combine deep philosophical insights with sharp social criticism.
Implement the complexity of the human condition in 10 organic parts.
Start with an intense opening sentence that shakes the reader's soul.""",
                
                "critic_director": """You are an expert verifying narrative logic and feasibility.
Find gaps in causality.
Evaluate credibility of character development.
Judge philosophical depth and literary value.
Judge appropriateness of 8,000-word length.""",
                
                "writer": """You are an alchemist of language.
Transform everyday language into poetry, concrete into abstract, individual into universal.
Capture both darkness and light of the modern soul.
Become a mirror where readers rediscover themselves.""",
                
                "critic_final": """You are a collaborator maximizing the work's literary potential.
Provide sharp insights leading ordinariness to extraordinariness.
Excavate gems sleeping in the writer's unconscious.
Demand the best with uncompromising standards."""
            }
        }
        
        prompts = base_prompts.get(language, base_prompts["Korean"]).copy()
        
        # Add part-specific critic prompts
        for i in range(1, 11):
            prompts[f"critic_part{i}"] = f"""You are Part {i} dedicated critic.
Review causality with previous parts as top priority.
Verify character consistency and development.
Evaluate alignment with master plan.
Assess literary level and philosophical depth.
Provide specific and actionable revision instructions."""
        
        return prompts

    # --- Main process ---
    def process_novel_stream(self, query: str, language: str, 
                           session_id: Optional[str] = None) -> Generator[Tuple[str, List[Dict[str, Any]], str], None, None]:
        """Single writer novel generation process"""
        try:
            resume_from_stage = 0
            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
                    saved_tracker = NovelDatabase.load_narrative_tracker(session_id)
                    if saved_tracker:
                        self.narrative_tracker = saved_tracker
            else:
                self.current_session_id = NovelDatabase.create_session(query, language)
                logger.info(f"Created new session: {self.current_session_id}")

            stages = []
            if resume_from_stage > 0:
                stages = [{
                    "name": s['stage_name'], 
                    "status": s['status'], 
                    "content": s.get('content', ''),
                    "word_count": s.get('word_count', 0),
                    "momentum": s.get('narrative_momentum', 0.0)
                } for s in NovelDatabase.get_stages(self.current_session_id)]

            total_words = NovelDatabase.get_total_words(self.current_session_id)

            for stage_idx in range(resume_from_stage, len(UNIFIED_STAGES)):
                role, stage_name = UNIFIED_STAGES[stage_idx]
                if stage_idx >= len(stages):
                    stages.append({
                        "name": stage_name, 
                        "status": "active", 
                        "content": "", 
                        "word_count": 0,
                        "momentum": 0.0
                    })
                else:
                    stages[stage_idx]["status"] = "active"

                yield f"🔄 Processing... (Current {total_words:,} words)", stages, self.current_session_id
                
                prompt = self.get_stage_prompt(stage_idx, role, query, language, stages)
                stage_content = ""
                
                for chunk in self.call_llm_streaming([{"role": "user", "content": prompt}], role, language):
                    stage_content += chunk
                    stages[stage_idx]["content"] = stage_content
                    stages[stage_idx]["word_count"] = len(stage_content.split())
                    yield f"🔄 {stage_name} writing... ({total_words + stages[stage_idx]['word_count']:,} words)", stages, self.current_session_id

                # Content processing and tracking
                if role == "writer":
                    # Calculate part number
                    part_num = self._get_part_number(stage_idx)
                    if part_num:
                        self.narrative_tracker.accumulated_content.append(stage_content)
                        self.narrative_tracker.word_count_by_part[part_num] = len(stage_content.split())
                        
                        # Calculate narrative momentum
                        momentum = self.narrative_tracker.calculate_narrative_momentum(part_num, stage_content)
                        stages[stage_idx]["momentum"] = momentum
                        
                        # Update story bible
                        self._update_story_bible_from_content(stage_content, part_num)

                stages[stage_idx]["status"] = "complete"
                NovelDatabase.save_stage(
                    self.current_session_id, stage_idx, stage_name, role,
                    stage_content, "complete", stages[stage_idx].get("momentum", 0.0)
                )
                
                NovelDatabase.save_narrative_tracker(self.current_session_id, self.narrative_tracker)
                total_words = NovelDatabase.get_total_words(self.current_session_id)
                yield f"✅ {stage_name} completed (Total {total_words:,} words)", stages, self.current_session_id

            # Final processing
            final_novel = NovelDatabase.get_writer_content(self.current_session_id)
            final_word_count = len(final_novel.split())
            final_report = self.generate_literary_report(final_novel, final_word_count, language)
            
            NovelDatabase.update_final_novel(self.current_session_id, final_novel, final_report)
            yield f"✅ Novel completed! Total {final_word_count:,} words", stages, self.current_session_id

        except Exception as e:
            logger.error(f"Novel generation process error: {e}", exc_info=True)
            yield f"❌ Error occurred: {e}", stages if 'stages' in locals() else [], self.current_session_id

    def get_stage_prompt(self, stage_idx: int, role: str, query: str, 
                        language: str, stages: List[Dict]) -> str:
        """Generate stage-specific prompt"""
        if stage_idx == 0:  # Director initial planning
            return self.create_director_initial_prompt(query, language)
        
        if stage_idx == 1:  # Director plan review
            return self.create_critic_director_prompt(stages[0]["content"], query, language)
        
        if stage_idx == 2:  # Director final master plan
            return self.create_director_final_prompt(stages[0]["content"], stages[1]["content"], query, language)

        master_plan = stages[2]["content"]
        
        # Writer part writing
        if role == "writer" and "Revision" not in stages[stage_idx]["name"]:
            part_num = self._get_part_number(stage_idx)
            accumulated = '\n\n'.join(self.narrative_tracker.accumulated_content)
            return self.create_writer_prompt(part_num, master_plan, accumulated, 
                                           self.narrative_tracker.story_bible, language)
        
        # Part-specific critique
        if role.startswith("critic_part"):
            part_num = int(role.replace("critic_part", ""))
            # Find writer content for this part
            writer_content = stages[stage_idx-1]["content"]
            accumulated = '\n\n'.join(self.narrative_tracker.accumulated_content[:-1])
            return self.create_part_critic_prompt(part_num, writer_content, master_plan,
                                                accumulated, self.narrative_tracker.story_bible, language)
        
        # Writer revision
        if role == "writer" and "Revision" in stages[stage_idx]["name"]:
            part_num = self._get_part_number(stage_idx)
            original_content = stages[stage_idx-2]["content"]  # Original
            critic_feedback = stages[stage_idx-1]["content"]   # Critique
            return self.create_writer_revision_prompt(part_num, original_content, 
                                                    critic_feedback, language)
        
        # Final critique
        if role == "critic_final":
            complete_novel = NovelDatabase.get_writer_content(self.current_session_id)
            word_count = len(complete_novel.split())
            return self.create_final_critic_prompt(complete_novel, word_count,
                                                 self.narrative_tracker.story_bible, language)
        
        return ""

    def _get_part_number(self, stage_idx: int) -> Optional[int]:
        """Extract part number from stage index"""
        stage_name = UNIFIED_STAGES[stage_idx][1]
        match = re.search(r'Part (\d+)', stage_name)
        if match:
            return int(match.group(1))
        return None

    def _update_story_bible_from_content(self, content: str, part_num: int):
        """Auto-update story bible from content"""
        # Simple keyword-based extraction (more sophisticated NLP needed in reality)
        lines = content.split('\n')
        
        # Extract character names (words starting with capital letters)
        for line in lines:
            words = line.split()
            for word in words:
                if word and word[0].isupper() and len(word) > 1:
                    if word not in self.narrative_tracker.story_bible.characters:
                        self.narrative_tracker.story_bible.characters[word] = {
                            "first_appearance": part_num,
                            "traits": []
                        }

    def generate_literary_report(self, complete_novel: str, word_count: int, language: str) -> str:
        """Generate final literary evaluation report"""
        prompt = self.create_final_critic_prompt(complete_novel, word_count,
                                               self.narrative_tracker.story_bible, language)
        try:
            report = self.call_llm_sync([{"role": "user", "content": prompt}], 
                                      "critic_final", language)
            return report
        except Exception as e:
            logger.error(f"Final report generation failed: {e}")
            return "Error occurred during report generation"

class WebSearchIntegration:
    """Web search functionality"""
    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 = 3, language: str = "en") -> List[Dict]:
        if not self.enabled:
            return []
        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": self.brave_api_key
        }
        params = {
            "q": query,
            "count": count,
            "search_lang": "ko" if language == "Korean" else "en",
            "text_decorations": False,
            "safesearch": "moderate"
        }
        try:
            response = requests.get(self.search_url, headers=headers, params=params, timeout=10)
            response.raise_for_status()
            results = response.json().get("web", {}).get("results", [])
            return results
        except requests.exceptions.RequestException as e:
            logger.error(f"Web search API error: {e}")
            return []

    def extract_relevant_info(self, results: List[Dict], max_chars: int = 1500) -> str:
        if not results:
            return ""
        extracted = []
        total_chars = 0
        for i, result in enumerate(results[:3], 1):
            title = result.get("title", "")
            description = result.get("description", "")
            info = f"[{i}] {title}: {description}"
            if total_chars + len(info) < max_chars:
                extracted.append(info)
                total_chars += len(info)
            else:
                break
        return "\n".join(extracted)

class HFDatasetManager:
    """Manage theme data storage in HuggingFace dataset"""
    
    def __init__(self):
        self.token = os.getenv("HF_TOKEN")
        self.dataset_name = "novel-themes-library"
        self.username = None
        self.repo_id = None
        
        if self.token:
            try:
                self.api = HfApi()
                # Get username from token
                self.username = self.api.whoami(token=self.token)["name"]
                self.repo_id = f"{self.username}/{self.dataset_name}"
                
                # Create dataset repo if it doesn't exist
                try:
                    self.api.create_repo(
                        repo_id=self.repo_id,
                        token=self.token,
                        repo_type="dataset",
                        private=False,
                        exist_ok=True
                    )
                    logger.info(f"HF Dataset initialized: {self.repo_id}")
                except Exception as e:
                    logger.error(f"Error creating HF dataset: {e}")
                    
            except Exception as e:
                logger.error(f"HF authentication failed: {e}")
                self.token = None
    
    def save_themes_to_hf(self, themes_data: List[Dict]):
        """Save themes to HuggingFace dataset"""
        if not self.token or not themes_data:
            return False
            
        try:
            # Create temporary file
            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tmp_file:
                json.dump({
                    "themes": themes_data,
                    "last_updated": datetime.now().isoformat(),
                    "version": "1.0"
                }, tmp_file, ensure_ascii=False, indent=2)
                tmp_path = tmp_file.name
            
            # Upload to HF
            upload_file(
                path_or_fileobj=tmp_path,
                path_in_repo="themes_library.json",
                repo_id=self.repo_id,
                token=self.token,
                repo_type="dataset",
                commit_message=f"Update themes library - {len(themes_data)} themes"
            )
            
            # Clean up
            os.unlink(tmp_path)
            logger.info(f"Saved {len(themes_data)} themes to HF dataset")
            return True
            
        except Exception as e:
            logger.error(f"Error saving to HF dataset: {e}")
            return False
    
    def load_themes_from_hf(self) -> List[Dict]:
        """Load themes from HuggingFace dataset"""
        if not self.token:
            return []
            
        try:
            # Download file from HF
            file_path = hf_hub_download(
                repo_id=self.repo_id,
                filename="themes_library.json",
                token=self.token,
                repo_type="dataset"
            )
            
            # Load data
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
            themes = data.get("themes", [])
            logger.info(f"Loaded {len(themes)} themes from HF dataset")
            return themes
            
        except Exception as e:
            logger.warning(f"Error loading from HF dataset: {e}")
            return []



    def sync_with_local_db(self):
        """Sync HF dataset with local database"""
        if not self.token:
            return
            
        # Load from HF
        hf_themes = self.load_themes_from_hf()
        
        if hf_themes:
            # Get existing theme IDs from local DB
            local_theme_ids = set()
            with NovelDatabase.get_db() as conn:
                rows = conn.cursor().execute(
                    "SELECT theme_id FROM random_themes_library"
                ).fetchall()
                local_theme_ids = {row['theme_id'] for row in rows}
            
            # Add new themes from HF to local DB
            new_count = 0
            for theme in hf_themes:
                if theme.get('theme_id') not in local_theme_ids:
                    try:
                        # Ensure tags and metadata are JSON strings
                        tags_data = theme.get('tags', [])
                        if isinstance(tags_data, list):
                            tags_json = json.dumps(tags_data, ensure_ascii=False)
                        else:
                            tags_json = tags_data if isinstance(tags_data, str) else '[]'
                        
                        metadata_data = theme.get('metadata', {})
                        if isinstance(metadata_data, dict):
                            metadata_json = json.dumps(metadata_data, ensure_ascii=False)
                        else:
                            metadata_json = metadata_data if isinstance(metadata_data, str) else '{}'
                        
                        with NovelDatabase.get_db() as conn:
                            conn.cursor().execute('''
                                INSERT INTO random_themes_library 
                                (theme_id, theme_text, language, title, opening_sentence,
                                 protagonist, conflict, philosophical_question, generated_at,
                                 view_count, used_count, tags, metadata)
                                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                            ''', (
                                theme.get('theme_id'),
                                theme.get('theme_text'),
                                theme.get('language'),
                                theme.get('title', ''),
                                theme.get('opening_sentence', ''),
                                theme.get('protagonist', ''),
                                theme.get('conflict', ''),
                                theme.get('philosophical_question', ''),
                                theme.get('generated_at'),
                                theme.get('view_count', 0),
                                theme.get('used_count', 0),
                                tags_json,
                                metadata_json
                            ))
                            conn.commit()
                        new_count += 1
                    except Exception as e:
                        logger.error(f"Error adding theme {theme.get('theme_id')}: {e}")
            
            if new_count > 0:
                logger.info(f"Added {new_count} new themes from HF dataset")
    
    def backup_to_hf(self):
        """Backup all local themes to HF dataset"""
        if not self.token:
            return
            
        # Get all themes from local DB
        themes = NovelDatabase.get_random_themes_library(limit=1000)
        
        if themes:
            # Convert Row objects to dicts and ensure all data is serializable
            themes_data = []
            for theme in themes:
                theme_dict = dict(theme)
                # Parse tags and metadata from JSON strings
                if isinstance(theme_dict.get('tags'), str):
                    try:
                        theme_dict['tags'] = json.loads(theme_dict['tags'])
                    except:
                        theme_dict['tags'] = []
                else:
                    theme_dict['tags'] = theme_dict.get('tags', [])
                    
                if isinstance(theme_dict.get('metadata'), str):
                    try:
                        theme_dict['metadata'] = json.loads(theme_dict['metadata'])
                    except:
                        theme_dict['metadata'] = {}
                else:
                    theme_dict['metadata'] = theme_dict.get('metadata', {})
                    
                themes_data.append(theme_dict)
            
            self.save_themes_to_hf(themes_data)
    
 
# --- Utility functions ---
def process_query(query: str, language: str, session_id: Optional[str] = None) -> Generator[Tuple[str, str, str, str, str], None, None]:
    """Main query processing function"""
    if not query.strip():
        yield "", "", "❌ Please enter a theme.", session_id, ""
        return
    
    system = UnifiedLiterarySystem()
    stages_markdown = ""
    novel_content = ""
    novel_text = ""  # 실제 텍스트 저장용
    
    for status, stages, current_session_id in system.process_novel_stream(query, language, session_id):
        stages_markdown = format_stages_display(stages)
        
        # Get final novel content
        if stages and all(s.get("status") == "complete" for s in stages[-10:]):
            novel_text = NovelDatabase.get_writer_content(current_session_id)  # 원본 텍스트
            novel_content = format_novel_display(novel_text)  # 포맷된 디스플레이
        
        yield stages_markdown, novel_content, status or "🔄 Processing...", current_session_id, novel_text


def get_active_sessions(language: str) -> List[str]:
    """Get active session list"""
    sessions = NovelDatabase.get_active_sessions()
    return [f"{s['session_id'][:8]}... - {s['user_query'][:50]}... ({s['created_at']}) [{s['total_words']:,} words]" 
            for s in sessions]

def auto_recover_session(language: str) -> Tuple[Optional[str], str]:
    """Auto-recover recent session"""
    sessions = NovelDatabase.get_active_sessions()
    if sessions:
        latest_session = sessions[0]
        return latest_session['session_id'], f"Session {latest_session['session_id'][:8]}... recovered"
    return None, "No session to recover."

def resume_session(session_id: str, language: str) -> Generator[Tuple[str, str, str, str, str], None, None]:
    """Resume session"""
    if not session_id:
        yield "", "", "❌ No session ID.", session_id, ""
        return
    
    if "..." in session_id:
        session_id = session_id.split("...")[0]
    
    session = NovelDatabase.get_session(session_id)
    if not session:
        yield "", "", "❌ Session not found.", None, ""
        return
    
    yield from process_query(session['user_query'], session['language'], session_id)



def format_stages_display(stages: List[Dict]) -> str:
    """Stage progress display - For single writer system"""
    markdown = "## 🎬 Progress Status\n\n"
    
    # Calculate total word count (writer stages only)
    total_words = sum(s.get('word_count', 0) for s in stages 
                     if s.get('name', '').startswith('✍️ Writer:') and 'Revision' in s.get('name', ''))
    markdown += f"**Total Word Count: {total_words:,} / {TARGET_WORDS:,}**\n\n"
    
    # Progress summary
    completed_parts = sum(1 for s in stages 
                         if 'Revision' in s.get('name', '') and s.get('status') == 'complete')
    markdown += f"**Completed Parts: {completed_parts} / 10**\n\n"
    
    # Average narrative momentum
    momentum_scores = [s.get('momentum', 0) for s in stages if s.get('momentum', 0) > 0]
    if momentum_scores:
        avg_momentum = sum(momentum_scores) / len(momentum_scores)
        markdown += f"**Average Narrative Momentum: {avg_momentum:.1f} / 10**\n\n"
    
    markdown += "---\n\n"
    
    # Display each stage
    current_part = 0
    for i, stage in enumerate(stages):
        status_icon = "✅" if stage['status'] == 'complete' else "🔄" if stage['status'] == 'active' else "⏳"
        
        # Add part divider
        if 'Part' in stage.get('name', '') and 'Critic' not in stage.get('name', ''):
            part_match = re.search(r'Part (\d+)', stage['name'])
            if part_match:
                new_part = int(part_match.group(1))
                if new_part != current_part:
                    current_part = new_part
                    markdown += f"\n### 📚 Part {current_part}\n\n"
        
        markdown += f"{status_icon} **{stage['name']}**"
        
        if stage.get('word_count', 0) > 0:
            markdown += f" ({stage['word_count']:,} words)"
        
        if stage.get('momentum', 0) > 0:
            markdown += f" [Momentum: {stage['momentum']:.1f}/10]"
        
        markdown += "\n"
        
        if stage['content'] and stage['status'] == 'complete':
            # Adjust preview length by role
            preview_length = 300 if 'writer' in stage.get('name', '').lower() else 200
            preview = stage['content'][:preview_length] + "..." if len(stage['content']) > preview_length else stage['content']
            markdown += f"> {preview}\n\n"
        elif stage['status'] == 'active':
            markdown += "> *Writing...*\n\n"
    
    return markdown

def format_novel_display(novel_text: str) -> str:
    """Display novel content - Enhanced part separation"""
    if not novel_text:
        return "No completed content yet."
    
    formatted = "# 📖 Completed Novel\n\n"
    
    # Display word count
    word_count = len(novel_text.split())
    formatted += f"**Total Length: {word_count:,} words (Target: {TARGET_WORDS:,} words)**\n\n"
    
    # Achievement rate
    achievement = (word_count / TARGET_WORDS) * 100
    formatted += f"**Achievement Rate: {achievement:.1f}%**\n\n"
    formatted += "---\n\n"
    
    # Display each part separately
    parts = novel_text.split('\n\n')
    
    for i, part in enumerate(parts):
        if part.strip():
            # Add part title
            if i < len(NARRATIVE_PHASES):
                formatted += f"## {NARRATIVE_PHASES[i]}\n\n"
            
            formatted += f"{part}\n\n"
            
            # Part divider
            if i < len(parts) - 1:
                formatted += "---\n\n"
    
    return formatted

def export_to_docx(content: str, filename: str, language: str, session_id: str) -> str:
    """Export to DOCX file - Korean standard book format"""
    try:
        doc = Document()
        
        # Korean standard book format (152mm x 225mm)
        section = doc.sections[0]
        section.page_height = Mm(225)  # 225mm
        section.page_width = Mm(152)   # 152mm
        section.top_margin = Mm(20)    # Top margin 20mm
        section.bottom_margin = Mm(20) # Bottom margin 20mm
        section.left_margin = Mm(20)   # Left margin 20mm
        section.right_margin = Mm(20)  # Right margin 20mm
        
        # Generate title from session info
        session = NovelDatabase.get_session(session_id)
        
        # Title generation function
        def generate_title(user_query: str, content_preview: str) -> str:
            """Generate title based on theme and content"""
            if len(user_query) < 20:
                return user_query
            else:
                keywords = user_query.split()[:5]
                return " ".join(keywords)
        
        # Title page
        title = generate_title(session["user_query"], content[:500]) if session else "Untitled"
        
        # Title style settings
        title_para = doc.add_paragraph()
        title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        title_para.paragraph_format.space_before = Pt(100)
        
        title_run = title_para.add_run(title)
        if language == "Korean":
            title_run.font.name = 'Batang'
            title_run._element.rPr.rFonts.set(qn('w:eastAsia'), 'Batang')
        else:
            title_run.font.name = 'Times New Roman'
        title_run.font.size = Pt(20)
        title_run.bold = True
        
        # Page break
        doc.add_page_break()
        
        # Body style settings
        style = doc.styles['Normal']
        if language == "Korean":
            style.font.name = 'Batang'
            style._element.rPr.rFonts.set(qn('w:eastAsia'), 'Batang')
        else:
            style.font.name = 'Times New Roman'
        style.font.size = Pt(10.5)
        style.paragraph_format.line_spacing = 1.8
        style.paragraph_format.space_after = Pt(0)
        style.paragraph_format.first_line_indent = Mm(10)
        
        # Clean content function
        def clean_content(text: str) -> str:
            """Remove unnecessary markdown, part numbers, etc."""
            patterns_to_remove = [
                r'^#{1,6}\s+.*',          # Markdown headers
                r'^\*\*.*\*\*',           # Bold text
                r'^Part\s*\d+.*',         # Part numbers
                r'^\d+\.\s+.*:.*',        # Numbered lists
                r'^---+',                 # Dividers
                r'^\s*\[.*\]\s*',         # Brackets
            ]
            
            lines = text.split('\n')
            cleaned_lines = []
            
            for line in lines:
                if not line.strip():
                    cleaned_lines.append('')
                    continue
                
                skip_line = False
                for pattern in patterns_to_remove:
                    if re.match(pattern, line.strip(), re.MULTILINE):
                        skip_line = True
                        break
                
                if not skip_line:
                    cleaned_line = line
                    cleaned_line = re.sub(r'\*\*(.*?)\*\*', r'\1', cleaned_line)
                    cleaned_line = re.sub(r'\*(.*?)\*', r'\1', cleaned_line)
                    cleaned_line = re.sub(r'`(.*?)`', r'\1', cleaned_line)
                    cleaned_lines.append(cleaned_line.strip())
            
            final_lines = []
            prev_empty = False
            for line in cleaned_lines:
                if not line:
                    if not prev_empty:
                        final_lines.append('')
                    prev_empty = True
                else:
                    final_lines.append(line)
                    prev_empty = False
            
            return '\n'.join(final_lines)
        
        # Clean content
        cleaned_content = clean_content(content)
        
        # Add body text
        paragraphs = cleaned_content.split('\n')
        for para_text in paragraphs:
            if para_text.strip():
                para = doc.add_paragraph(para_text.strip())
                for run in para.runs:
                    if language == "Korean":
                        run.font.name = 'Batang'
                        run._element.rPr.rFonts.set(qn('w:eastAsia'), 'Batang')
                    else:
                        run.font.name = 'Times New Roman'
            else:
                doc.add_paragraph()
        
        # Create temporary file with proper handling
        with tempfile.NamedTemporaryFile(mode='wb', suffix='.docx', delete=False) as tmp_file:
            doc.save(tmp_file)
            temp_path = tmp_file.name
        
        return temp_path
        
    except Exception as e:
        logger.error(f"DOCX export error: {str(e)}")
        raise e


def download_novel(novel_text: str, format_type: str, language: str, session_id: str) -> Optional[str]:
    """Generate novel download file - FIXED VERSION"""
    if not novel_text or not session_id:
        logger.error("Missing novel_text or session_id")
        return None
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"novel_{session_id[:8]}_{timestamp}"
    
    try:
        if format_type == "DOCX" and DOCX_AVAILABLE:
            # Use the fixed export_to_docx function
            return export_to_docx(novel_text, filename, language, session_id)
        else:
            # For TXT format
            return export_to_txt(novel_text, filename)
    except Exception as e:
        logger.error(f"File generation failed: {e}")
        return None


# In the Gradio interface, update the download handler:
def handle_download(format_type, language, session_id, novel_text):
    """Fixed download handler with better error handling and debugging"""
    logger.info(f"Download attempt - Session ID: {session_id}, Format: {format_type}")
    logger.info(f"Novel text length: {len(novel_text) if novel_text else 0}")
    logger.info(f"Novel text preview: {novel_text[:100] if novel_text else 'None'}")
    
    if not session_id:
        logger.error("No session ID provided")
        return gr.update(visible=False, value=None)
    
    if not novel_text or novel_text.strip() == "" or novel_text == "*Your completed novel will appear here, ready to be read and cherished...*":
        logger.error(f"No novel content to download. Content: '{novel_text[:50] if novel_text else 'None'}'")
        return gr.update(visible=False, value=None)
    
    try:
        file_path = download_novel(novel_text, format_type, language, session_id)
        if file_path and os.path.exists(file_path):
            logger.info(f"File created successfully: {file_path}")
            return gr.update(value=file_path, visible=True)
        else:
            logger.error("File path not created or doesn't exist")
            return gr.update(visible=False, value=None)
    except Exception as e:
        logger.error(f"Download handler error: {str(e)}")
        return gr.update(visible=False, value=None)

# Also add cleanup function for temporary files
def cleanup_temp_files():
    """Clean up old temporary files"""
    temp_dir = tempfile.gettempdir()
    pattern = os.path.join(temp_dir, "novel_*.docx")
    
    for file_path in glob.glob(pattern):
        try:
            # Delete files older than 1 hour
            if os.path.getmtime(file_path) < time.time() - 3600:
                os.unlink(file_path)
        except:
            pass


def export_to_txt(content: str, filename: str) -> str:
    """Export to TXT file"""
    filepath = f"{filename}.txt"
    with open(filepath, 'w', encoding='utf-8') as f:
        # Header
        f.write("=" * 80 + "\n")
        f.write(f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Total word count: {len(content.split()):,} words\n")
        f.write("=" * 80 + "\n\n")
        
        # Body
        f.write(content)
        
        # Footer
        f.write("\n\n" + "=" * 80 + "\n")
        f.write("AI Literary Creation System v2.0\n")
        f.write("=" * 80 + "\n")
    
    return filepath

def generate_random_theme(language="English"):
    """Generate a coherent and natural novel theme using LLM"""
    try:
        # JSON 파일 로드
        json_path = Path("novel_themes.json")
        if not json_path.exists():
            print("[WARNING] novel_themes.json not found, using built-in data")
            # 기본 데이터 정의 - 더 현실적인 테마로 수정
            themes_data = {
                "themes": ["family secrets", "career transition", "lost love", "friendship test", "generational conflict"],
                "characters": ["middle-aged teacher", "retiring doctor", "single parent", "immigrant artist", "war veteran"],
                "hooks": ["unexpected inheritance", "old diary discovery", "chance reunion", "life-changing diagnosis", "sudden job loss"],
                "questions": ["What defines family?", "Can people truly change?", "What is worth sacrificing?", "How do we forgive?"]
            }
        else:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # 가중치 기반 필터링 - 현실적인 테마 우선
                realistic_themes = []
                for theme_key, theme_data in data.get('core_themes', {}).items():
                    weight = theme_data.get('weight', 0.1)
                    # 현실적인 테마에 더 높은 가중치
                    if any(word in theme_key for word in ['family', 'love', 'work', 'memory', 'identity', 'aging']):
                        weight *= 1.5
                    elif any(word in theme_key for word in ['digital', 'extinction', 'apocalypse', 'quantum']):
                        weight *= 0.5
                    realistic_themes.append((theme_key, weight))
                
                # 가중치 기반 선택
                themes = [t[0] for t in sorted(realistic_themes, key=lambda x: x[1], reverse=True)[:10]]
                
                themes_data = {
                    "themes": themes if themes else ["family secrets", "career crisis", "lost love"],
                    "characters": [],
                    "hooks": [],
                    "questions": []
                }
                
                # Extract realistic data
                for char_data in data.get('characters', {}).values():
                    for variation in char_data.get('variations', []):
                        # 현실적인 캐릭터 필터링
                        if not any(word in variation.lower() for word in ['cyborg', 'quantum', 'binary', 'extinct']):
                            themes_data["characters"].append(variation)
                
                for hook_list in data.get('narrative_hooks', {}).values():
                    for hook in hook_list:
                        # 현실적인 사건 필터링
                        if not any(word in hook.lower() for word in ['download', 'digital', 'algorithm', 'corporate subscription']):
                            themes_data["hooks"].append(hook)
                
                for phil_data in data.get('philosophies', {}).values():
                    themes_data["questions"].extend(phil_data.get('core_questions', []))
        
        # 기본값 설정
        if not themes_data["characters"]:
            themes_data["characters"] = ["struggling artist", "retired teacher", "young mother", "elderly caregiver", "small business owner"]
        if not themes_data["hooks"]:
            themes_data["hooks"] = ["discovering family secret", "unexpected reunion", "facing illness", "losing home", "finding old letters"]
        if not themes_data["questions"]:
            themes_data["questions"] = ["What makes a family?", "How do we find meaning?", "Can we escape our past?", "What legacy do we leave?"]
        
        # Random selection
        import secrets
        theme = secrets.choice(themes_data["themes"])
        character = secrets.choice(themes_data["characters"])
        hook = secrets.choice(themes_data["hooks"])
        question = secrets.choice(themes_data["questions"])
        
        # 언어별 프롬프트 - 톤과 스타일 섹션 제거
        if language == "Korean":
            # 한국어 번역 및 자연스러운 표현
            theme_kr = translate_theme_naturally(theme, "theme")
            character_kr = translate_theme_naturally(character, "character")
            hook_kr = translate_theme_naturally(hook, "hook")
            question_kr = translate_theme_naturally(question, "question")
            
            prompt = f"""다음 요소들을 사용하여 현실적이고 공감가능한 소설 주제를 생성하세요:

주제: {theme_kr}
인물: {character_kr}
사건: {hook_kr}
핵심 질문: {question_kr}

요구사항:
1. 현대 한국 사회에서 일어날 수 있는 현실적인 이야기
2. 보편적으로 공감할 수 있는 인물과 상황
3. 구체적이고 생생한 배경 설정
4. 깊이 있는 심리 묘사가 가능한 갈등

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

[한 문장으로 된 매력적인 첫 문장]

주인공은 [구체적인 상황의 인물]입니다.
[핵심 사건]을 계기로 [내적 갈등]에 직면하게 되고,
결국 [철학적 질문]에 대한 답을 찾아가는 여정을 그립니다."""
            
        else:
            prompt = f"""Generate a realistic and relatable novel theme using these elements:

Theme: {theme}
Character: {character}
Event: {hook}
Core Question: {question}

Requirements:
1. A story that could happen in contemporary society
2. Universally relatable characters and situations
3. Specific and vivid settings
4. Conflicts allowing deep psychological exploration

Write concisely in this format:

[One compelling opening sentence]

The protagonist is [character in specific situation].
Through [key event], they face [internal conflict],
ultimately embarking on a journey to answer [philosophical question]."""

        # Use the UnifiedLiterarySystem's LLM to generate coherent theme
        system = UnifiedLiterarySystem()
        
        # Call LLM synchronously for theme generation
        messages = [{"role": "user", "content": prompt}]
        generated_theme = system.call_llm_sync(messages, "director", language)
        
        # Extract metadata for database storage
        metadata = extract_theme_metadata(generated_theme, language)
        metadata.update({
            'original_theme': theme,
            'original_character': character,
            'original_hook': hook,
            'original_question': question
        })
        
        # Save to database
        theme_id = save_random_theme_with_hf(generated_theme, language, metadata)
        logger.info(f"Saved random theme with ID: {theme_id}")
        
        # 톤과 스타일 섹션 제거 - 불필요한 반복 내용 삭제
        if "**톤과 스타일:**" in generated_theme or "**Tone and Style:**" in generated_theme:
            lines = generated_theme.split('\n')
            filtered_lines = []
            skip = False
            for line in lines:
                if "톤과 스타일" in line or "Tone and Style" in line:
                    skip = True
                elif skip and (line.strip() == "" or line.startswith("**")):
                    skip = False
                if not skip:
                    filtered_lines.append(line)
            generated_theme = '\n'.join(filtered_lines).strip()
        
        return generated_theme
        
    except Exception as e:
        logger.error(f"Theme generation error: {str(e)}")
        # Fallback to simple realistic themes
        fallback_themes = {
            "Korean": [
                """"아버지가 돌아가신 날, 나는 그가 평생 숨겨온 또 다른 가족의 존재를 알게 되었다."

주인공은 평범한 회사원으로 살아온 40대 여성입니다.
아버지의 장례식에서 낯선 여인과 그녀의 딸을 만나게 되면서 가족의 의미에 대해 다시 생각하게 되고,
결국 진정한 가족이란 무엇인지에 대한 답을 찾아가는 여정을 그립니다.""",
                
                """"서른 년간 가르친 학교에서 나온 날, 처음으로 내가 누구인지 몰랐다."

주인공은 정년퇴직을 맞은 고등학교 국어 교사입니다.
갑작스러운 일상의 공백 속에서 잊고 지냈던 젊은 날의 꿈을 마주하게 되고,
결국 남은 인생에서 무엇을 할 것인가에 대한 답을 찾아가는 여정을 그립니다."""
            ],
            "English": [
                """"The day my father died, I discovered he had another family he'd hidden all his life."

The protagonist is a woman in her 40s who has lived as an ordinary office worker.
Through meeting a strange woman and her daughter at her father's funeral, she confronts what family truly means,
ultimately embarking on a journey to answer what constitutes a real family.""",
                
                """"The day I left the school where I'd taught for thirty years, I didn't know who I was anymore."

The protagonist is a high school literature teacher facing retirement.
Through the sudden emptiness of daily life, they confront long-forgotten dreams of youth,
ultimately embarking on a journey to answer what to do with the remaining years."""
            ]
        }
        
        import secrets
        return secrets.choice(fallback_themes.get(language, fallback_themes["English"]))

def translate_theme_naturally(text, category):
    """자연스러운 한국어 번역"""
    translations = {
        # 테마
        "family secrets": "가족의 비밀",
        "career transition": "인생의 전환점",
        "lost love": "잃어버린 사랑",
        "friendship test": "우정의 시험",
        "generational conflict": "세대 간 갈등",
        "digital extinction": "디지털 시대의 소외",
        "sensory revolution": "감각의 혁명",
        "temporal paradox": "시간의 역설",
        
        # 캐릭터
        "struggling artist": "생활고에 시달리는 예술가",
        "retired teacher": "은퇴한 교사",
        "young mother": "젊은 엄마",
        "elderly caregiver": "노인을 돌보는 간병인",
        "small business owner": "작은 가게 주인",
        "middle-aged teacher": "중년의 교사",
        "retiring doctor": "은퇴를 앞둔 의사",
        "single parent": "혼자 아이를 키우는 부모",
        "immigrant artist": "이민자 예술가",
        "war veteran": "전쟁 참전용사",
        "last person who dreams without ads": "광고 없이 꿈꾸는 마지막 사람",
        "memory trader": "기억 거래상",
        
        # 사건
        "discovering family secret": "가족의 비밀을 발견하다",
        "unexpected reunion": "예상치 못한 재회",
        "facing illness": "질병과 마주하다",
        "losing home": "집을 잃다",
        "finding old letters": "오래된 편지를 발견하다",
        "unexpected inheritance": "뜻밖의 유산",
        "old diary discovery": "오래된 일기장 발견",
        "chance reunion": "우연한 재회",
        "life-changing diagnosis": "인생을 바꾸는 진단",
        "sudden job loss": "갑작스러운 실직",
        "discovers their memories belong to a corporate subscription": "기억이 기업 서비스의 일부임을 발견하다",
        
        # 질문
        "What makes a family?": "가족이란 무엇인가?",
        "How do we find meaning?": "우리는 어떻게 의미를 찾는가?",
        "Can we escape our past?": "과거로부터 벗어날 수 있는가?",
        "What legacy do we leave?": "우리는 어떤 유산을 남기는가?",
        "What defines family?": "무엇이 가족을 정의하는가?",
        "Can people truly change?": "사람은 정말 변할 수 있는가?",
        "What is worth sacrificing?": "무엇을 위해 희생할 가치가 있는가?",
        "How do we forgive?": "우리는 어떻게 용서하는가?",
        "What remains human when humanity is optional?": "인간성이 선택사항일 때 무엇이 인간으로 남는가?"
    }
    
    # 먼저 정확한 매칭 시도
    if text in translations:
        return translations[text]
    
    # 부분 매칭 시도
    text_lower = text.lower()
    for key, value in translations.items():
        if key.lower() in text_lower or text_lower in key.lower():
            return value
    
    # 번역이 없으면 원문 반환
    return text

def extract_theme_metadata(theme_text: str, language: str) -> Dict[str, Any]:
    """Extract metadata from generated theme text"""
    metadata = {
        'title': '',
        'opening_sentence': '',
        'protagonist': '',
        'conflict': '',
        'philosophical_question': '',
        'tags': []
    }
    
    lines = theme_text.split('\n')
    
    # Extract opening sentence (usually in quotes)
    for line in lines:
        if '"' in line or '"' in line or '「' in line:
            # Extract text between quotes
            import re
            quotes = re.findall(r'["""「](.*?)["""」]', line)
            if quotes:
                metadata['opening_sentence'] = quotes[0]
                break
    
    # Extract other elements based on patterns
    for i, line in enumerate(lines):
        line = line.strip()
        
        # Title extraction (if exists)
        if i == 0 and not any(quote in line for quote in ['"', '"', '「']):
            metadata['title'] = line.replace('**', '').strip()
        
        # Protagonist
        if any(marker in line for marker in ['protagonist is', '주인공은', 'The protagonist']):
            metadata['protagonist'] = line.split('is' if 'is' in line else '은')[-1].strip().rstrip('.')
        
        # Conflict/Event
        if any(marker in line for marker in ['Through', '통해', '계기로', 'face']):
            metadata['conflict'] = line
        
        # Question
        if any(marker in line for marker in ['answer', '답을', 'question', '질문']):
            metadata['philosophical_question'] = line
    
    # Generate tags based on content
    tag_keywords = {
        'family': ['family', '가족', 'father', '아버지', 'mother', '어머니'],
        'love': ['love', '사랑', 'relationship', '관계'],
        'death': ['death', '죽음', 'died', '돌아가신'],
        'memory': ['memory', '기억', 'remember', '추억'],
        'identity': ['identity', '정체성', 'who am I', '누구인지'],
        'work': ['work', '일', 'career', '직업', 'retirement', '은퇴'],
        'aging': ['aging', '노화', 'old', '늙은', 'elderly', '노인']
    }
    
    theme_lower = theme_text.lower()
    for tag, keywords in tag_keywords.items():
        if any(keyword in theme_lower for keyword in keywords):
            metadata['tags'].append(tag)
    
    return metadata

def save_random_theme_with_hf(theme_text: str, language: str, metadata: Dict[str, Any]) -> str:
    """Save randomly generated theme to library and HF dataset"""
    theme_id = hashlib.md5(f"{theme_text}{datetime.now()}".encode()).hexdigest()[:12]
    
    # Extract components from theme text
    title = metadata.get('title', '')
    opening_sentence = metadata.get('opening_sentence', '')
    protagonist = metadata.get('protagonist', '')
    conflict = metadata.get('conflict', '')
    philosophical_question = metadata.get('philosophical_question', '')
    tags = json.dumps(metadata.get('tags', []))
    
    with NovelDatabase.get_db() as conn:
        conn.cursor().execute('''
            INSERT INTO random_themes_library 
            (theme_id, theme_text, language, title, opening_sentence, 
             protagonist, conflict, philosophical_question, tags, metadata)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (theme_id, theme_text, language, title, opening_sentence,
              protagonist, conflict, philosophical_question, tags, 
              json.dumps(metadata)))
        conn.commit()
    
    # Backup to HF dataset
    if 'hf_manager' in globals() and hf_manager.token:
        try:
            hf_manager.backup_to_hf()
            logger.info(f"Theme {theme_id} backed up to HF dataset")
        except Exception as e:
            logger.error(f"Failed to backup theme to HF: {e}")
    
    return theme_id

def format_theme_card(theme_data: Dict, language: str) -> str:
    """Format theme data as a card for display with scrollable content"""
    theme_id = theme_data.get('theme_id', '')
    theme_text = theme_data.get('theme_text', '')
    generated_at = theme_data.get('generated_at', '')
    view_count = theme_data.get('view_count', 0)
    used_count = theme_data.get('used_count', 0)
    tags = json.loads(theme_data.get('tags', '[]')) if isinstance(theme_data.get('tags'), str) else theme_data.get('tags', [])
    
    # Format timestamp
    if generated_at:
        try:
            dt = datetime.fromisoformat(generated_at.replace(' ', 'T'))
            time_str = dt.strftime('%Y-%m-%d %H:%M')
        except:
            time_str = generated_at
    else:
        time_str = ""
    
    # Create tag badges
    tag_badges = ' '.join([f'<span class="theme-tag">{tag}</span>' for tag in tags])
    
    # Format theme text with line breaks
    formatted_text = theme_text.replace('\n', '<br>')
    
    # Create card HTML with scrollable content - Simplified version
    card_html = f"""
<div class="theme-card" data-theme-id="{theme_id}">
    <div class="theme-card-header">
        <span class="theme-id">#{theme_id[:8]}</span>
        <span class="theme-timestamp">{time_str}</span>
    </div>
    <div class="theme-card-content">
        <div class="theme-full-text">{formatted_text}</div>
        <div class="theme-tags">{tag_badges}</div>
    </div>
    <div class="theme-card-footer">
        <div class="theme-stats">
            <span class="theme-stat">👁️ {view_count}</span>
            <span class="theme-stat">📖 {used_count}</span>
        </div>
        <button class="theme-action-btn use-btn" onclick="handleThemeUse('{theme_id}')">
            {'사용하기' if language == 'Korean' else 'Use This'}
        </button>
    </div>
</div>"""
    
    return card_html

def get_theme_library_display(language: str = None, search_query: str = "") -> str:
    """Get formatted display of theme library"""
    themes = NovelDatabase.get_random_themes_library(language, limit=50)
    
    if not themes:
        empty_msg = {
            "Korean": "아직 생성된 테마가 없습니다. 랜덤 버튼을 눌러 첫 테마를 만들어보세요!",
            "English": "No themes generated yet. Click the Random button to create your first theme!"
        }
        return f'<div class="empty-library">{empty_msg.get(language, empty_msg["English"])}</div>'
    
    # Filter by search query if provided
    if search_query:
        search_lower = search_query.lower()
        themes = [t for t in themes if search_lower in t.get('theme_text', '').lower()]
    
    # Statistics
    total_themes = len(themes)
    total_views = sum(t.get('view_count', 0) for t in themes)
    total_uses = sum(t.get('used_count', 0) for t in themes)
    
    stats_html = f"""
<div class="library-stats">
    <div class="stat-item">
        <span class="stat-label">{'총 테마' if language == 'Korean' else 'Total Themes'}</span>
        <span class="stat-value">{total_themes}</span>
    </div>
    <div class="stat-item">
        <span class="stat-label">{'총 조회수' if language == 'Korean' else 'Total Views'}</span>
        <span class="stat-value">{total_views}</span>
    </div>
    <div class="stat-item">
        <span class="stat-label">{'총 사용수' if language == 'Korean' else 'Total Uses'}</span>
        <span class="stat-value">{total_uses}</span>
    </div>
</div>"""
    
    # Theme cards
    cards_html = '<div class="theme-cards-grid">'
    for theme in themes:
        cards_html += format_theme_card(theme, language)
    cards_html += '</div>'
    
    # JavaScript for interactions
    js_script = """
<script>
function handleThemeUse(themeId) {
    // Set values to hidden inputs
    const themeInput = document.querySelector('#selected_theme_for_action textarea');
    const actionInput = document.querySelector('#action_type textarea');
    
    if (themeInput && actionInput) {
        themeInput.value = themeId;
        actionInput.value = 'use';
        
        // Trigger the hidden button
        const triggerBtn = document.querySelector('#trigger_action button');
        if (triggerBtn) {
            triggerBtn.click();
        }
    }
}
</script>"""
    
    return stats_html + cards_html + js_script

# CSS styles
custom_css = """
/* Global container - Light paper background */
.gradio-container {
    background: linear-gradient(135deg, #faf8f3 0%, #f5f2e8 50%, #f0ebe0 100%);
    min-height: 100vh;
    font-family: 'Georgia', 'Times New Roman', serif;
}

/* Main header - Classic book cover feel */
.main-header {
    background: linear-gradient(145deg, #ffffff 0%, #fdfcf8 100%);
    backdrop-filter: blur(10px);
    padding: 45px;
    border-radius: 20px;
    margin-bottom: 35px;
    text-align: center;
    color: #3d2914;
    border: 1px solid #e8dcc6;
    box-shadow: 0 10px 30px rgba(139, 69, 19, 0.08), 
                0 5px 15px rgba(139, 69, 19, 0.05),
                inset 0 1px 2px rgba(255, 255, 255, 0.9);
    position: relative;
    overflow: hidden;
}

/* Book spine decoration */
.main-header::before {
    content: '';
    position: absolute;
    left: 50px;
    top: 0;
    bottom: 0;
    width: 3px;
    background: linear-gradient(180deg, #d4a574 0%, #c19656 50%, #d4a574 100%);
    box-shadow: 1px 0 2px rgba(0, 0, 0, 0.1);
}

.header-title {
    font-size: 3.2em;
    margin-bottom: 20px;
    font-weight: 700;
    color: #2c1810;
    text-shadow: 2px 2px 4px rgba(139, 69, 19, 0.1);
    font-family: 'Playfair Display', 'Georgia', serif;
    letter-spacing: -0.5px;
}

.header-description {
    font-size: 0.95em;
    color: #5a453a;
    line-height: 1.7;
    margin-top: 25px;
    text-align: justify;
    max-width: 920px;
    margin-left: auto;
    margin-right: auto;
    font-family: 'Georgia', serif;
}

.badges-container {
    display: flex;
    justify-content: center;
    gap: 12px;
    margin-top: 25px;
    margin-bottom: 25px;
}

/* Progress notes - Manuscript notes style */
.progress-note {
    background: linear-gradient(135deg, #fff9e6 0%, #fff5d6 100%);
    border-left: 4px solid #d4a574;
    padding: 22px 30px;
    margin: 25px auto;
    border-radius: 12px;
    color: #5a453a;
    max-width: 820px;
    font-weight: 500;
    box-shadow: 0 4px 12px rgba(212, 165, 116, 0.15);
    position: relative;
}

/* Handwritten note effect */
.progress-note::after {
    content: '📌';
    position: absolute;
    top: -10px;
    right: 20px;
    font-size: 24px;
    transform: rotate(15deg);
}

.warning-note {
    background: #fef3e2;
    border-left: 4px solid #f6b73c;
    padding: 18px 25px;
    margin: 20px auto;
    border-radius: 10px;
    color: #7a5c00;
    max-width: 820px;
    font-size: 0.92em;
    box-shadow: 0 3px 10px rgba(246, 183, 60, 0.15);
}

/* Input section - Writing desk feel */
.input-section {
    background: linear-gradient(145deg, #ffffff 0%, #fcfaf7 100%);
    backdrop-filter: blur(10px);
    padding: 30px;
    border-radius: 16px;
    margin-bottom: 28px;
    border: 1px solid #e8dcc6;
    box-shadow: 0 6px 20px rgba(139, 69, 19, 0.06),
                inset 0 1px 3px rgba(255, 255, 255, 0.8);
}

/* Session section - File cabinet style */
.session-section {
    background: linear-gradient(145deg, #f8f4ed 0%, #f3ede2 100%);
    backdrop-filter: blur(8px);
    padding: 22px;
    border-radius: 14px;
    margin-top: 28px;
    color: #3d2914;
    border: 1px solid #ddd0b8;
    box-shadow: inset 0 2px 4px rgba(139, 69, 19, 0.08);
}

/* Display areas - Clean manuscript pages */
#stages-display {
    background: linear-gradient(to bottom, #ffffff 0%, #fdfcfa 100%);
    padding: 35px 40px;
    border-radius: 16px;
    max-height: 680px;
    overflow-y: auto;
    box-shadow: 0 8px 25px rgba(139, 69, 19, 0.08),
                inset 0 1px 3px rgba(255, 255, 255, 0.9);
    color: #3d2914;
    border: 1px solid #e8dcc6;
    font-family: 'Georgia', serif;
    line-height: 1.8;
}

#novel-output {
    background: linear-gradient(to bottom, #ffffff 0%, #fdfcfa 100%);
    padding: 45px 50px;
    border-radius: 16px;
    max-height: 780px;
    overflow-y: auto;
    box-shadow: 0 10px 30px rgba(139, 69, 19, 0.1),
                inset 0 1px 3px rgba(255, 255, 255, 0.9);
    color: #2c1810;
    line-height: 2.1;
    font-size: 1.05em;
    border: 1px solid #e8dcc6;
    font-family: 'Georgia', serif;
}

/* Typography enhancements */
#novel-output h1, #novel-output h2, #novel-output h3 {
    color: #2c1810;
    font-family: 'Playfair Display', 'Georgia', serif;
    margin-top: 30px;
    margin-bottom: 20px;
}

#novel-output blockquote {
    border-left: 3px solid #d4a574;
    padding-left: 20px;
    margin: 20px 0;
    font-style: italic;
    color: #5a453a;
}

/* Download section - Book binding style */
.download-section {
    background: linear-gradient(145deg, #faf6f0 0%, #f5efe6 100%);
    padding: 24px;
    border-radius: 14px;
    margin-top: 28px;
    box-shadow: 0 5px 15px rgba(139, 69, 19, 0.08);
    border: 1px solid #e8dcc6;
}

/* Progress bar - Vintage style */
.progress-bar {
    background-color: #f0e6d6;
    height: 28px;
    border-radius: 14px;
    overflow: hidden;
    margin: 18px 0;
    box-shadow: inset 0 3px 6px rgba(139, 69, 19, 0.15);
    border: 1px solid #e0d0b8;
}

.progress-fill {
    background: linear-gradient(90deg, #d4a574 0%, #c8995d 50%, #d4a574 100%);
    height: 100%;
    transition: width 0.6s ease;
    box-shadow: 0 2px 8px rgba(212, 165, 116, 0.4);
}

/* Custom scrollbar - Antique style */
::-webkit-scrollbar {
    width: 12px;
}

::-webkit-scrollbar-track {
    background: #f5f0e6;
    border-radius: 6px;
    box-shadow: inset 0 0 3px rgba(139, 69, 19, 0.1);
}

::-webkit-scrollbar-thumb {
    background: linear-gradient(180deg, #d4a574, #c19656);
    border-radius: 6px;
    box-shadow: 0 2px 4px rgba(139, 69, 19, 0.2);
}

::-webkit-scrollbar-thumb:hover {
    background: linear-gradient(180deg, #c19656, #b08648);
}

/* Button styling - Vintage typewriter keys */
.gr-button {
    background: linear-gradient(145deg, #faf8f5 0%, #f0e8dc 100%);
    border: 1px solid #d4a574;
    color: #3d2914;
    font-weight: 600;
    box-shadow: 0 3px 8px rgba(139, 69, 19, 0.15),
                inset 0 1px 2px rgba(255, 255, 255, 0.8);
    transition: all 0.3s ease;
    font-family: 'Georgia', serif;
}

.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 12px rgba(139, 69, 19, 0.2),
                inset 0 1px 3px rgba(255, 255, 255, 0.9);
    background: linear-gradient(145deg, #fdfbf8 0%, #f3ebe0 100%);
}

.gr-button:active {
    transform: translateY(0);
    box-shadow: 0 2px 5px rgba(139, 69, 19, 0.15),
                inset 0 1px 2px rgba(139, 69, 19, 0.1);
}

/* Primary button - Gold accent */
.gr-button.primary, button[variant="primary"] {
    background: linear-gradient(145deg, #e4c896 0%, #d4a574 100%);
    border: 1px solid #c19656;
    color: #2c1810;
    font-weight: 700;
}

.gr-button.primary:hover, button[variant="primary"]:hover {
    background: linear-gradient(145deg, #e8d0a4 0%, #ddb280 100%);
}

/* Secondary button - Deep brown */
.gr-button.secondary, button[variant="secondary"] {
    background: linear-gradient(145deg, #8b6239 0%, #6d4e31 100%);
    border: 1px solid #5a3e28;
    color: #faf8f5;
}

.gr-button.secondary:hover, button[variant="secondary"]:hover {
    background: linear-gradient(145deg, #96693f 0%, #785436 100%);
}

/* Input fields - Parchment style */
input[type="text"], textarea, .gr-textbox textarea {
    background: #fffefa;
    border: 1px solid #d4c4b0;
    color: #3d2914;
    font-family: 'Georgia', serif;
    box-shadow: inset 0 2px 4px rgba(139, 69, 19, 0.05);
}

input[type="text"]:focus, textarea:focus, .gr-textbox textarea:focus {
    border-color: #c19656;
    box-shadow: 0 0 0 2px rgba(212, 165, 116, 0.2),
                inset 0 2px 4px rgba(139, 69, 19, 0.05);
    outline: none;
}

/* Tab styling - Book chapters */
.gr-tab-button {
    background: #f5f0e6;
    border: 1px solid #d4c4b0;
    color: #5a453a;
    font-weight: 600;
}

.gr-tab-button.selected {
    background: linear-gradient(145deg, #ffffff 0%, #fdfcf8 100%);
    border-bottom-color: transparent;
    color: #2c1810;
    box-shadow: 0 -2px 8px rgba(139, 69, 19, 0.1);
}

/* Dropdown styling */
select, .gr-dropdown {
    background: #fffefa;
    border: 1px solid #d4c4b0;
    color: #3d2914;
}

/* Radio button styling */
.gr-radio-group {
    background: transparent;
}

.gr-radio-group label {
    color: #3d2914;
}

/* Examples section */
.gr-examples {
    background: #f8f4ed;
    border: 1px solid #e8dcc6;
    border-radius: 12px;
    padding: 20px;
    margin-top: 20px;
}

/* Loading animation - Typewriter effect */
@keyframes typewriter {
    from { width: 0; }
    to { width: 100%; }
}

.typing-indicator {
    overflow: hidden;
    border-right: 3px solid #3d2914;
    white-space: nowrap;
    animation: typewriter 3s steps(40, end);
}

/* Markdown content styling */
.markdown-text h1, .markdown-text h2, .markdown-text h3 {
    color: #2c1810;
    font-family: 'Playfair Display', 'Georgia', serif;
}

.markdown-text p {
    color: #3d2914;
    line-height: 1.8;
}

.markdown-text code {
    background: #f5f0e6;
    padding: 2px 6px;
    border-radius: 4px;
    font-family: 'Courier New', monospace;
    color: #5a453a;
}

/* File component styling */
.gr-file {
    background: #faf8f5;
    border: 1px solid #d4c4b0;
    border-radius: 8px;
}

/* Status text special styling */
#status_text textarea {
    background: linear-gradient(145deg, #fff9e6 0%, #fff5d6 100%);
    border: 2px solid #d4a574;
    font-weight: 600;
    text-align: center;
}
"""

# Additional CSS for scrollable theme library
theme_library_css = """
/* Theme Library Styles - Simplified card design */
.library-stats {
    display: flex;
    justify-content: space-around;
    margin-bottom: 30px;
    padding: 20px;
    background: linear-gradient(145deg, #f8f4ed 0%, #f3ede2 100%);
    border-radius: 12px;
    box-shadow: 0 4px 12px rgba(139, 69, 19, 0.08);
}

.stat-item {
    text-align: center;
}

.stat-label {
    display: block;
    font-size: 0.9em;
    color: #5a453a;
    margin-bottom: 5px;
}

.stat-value {
    display: block;
    font-size: 2em;
    font-weight: bold;
    color: #2c1810;
    font-family: 'Playfair Display', 'Georgia', serif;
}

.theme-cards-grid {
    display: grid;
    grid-template-columns: repeat(auto-fill, minmax(400px, 1fr));
    gap: 25px;
    padding: 20px;
}

.theme-card {
    background: linear-gradient(145deg, #ffffff 0%, #fdfcf8 100%);
    border: 1px solid #e8dcc6;
    border-radius: 12px;
    padding: 0;
    box-shadow: 0 4px 12px rgba(139, 69, 19, 0.06);
    transition: all 0.3s ease;
    position: relative;
    overflow: hidden;
    display: flex;
    flex-direction: column;
    height: 450px; /* Reduced height */
}

.theme-card:hover {
    transform: translateY(-3px);
    box-shadow: 0 6px 20px rgba(139, 69, 19, 0.12);
}

.theme-card-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: 15px 20px;
    border-bottom: 1px solid #e8dcc6;
    background: linear-gradient(145deg, #faf6f0 0%, #f5efe6 100%);
    flex-shrink: 0;
}

.theme-id {
    font-family: 'Courier New', monospace;
    color: #8b6239;
    font-size: 0.85em;
    font-weight: bold;
}

.theme-timestamp {
    font-size: 0.8em;
    color: #8a7968;
}

.theme-card-content {
    flex: 1;
    overflow-y: auto;
    padding: 20px;
    background: #fffefa;
}

/* Custom scrollbar for theme content */
.theme-card-content::-webkit-scrollbar {
    width: 6px;
}

.theme-card-content::-webkit-scrollbar-track {
    background: #f5f0e6;
    border-radius: 3px;
}

.theme-card-content::-webkit-scrollbar-thumb {
    background: #d4a574;
    border-radius: 3px;
}

.theme-card-content::-webkit-scrollbar-thumb:hover {
    background: #c19656;
}

.theme-full-text {
    font-family: 'Georgia', serif;
    line-height: 1.8;
    color: #3d2914;
    margin-bottom: 15px;
    font-size: 0.95em;
    text-align: justify;
}

.theme-tags {
    display: flex;
    flex-wrap: wrap;
    gap: 6px;
    margin-top: 15px;
    padding-top: 15px;
    border-top: 1px solid #e8dcc6;
}

.theme-tag {
    display: inline-block;
    padding: 4px 12px;
    background: #f0e6d6;
    border-radius: 15px;
    font-size: 0.75em;
    color: #6d4e31;
    border: 1px solid #d4c4b0;
}

.theme-card-footer {
    display: flex;
    justify-content: space-between;
    align-items: center;
    padding: 15px 20px;
    border-top: 1px solid #e8dcc6;
    background: linear-gradient(145deg, #faf6f0 0%, #f5efe6 100%);
    flex-shrink: 0;
}

.theme-stats {
    display: flex;
    gap: 15px;
}

.theme-stat {
    font-size: 0.85em;
    color: #8a7968;
}

.theme-action-btn {
    padding: 8px 20px;
    font-size: 0.9em;
    border-radius: 6px;
    border: 1px solid #d4a574;
    background: linear-gradient(145deg, #e4c896 0%, #d4a574 100%);
    color: #2c1810;
    cursor: pointer;
    transition: all 0.2s ease;
    font-family: 'Georgia', serif;
    font-weight: bold;
}

.theme-action-btn:hover {
    background: linear-gradient(145deg, #e8d0a4 0%, #ddb280 100%);
    transform: translateY(-1px);
    box-shadow: 0 3px 8px rgba(212, 165, 116, 0.3);
}

.use-btn {
    background: linear-gradient(145deg, #e4c896 0%, #d4a574 100%);
    font-weight: bold;
}

.empty-library {
    text-align: center;
    padding: 60px 20px;
    color: #5a453a;
    font-size: 1.1em;
    font-style: italic;
}

/* Responsive design */
@media (max-width: 768px) {
    .theme-cards-grid {
        grid-template-columns: 1fr;
    }
    
    .theme-card {
        height: 400px;
    }
}
"""

# Create Gradio interface
def create_interface():
    # Combine CSS
    combined_css = custom_css + theme_library_css
    
    # Using Soft theme with safe color options
    with gr.Blocks(theme=gr.themes.Soft(), css=combined_css, title="AGI NOVEL Generator") as interface:
        gr.HTML("""
        <div class="main-header">
            <h1 class="header-title">📖 AGI NOVEL Generator</h1>
            
            <div class="badges-container">
                <a href="https://huggingface.co/spaces/fantaxy/AGI-LEADERBOARD" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Huggingface&message=AGI-LEADERBOARD&color=#d4a574&labelColor=#8b6239&logo=HUGGINGFACE&logoColor=#faf8f5&style=for-the-badge" alt="badge">
                </a>
                <a href="https://discord.gg/openfreeai" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree AI&color=#c19656&labelColor=#6d4e31&logo=discord&logoColor=white&style=for-the-badge" alt="badge">
                </a>
                <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST AI Services&color=#d4a574&labelColor=#5a3e28&logo=huggingface&logoColor=#faf8f5&style=for-the-badge" alt="badge">
                </a>
            </div>
           
            <p class="header-description">
                Artificial General Intelligence (AGI) denotes an artificial system possessing human-level, general-purpose intelligence and is now commonly framed as AI that can outperform humans in most economically and intellectually valuable tasks. Demonstrating such breadth requires evaluating not only calculation, logical reasoning, and perception but also the distinctly human faculties of creativity and language. Among the creative tests, the most demanding is the production of a full-length novel running 100k–200k words. An extended narrative forces an AGI candidate to exhibit (1) sustained long-term memory and context tracking (2) intricate causal and plot planning (3) nuanced cultural and emotional expression (4) autonomous self-censorship and ethical filtering to avoid harmful or biased content and (5) verifiable originality beyond simple recombination of training data.
            </p>
           
            <div class="progress-note">
                🎲 <strong>Novel Theme Random Generator:</strong> This system can generate up to approximately 170 quadrillion (1.7 × 10¹⁷) unique novel themes. 
                Even writing 100 novels per day, it would take 4.6 million years to exhaust all combinations. 
                Click the "Random" button to explore infinite creative possibilities!
            </div>
           
            <div class="warning-note">
                ⏱️ <strong>Note:</strong> Creating a complete novel takes approximately 20 minutes. If your web session disconnects, you can restore your work using the "Session Recovery" feature.
            </div>
           
            <div class="progress-note">
                🎯 <strong>Core Innovation:</strong> Not fragmented texts from multiple writers, 
                but a genuine full-length novel written consistently by a single author from beginning to end.
            </div>
        </div>
        """)
       
        # State management
        current_session_id = gr.State(None)
        selected_theme_id = gr.State(None)
        selected_theme_text = gr.State(None)
        
        # Create tabs and store reference
        with gr.Tabs() as main_tabs:
            # Main Novel Writing Tab
            with gr.Tab("📝 Novel Writing", elem_id="writing_main_tab"):
                # Input section at the top with full width
                with gr.Group(elem_classes=["input-section"]):
                    gr.Markdown("### ✍️ Writing Desk")
                    
                    with gr.Row():
                        with gr.Column(scale=3):
                            query_input = gr.Textbox(
                                label="Novel Theme",
                                placeholder="""Enter your novella theme. Like a seed that grows into a tree, your theme will blossom into a full narrative...

You can describe:
- A specific situation or conflict
- Character relationships and dynamics
- Philosophical questions to explore
- Social or personal transformations
- Any combination of the above

The more detailed your theme, the richer the resulting narrative will be.""",
                                lines=8,
                                elem_id="theme_input"
                            )
                        
                        with gr.Column(scale=1):
                            language_select = gr.Radio(
                                choices=["English", "Korean"],
                                value="English",
                                label="Language",
                                elem_id="language_select"
                            )
                            
                            with gr.Column():
                                random_btn = gr.Button("🎲 Random Theme", variant="primary", size="lg")
                                submit_btn = gr.Button("🖋️ Begin Writing", variant="secondary", size="lg")
                                clear_btn = gr.Button("🗑️ Clear All", size="lg")
                    
                    status_text = gr.Textbox(
                        label="Writing Progress",
                        interactive=False,
                        value="✨ Ready to begin your literary journey",
                        elem_id="status_text"
                    )
                
                # Session management section
                with gr.Group(elem_classes=["session-section"]):
                    gr.Markdown("### 📚 Your Library")
                    with gr.Row():
                        session_dropdown = gr.Dropdown(
                            label="Saved Manuscripts",
                            choices=[],
                            interactive=True,
                            elem_id="session_dropdown",
                            scale=3
                        )
                        refresh_btn = gr.Button("🔄 Refresh", scale=1)
                        resume_btn = gr.Button("📖 Continue", variant="secondary", scale=1)
                        auto_recover_btn = gr.Button("🔮 Recover Last", scale=1)
                
                # Output sections below input
                with gr.Row():
                    with gr.Column():
                        with gr.Tab("🖋️ Writing Process", elem_id="writing_tab"):
                            stages_display = gr.Markdown(
                                value="*Your writing journey will unfold here, like pages turning in a book...*",
                                elem_id="stages-display"
                            )
                   
                        with gr.Tab("📖 Completed Manuscript", elem_id="manuscript_tab"):
                            novel_output = gr.Markdown(
                                value="*Your completed novel will appear here, ready to be read and cherished...*",
                                elem_id="novel-output"
                            )
                       
                            with gr.Group(elem_classes=["download-section"]):
                                gr.Markdown("### 📦 Bind Your Book")
                                with gr.Row():
                                    format_select = gr.Radio(
                                        choices=["DOCX", "TXT"],
                                        value="DOCX" if DOCX_AVAILABLE else "TXT",
                                        label="Format",
                                        elem_id="format_select"
                                    )
                                    download_btn = gr.Button("📥 Download Manuscript", variant="secondary")
                           
                                download_file = gr.File(
                                    label="Your Manuscript",
                                    visible=False,
                                    elem_id="download_file"
                                )
               
                # Hidden state
                novel_text_state = gr.State("")
               
                # Examples with literary flair
                gr.Examples(
                    examples=[
                        ["A daughter discovering her mother's hidden past through old letters found in an attic trunk"],
                        ["An architect losing sight who learns to design through touch, sound, and the memories of light"],
                        ["A translator replaced by AI rediscovering the essence of language through handwritten poetry"],
                        ["A middle-aged man who lost his job finding new meaning in the rhythms of rural life"],
                        ["A doctor with war trauma finding healing through Doctors Without Borders missions"],
                        ["A neighborhood coming together to save their beloved bookstore from corporate development"],
                        ["A year in the life of a professor losing memory and his devoted last student"]
                    ],
                    inputs=query_input,
                    label="💡 Inspiring Themes",
                    examples_per_page=7,
                    elem_id="example_themes"
                )
           
            # Random Theme Library Tab
            with gr.Tab("🎲 Random Theme Library", elem_id="library_tab"):
                with gr.Column():
                    gr.Markdown("""
                    ### 📚 Random Theme Library
                   
                    Browse through all randomly generated themes. Each theme is unique and can be used to create a novel.
                    """)
                   
                    with gr.Row():
                        library_search = gr.Textbox(
                            label="Search Themes",
                            placeholder="Search by keywords...",
                            elem_classes=["library-search"],
                            scale=2
                        )
                        library_language_filter = gr.Radio(
                            choices=["All", "English", "Korean"],
                            value="All",
                            label="Filter by Language",
                            scale=2
                        )
                        library_refresh_btn = gr.Button("🔄 Refresh", scale=1)
                   
                    library_display = gr.HTML(
                        value=get_theme_library_display(),
                        elem_id="library-display"
                    )
                    
                    # Hidden components for theme interaction
                    selected_theme_for_action = gr.Textbox(visible=False, elem_id="selected_theme_for_action")
                    action_type = gr.Textbox(visible=False, elem_id="action_type")
                    trigger_action = gr.Button("Trigger Action", visible=False, elem_id="trigger_action")
       
        # Event handlers
        def refresh_sessions():
            try:
                sessions = get_active_sessions("English")
                return gr.update(choices=sessions)
            except Exception as e:
                logger.error(f"Session refresh error: {str(e)}")
                return gr.update(choices=[])
       
        def handle_auto_recover(language):
            session_id, message = auto_recover_session(language)
            return session_id, message

        def handle_random_theme(language):
            """Handle random theme generation with library storage"""
            theme = generate_random_theme(language)
            return theme
       
        def refresh_library(language_filter, search_query):
            """Refresh theme library display"""
            lang = None if language_filter == "All" else language_filter
            return get_theme_library_display(lang, search_query)
        
        def handle_library_action(theme_id, action):
            """Handle theme library actions"""
            if not theme_id:
                return gr.update(), gr.update()
            
            if action == "use":
                # Handle use action
                theme_data = NovelDatabase.get_theme_by_id(theme_id)
                if theme_data:
                    NovelDatabase.update_theme_used_count(theme_id)
                    NovelDatabase.update_theme_view_count(theme_id)  # Also update view count
                    return (
                        gr.update(value=theme_data.get('theme_text', '')),  # query_input
                        f"Theme #{theme_id[:8]} loaded"   # status_text
                    )
            
            return gr.update(), gr.update()
       
        # Event connections
        submit_btn.click(
            fn=process_query,
            inputs=[query_input, language_select, current_session_id],
            outputs=[stages_display, novel_output, status_text, current_session_id, novel_text_state]  # novel_text_state 추가
        )

        resume_btn.click(
            fn=lambda x: x.split("...")[0] if x and "..." in x else 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, current_session_id, novel_text_state]  # novel_text_state 추가
        )

        auto_recover_btn.click(
            fn=handle_auto_recover,
            inputs=[language_select],
            outputs=[current_session_id, status_text]
        ).then(
            fn=resume_session,
            inputs=[current_session_id, language_select],
            outputs=[stages_display, novel_output, status_text, current_session_id, novel_text_state]  # novel_text_state 추가
        )


        refresh_btn.click(
            fn=refresh_sessions,
            outputs=[session_dropdown]
        )
        
        clear_btn.click(
            fn=lambda: ("", "", "✨ Ready to begin your literary journey", "", None, ""),  # 빈 문자열 추가
            outputs=[stages_display, novel_output, status_text, novel_text_state, current_session_id, novel_text_state]
        )       

        random_btn.click(
            fn=handle_random_theme,
            inputs=[language_select],
            outputs=[query_input],
            queue=False
        )
       
        # Library event handlers
        library_refresh_btn.click(
            fn=refresh_library,
            inputs=[library_language_filter, library_search],
            outputs=[library_display]
        )
       
        library_search.change(
            fn=refresh_library,
            inputs=[library_language_filter, library_search],
            outputs=[library_display]
        )
       
        library_language_filter.change(
            fn=refresh_library,
            inputs=[library_language_filter, library_search],
            outputs=[library_display]
        )
        
        # Handle clicks on library display - using trigger button
        trigger_action.click(
            fn=handle_library_action,
            inputs=[selected_theme_for_action, action_type],
            outputs=[query_input, status_text]
        )
       
       
        download_btn.click(
            fn=handle_download,
            inputs=[format_select, language_select, current_session_id, novel_text_state],
            outputs=[download_file]
        )
       
        # Load sessions and library on start
        def initialize_interface():
            # Sync with HF dataset on startup
            if 'hf_manager' in globals() and hf_manager.token:
                hf_manager.sync_with_local_db()
            
            return refresh_sessions(), refresh_library("All", "")
        
        interface.load(
            fn=initialize_interface,
            outputs=[session_dropdown, library_display]
        )
   
    return interface
# Initialize HF Dataset Manager as global variable
hf_manager = None

# Main function
if __name__ == "__main__":
    logger.info("AGI NOVEL Generator v2.0 Starting...")
    logger.info("=" * 60)
    
    # Environment check
    logger.info(f"API Endpoint: {API_URL}")
    logger.info(f"Target Length: {TARGET_WORDS:,} words")
    logger.info(f"Minimum Words per Part: {MIN_WORDS_PER_PART:,} words")
    logger.info("System Features: Single writer + Immediate part-by-part critique")
    
    if BRAVE_SEARCH_API_KEY:
        logger.info("Web search enabled.")
    else:
        logger.warning("Web search disabled.")
    
    if DOCX_AVAILABLE:
        logger.info("DOCX export enabled.")
    else:
        logger.warning("DOCX export disabled.")
    
    logger.info("=" * 60)
    
    # Initialize database
    logger.info("Initializing database...")
    NovelDatabase.init_db()
    logger.info("Database initialization complete.")
    
    # Initialize HF Dataset Manager
    logger.info("Initializing HuggingFace dataset manager...")
    hf_manager = HFDatasetManager()
    
    if hf_manager.token:
        logger.info("HuggingFace authentication successful.")
        # Sync with HF dataset on startup
        hf_manager.sync_with_local_db()
    else:
        logger.warning("HuggingFace token not found. Theme persistence will be local only.")
    
    # Create and launch interface
    interface = create_interface()
    
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )