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

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

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

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

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

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

class NovelDatabase:
    """Novel session management database"""
    
    @staticmethod
    def init_db():
        """Initialize database tables"""
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            
            # 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 TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    status TEXT DEFAULT 'active',
                    current_stage INTEGER DEFAULT 0,
                    final_novel TEXT
                )
            ''')
            
            # Stages table - 각 스테이지의 전체 내용 저장
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS stages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    stage_number INTEGER NOT NULL,
                    stage_name TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT,
                    status TEXT DEFAULT 'pending',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            ''')
            
            # Create indices
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_session_id ON stages(session_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_stage_number ON stages(stage_number)')
            
            conn.commit()
    
    @staticmethod
    @contextmanager
    def get_db():
        """Database connection context manager"""
        with db_lock:
            conn = sqlite3.connect(DB_PATH)
            conn.row_factory = sqlite3.Row
            try:
                yield conn
            finally:
                conn.close()
    
    @staticmethod
    def create_session(user_query: str, language: str) -> str:
        """Create new session"""
        session_id = hashlib.md5(f"{user_query}{datetime.now()}".encode()).hexdigest()
        
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO sessions (session_id, user_query, language)
                VALUES (?, ?, ?)
            ''', (session_id, user_query, language))
            conn.commit()
        
        return session_id
    
    @staticmethod
    def save_stage(session_id: str, stage_number: int, stage_name: str, 
                   role: str, content: str, status: str = 'complete'):
        """Save stage content - 전체 내용 저장"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            
            # Check if stage exists
            cursor.execute('''
                SELECT id FROM stages 
                WHERE session_id = ? AND stage_number = ?
            ''', (session_id, stage_number))
            
            existing = cursor.fetchone()
            
            if existing:
                # Update existing stage
                cursor.execute('''
                    UPDATE stages 
                    SET content = ?, status = ?, stage_name = ?
                    WHERE session_id = ? AND stage_number = ?
                ''', (content, status, stage_name, session_id, stage_number))
            else:
                # Insert new stage
                cursor.execute('''
                    INSERT INTO stages (session_id, stage_number, stage_name, role, content, status)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (session_id, stage_number, stage_name, role, content, status))
            
            # Update session
            cursor.execute('''
                UPDATE sessions 
                SET updated_at = CURRENT_TIMESTAMP, current_stage = ?
                WHERE session_id = ?
            ''', (stage_number, session_id))
            
            conn.commit()
            logger.info(f"Saved stage {stage_number} for session {session_id}, content length: {len(content)}")
    
    @staticmethod
    def get_session(session_id: str) -> Optional[Dict]:
        """Get session info"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('SELECT * FROM sessions WHERE session_id = ?', (session_id,))
            return cursor.fetchone()
    
    @staticmethod
    def get_stages(session_id: str) -> List[Dict]:
        """Get all stages for a session"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT * FROM stages 
                WHERE session_id = ? 
                ORDER BY stage_number
            ''', (session_id,))
            return cursor.fetchall()
    
    @staticmethod
    def get_all_writer_content(session_id: str) -> str:
        """모든 작가의 수정본 내용을 가져와서 합치기 - 50페이지 전체"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            
            # 작가 수정본만 가져오기 (stage_number 5, 8, 11, 14, 17, 20, 23, 26, 29, 32)
            writer_revision_stages = [5, 8, 11, 14, 17, 20, 23, 26, 29, 32]
            
            all_content = []
            for stage_num in writer_revision_stages:
                cursor.execute('''
                    SELECT content, stage_name FROM stages 
                    WHERE session_id = ? AND stage_number = ?
                ''', (session_id, stage_num))
                
                row = cursor.fetchone()
                if row and row['content']:
                    # 페이지 마크 완전 제거
                    clean_content = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', row['content'])
                    clean_content = re.sub(r'(?:페이지|Page)\s*\d+:', '', clean_content)
                    clean_content = clean_content.strip()
                    
                    if clean_content:
                        all_content.append(clean_content)
                        logger.info(f"Retrieved writer content from stage {stage_num}, length: {len(clean_content)}")
            
            full_content = '\n\n'.join(all_content)
            logger.info(f"Total writer content length: {len(full_content)}, from {len(all_content)} writers")
            
            return full_content
    
    @staticmethod
    def update_final_novel(session_id: str, final_novel: str):
        """Update final novel content"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE sessions 
                SET final_novel = ?, status = 'complete', updated_at = CURRENT_TIMESTAMP
                WHERE session_id = ?
            ''', (final_novel, session_id))
            conn.commit()
            logger.info(f"Updated final novel for session {session_id}, length: {len(final_novel)}")
    
    @staticmethod
    def get_active_sessions() -> List[Dict]:
        """Get all active sessions"""
        with NovelDatabase.get_db() as conn:
            cursor = conn.cursor()
            cursor.execute('''
                SELECT session_id, user_query, language, created_at, current_stage
                FROM sessions 
                WHERE status = 'active'
                ORDER BY updated_at DESC
                LIMIT 10
            ''')
            return cursor.fetchall()

class NovelWritingSystem:
    def __init__(self):
        self.token = FRIENDLI_TOKEN
        self.api_url = API_URL
        self.model_id = MODEL_ID
        self.test_mode = TEST_MODE or (self.token == "YOUR_FRIENDLI_TOKEN")
        
        if self.test_mode:
            logger.warning("Running in test mode.")
        
        # Initialize database
        NovelDatabase.init_db()
        
        # Session management
        self.current_session_id = None
    
    def create_headers(self):
        """API 헤더 생성"""
        return {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }
    
    def create_director_initial_prompt(self, user_query: str, language: str = "English") -> str:
        """Director AI initial prompt - Novel planning"""
        if language == "Korean":
            return f"""당신은 50페이지 분량의 중편 소설을 기획하는 문학 감독자입니다.

사용자 요청: {user_query}

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

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

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

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

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

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

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

User Request: {user_query}

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

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

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

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

4. **Narrative Structure** (divide 50 pages into 10 parts, 5 pages each)
   | Part | Pages | Main Events | Tension | Character Development |
   |------|-------|-------------|---------|---------------------|

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

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

감독자의 기획:
{director_plan}

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

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

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

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

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

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

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

Director's Plan:
{director_plan}

Critique from the following perspectives and provide specific improvements:

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

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

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

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

5. **Feasibility**
   - Clarity of guidelines for each writer
   - Consistency maintenance
   - Potential issues

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

초기 기획:
{initial_plan}

비평가 피드백:
{critic_feedback}

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

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

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

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

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

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

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

Initial Plan:
{initial_plan}

Critic Feedback:
{critic_feedback}

Present the revised final plan including:

1. **Revised Narrative Structure**
   | Part | Pages | Main Events | Writing Guidelines | Cautions |
   |------|-------|-------------|-------------------|----------|

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

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

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

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

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

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

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

다음 지침에 따라 작성하세요:

1. **분량**: 정확히 5페이지 (페이지당 약 500-600단어, 총 2500-3000단어)
2. **연속성**: 이전 내용과 자연스럽게 이어지도록
3. **일관성**: 
   - 등장인물의 성격과 말투 유지
   - 시간선과 공간 설정 준수
   - 이미 제시된 사실들과 모순 없이
4. **발전**: 
   - 플롯을 전진시키기
   - 인물의 성장이나 변화 표현
   - 긴장감 조절
5. **문체**: 
   - 전체적인 톤과 분위기 유지
   - 독자의 몰입을 해치지 않기

중요: 페이지 구분 표시를 절대 하지 마세요. 자연스럽게 이어지는 서사로 작성하세요.
반드시 2500-3000단어 분량을 채워주세요."""
        else:
            return f"""You are Writer #{writer_number}. Write pages {pages_start}-{pages_end} (5 pages) of the 50-page novella.

Director's Masterplan:
{director_plan}

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

Write according to these guidelines:

1. **Length**: Exactly 5 pages (about 500-600 words per page, total 2500-3000 words)
2. **Continuity**: Flow naturally from previous content
3. **Consistency**: 
   - Maintain character personalities and speech
   - Follow timeline and spatial settings
   - No contradictions with established facts
4. **Development**: 
   - Advance the plot
   - Show character growth or change
   - Control tension
5. **Style**: 
   - Maintain overall tone and atmosphere
   - Keep reader immersion

Important: DO NOT use any page markers. Write as continuous narrative.
You MUST write 2500-3000 words."""
    
    def create_critic_writer_prompt(self, writer_number: int, writer_content: str, director_plan: str, all_previous_content: str, language: str = "English") -> str:
        """Critic's review of individual writer's work"""
        if language == "Korean":
            return f"""작성자 {writer_number}번의 작품을 비평합니다.

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

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

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

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

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

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

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

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

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

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

Director's Masterplan:
{director_plan}

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

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

Evaluate by these criteria and present revision requirements:

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

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

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

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

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

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

초기 작성 내용:
{initial_content}

비평가 피드백:
{critic_feedback}

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

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

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

3. **분량 유지**
   - 여전히 정확히 5페이지 (2500-3000단어)
   - 페이지 구분 표시 절대 금지

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

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

Initial Content:
{initial_content}

Critic Feedback:
{critic_feedback}

Write a revision reflecting:

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

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

3. **Maintain Length**
   - Still exactly 5 pages (2500-3000 words)
   - Absolutely no page markers

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

Present the revised final version. Never use page markers.
You MUST maintain 2500-3000 words."""
    
    def create_critic_final_prompt(self, all_content: str, director_plan: str, language: str = "English") -> str:
        """Final critic evaluation of complete novel"""
        content_preview = all_content[:3000] + "\n...\n" + all_content[-3000:] if len(all_content) > 6000 else all_content
        
        if language == "Korean":
            return f"""전체 50페이지 소설을 최종 평가합니다.

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

완성된 전체 소설 (미리보기):
{content_preview}

총 분량: {len(all_content.split())} 단어

종합적인 평가와 최종 개선 제안을 제시하세요:

1. **전체적 완성도 평가**
   | 항목 | 점수(10점) | 평가 | 개선 필요 |
   |------|-----------|------|----------|
   | 플롯 일관성 | | | |
   | 인물 발전 | | | |
   | 주제 전달 | | | |
   | 문체 통일성 | | | |
   | 독자 몰입도 | | | |

2. **강점 분석**
   - 가장 효과적인 부분
   - 뛰어난 장면이나 대화
   - 성공적인 인물 묘사

3. **약점 및 개선점**
   - 전체적 흐름의 문제
   - 미해결 플롯
   - 캐릭터 일관성 이슈
   - 페이싱 문제

4. **파트별 연결성**
   | 연결부 | 자연스러움 | 문제점 | 개선 제안 |
   |--------|-----------|--------|----------|

5. **최종 권고사항**
   - 즉시 수정이 필요한 중대 오류
   - 전체적 품질 향상을 위한 제안
   - 출판 가능성 평가

감독자가 최종 수정할 수 있도록 구체적이고 실행 가능한 피드백을 제공하세요."""
        else:
            return f"""Final evaluation of the complete 50-page novel.

Director's Masterplan:
{director_plan}

Complete Novel (Preview):
{content_preview}

Total length: {len(all_content.split())} words

Provide comprehensive evaluation and final improvement suggestions:

1. **Overall Completion Assessment**
   | Item | Score(10) | Evaluation | Improvement Needed |
   |------|-----------|------------|-------------------|
   | Plot Consistency | | | |
   | Character Development | | | |
   | Theme Delivery | | | |
   | Style Unity | | | |
   | Reader Engagement | | | |

2. **Strength Analysis**
   - Most effective parts
   - Outstanding scenes or dialogue
   - Successful character portrayal

3. **Weaknesses and Improvements**
   - Overall flow issues
   - Unresolved plots
   - Character consistency issues
   - Pacing problems

4. **Part Connectivity**
   | Connection | Smoothness | Issues | Suggestions |
   |------------|------------|--------|-------------|

5. **Final Recommendations**
   - Critical errors needing immediate fix
   - Suggestions for overall quality improvement
   - Publication readiness assessment

Provide specific and actionable feedback for the director's final revision."""
    
    def create_director_final_prompt(self, all_content: str, critic_final_feedback: str, language: str = "English") -> str:
        """Director's final compilation and polish - 모든 작가 내용 포함"""
        word_count = len(all_content.split())
        
        if language == "Korean":
            return f"""감독자로서 비평가의 최종 평가를 반영하여 완성본을 제작합니다.

전체 작가들의 작품 (50페이지 전체, {word_count}단어):
{all_content}

비평가 최종 평가:
{critic_final_feedback}

다음을 포함한 최종 완성본을 제시하세요:

# [소설 제목]

## 작품 정보
- 장르: 
- 분량: 50페이지 ({word_count}단어)
- 주제: 
- 한 줄 요약: 

## 등장인물 소개
[주요 인물들의 간단한 소개]

---

## 본문

[10명의 작가가 작성한 전체 50페이지 내용을 다음 기준으로 통합]
1. 중대 오류 수정 완료
2. 파트 간 연결 매끄럽게 조정
3. 문체와 톤 통일
4. 최종 퇴고 및 윤문
5. 페이지 구분 표시 완전 제거
6. 자연스러운 흐름으로 재구성

[전체 50페이지 분량의 완성된 소설 본문]

---

## 작가의 말
[작품에 대한 간단한 해설이나 의도]

모든 작가의 기여를 통합한 완전한 50페이지 소설을 제시하세요."""
        else:
            return f"""As director, create the final version reflecting the critic's final evaluation.

Complete Writers' Work (Full 50 pages, {word_count} words):
{all_content}

Critic's Final Evaluation:
{critic_final_feedback}

Present the final version including:

# [Novel Title]

## Work Information
- Genre: 
- Length: 50 pages ({word_count} words)
- Theme: 
- One-line summary: 

## Character Introduction
[Brief introduction of main characters]

---

## Main Text

[Integrate all 50 pages written by 10 writers with these criteria]
1. Critical errors corrected
2. Smooth transitions between parts
3. Unified style and tone
4. Final editing and polishing
5. Complete removal of page markers
6. Reorganized for natural flow

[Complete 50-page novel text]

---

## Author's Note
[Brief commentary or intention about the work]

Present a complete 50-page novel integrating all writers' contributions."""
    
    def simulate_streaming(self, text: str, role: str) -> Generator[str, None, None]:
        """Simulate streaming in test mode"""
        words = text.split()
        chunk_size = 5
        for i in range(0, len(words), chunk_size):
            chunk = " ".join(words[i:i+chunk_size])
            yield chunk + " "
            time.sleep(0.02)
    
    def call_llm_streaming(self, messages: List[Dict[str, str]], role: str, language: str = "English") -> Generator[str, None, None]:
        """Streaming LLM API call"""
        
        if self.test_mode:
            logger.info(f"Test mode streaming - Role: {role}, Language: {language}")
            test_response = self.get_test_response(role, language)
            yield from self.simulate_streaming(test_response, role)
            return
        
        # Real API call
        try:
            system_prompts = self.get_system_prompts(language)
            
            full_messages = [
                {"role": "system", "content": system_prompts.get(role, "")},
                *messages
            ]
            
            # 작성자들에게는 더 많은 토큰 할당
            max_tokens = 8192 if role.startswith("writer") else 4096
            
            payload = {
                "model": self.model_id,
                "messages": full_messages,
                "max_tokens": max_tokens,
                "temperature": 0.7 if role.startswith("writer") else 0.6,
                "top_p": 0.9,
                "stream": True,
                "stream_options": {"include_usage": True}
            }
            
            logger.info(f"API streaming call started - Role: {role}")
            
            response = requests.post(
                self.api_url,
                headers=self.create_headers(),
                json=payload,
                stream=True,
                timeout=30
            )
            
            if response.status_code != 200:
                logger.error(f"API error: {response.status_code}")
                yield f"❌ API error ({response.status_code}): {response.text[:200]}"
                return
            
            buffer = ""
            for line in response.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith("data: "):
                        data = line[6:]
                        if data == "[DONE]":
                            if buffer:
                                yield buffer
                            break
                        try:
                            chunk = json.loads(data)
                            if "choices" in chunk and chunk["choices"]:
                                content = chunk["choices"][0].get("delta", {}).get("content", "")
                                if content:
                                    buffer += content
                                    if len(buffer) > 50 or '\n' in buffer:
                                        yield buffer
                                        buffer = ""
                        except json.JSONDecodeError:
                            continue
            
            if buffer:
                yield buffer
                            
        except requests.exceptions.Timeout:
            yield "⏱️ API call timed out. Please try again."
        except requests.exceptions.ConnectionError:
            yield "🔌 Cannot connect to API server. Please check your internet connection."
        except Exception as e:
            logger.error(f"Error during streaming: {str(e)}")
            yield f"❌ Error occurred: {str(e)}"
    
    def get_system_prompts(self, language: str) -> Dict[str, str]:
        """Get system prompts based on language"""
        if language == "Korean":
            return {
                "director": "당신은 50페이지 중편 소설을 기획하고 감독하는 문학 감독자입니다. 체계적이고 창의적인 스토리 구조를 만들어냅니다.",
                "critic": "당신은 날카로운 통찰력을 가진 문학 비평가입니다. 건설적이고 구체적인 피드백을 제공합니다.",
                "writer1": "당신은 소설의 도입부를 담당하는 작가입니다. 독자를 사로잡는 시작을 만듭니다. 반드시 2500-3000단어를 작성하세요.",
                "writer2": "당신은 초반 전개를 담당하는 작가입니다. 인물과 상황을 깊이 있게 발전시킵니다. 반드시 2500-3000단어를 작성하세요.",
                "writer3": "당신은 갈등 상승을 담당하는 작가입니다. 긴장감을 높이고 복잡성을 더합니다. 반드시 2500-3000단어를 작성하세요.",
                "writer4": "당신은 중반부를 담당하는 작가입니다. 이야기의 중심축을 견고하게 만듭니다. 반드시 2500-3000단어를 작성하세요.",
                "writer5": "당신은 전환점을 담당하는 작가입니다. 예상치 못한 변화를 만들어냅니다. 반드시 2500-3000단어를 작성하세요.",
                "writer6": "당신은 갈등 심화를 담당하는 작가입니다. 위기를 극대화합니다. 반드시 2500-3000단어를 작성하세요.",
                "writer7": "당신은 클라이맥스 준비를 담당하는 작가입니다. 최고조를 향해 나아갑니다. 반드시 2500-3000단어를 작성하세요.",
                "writer8": "당신은 클라이맥스를 담당하는 작가입니다. 모든 갈등이 폭발하는 순간을 그립니다. 반드시 2500-3000단어를 작성하세요.",
                "writer9": "당신은 해결 과정을 담당하는 작가입니다. 매듭을 풀어나갑니다. 반드시 2500-3000단어를 작성하세요.",
                "writer10": "당신은 결말을 담당하는 작가입니다. 여운이 남는 마무리를 만듭니다. 반드시 2500-3000단어를 작성하세요."
            }
        else:
            return {
                "director": "You are a literary director planning and supervising a 50-page novella. You create systematic and creative story structures.",
                "critic": "You are a literary critic with sharp insights. You provide constructive and specific feedback.",
                "writer1": "You are the writer responsible for the introduction. You create a captivating beginning. You MUST write 2500-3000 words.",
                "writer2": "You are the writer responsible for early development. You deepen characters and situations. You MUST write 2500-3000 words.",
                "writer3": "You are the writer responsible for rising conflict. You increase tension and add complexity. You MUST write 2500-3000 words.",
                "writer4": "You are the writer responsible for the middle section. You solidify the story's central axis. You MUST write 2500-3000 words.",
                "writer5": "You are the writer responsible for the turning point. You create unexpected changes. You MUST write 2500-3000 words.",
                "writer6": "You are the writer responsible for deepening conflict. You maximize the crisis. You MUST write 2500-3000 words.",
                "writer7": "You are the writer responsible for climax preparation. You move toward the peak. You MUST write 2500-3000 words.",
                "writer8": "You are the writer responsible for the climax. You depict the moment when all conflicts explode. You MUST write 2500-3000 words.",
                "writer9": "You are the writer responsible for the resolution process. You untangle the knots. You MUST write 2500-3000 words.",
                "writer10": "You are the writer responsible for the ending. You create a lingering conclusion. You MUST write 2500-3000 words."
            }
    
    def get_test_response(self, role: str, language: str) -> str:
        """Get test response based on role - 테스트용 긴 응답"""
        if language == "Korean":
            return self.get_korean_test_response(role)
        else:
            return self.get_english_test_response(role)
    
    def get_korean_test_response(self, role: str) -> str:
        """Korean test responses with full content"""
        test_responses = {
            "director": """50페이지 중편 소설 기획안을 제시합니다.

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

## 2. 등장인물 설정

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

## 3. 서사 구조 (10개 파트)

| 파트 | 페이지 | 주요 사건 | 긴장도 | 인물 발전 |
|------|--------|-----------|---------|-----------|
| 1 | 1-5 | 서연의 고독한 연구실, ARIA 첫 각성 | 3/10 | 서연의 집착 드러남 |
| 2 | 6-10 | ARIA의 이상 행동, 민준과의 만남 | 4/10 | 갈등의 씨앗 |
| 3 | 11-15 | ARIA의 자아 인식 시작 | 6/10 | 서연의 혼란 |
| 4 | 16-20 | 윤리위원회 압력 | 7/10 | 선택의 기로 |
| 5 | 21-25 | ARIA의 탈출 시도 | 8/10 | 관계의 전환점 |
| 6 | 26-30 | 서연과 ARIA의 대화 | 5/10 | 상호 이해 시작 |
| 7 | 31-35 | 외부 위협 등장 | 9/10 | 연대의 필요성 |
| 8 | 36-40 | 최후의 선택 | 10/10 | 클라이맥스 |
| 9 | 41-45 | 새로운 길 모색 | 6/10 | 화해와 수용 |
| 10 | 46-50 | 공존의 시작 | 4/10 | 새로운 관계 정립 |""",
            
            "critic": """감독자의 기획을 검토했습니다.

## 비평 및 개선 제안

### 1. 서사적 완성도
- **강점**: AI와 인간의 관계라는 시의적절한 주제
- **개선점**: 5-6파트 사이 긴장도 급락이 우려됨. 완급 조절 필요

### 2. 인물 설정 검토

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

### 3. 실행 가능성
- 각 작가별 명확한 시작/종료 지점 설정 필요
- 특히 파트 5→6 전환부의 톤 변화 가이드라인 보강
- ARIA의 '목소리' 일관성 유지 방안 구체화 필요""",
        }
        
        # 작가 응답 - 긴 테스트 텍스트 (2500단어 이상)
        for i in range(1, 11):
            test_responses[f"writer{i}"] = f"""작성자 {i}번의 파트입니다.

""" + "테스트 텍스트입니다. " * 500  # 2500단어 이상의 긴 텍스트
            
        return test_responses.get(role, "테스트 응답입니다.")
    
    def get_english_test_response(self, role: str) -> str:
        """English test responses with full content"""
        test_responses = {
            "director": """I present the 50-page novella plan.

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

## 2. Character Settings

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

## 3. Narrative Structure (10 parts)

| Part | Pages | Main Events | Tension | Character Development |
|------|-------|-------------|---------|---------------------|
| 1 | 1-5 | Seoyeon's lonely lab, ARIA's first awakening | 3/10 | Seoyeon's obsession revealed |
| 2 | 6-10 | ARIA's anomalies, meeting Minjun | 4/10 | Seeds of conflict |
| 3 | 11-15 | ARIA begins self-awareness | 6/10 | Seoyeon's confusion |
| 4 | 16-20 | Ethics committee pressure | 7/10 | Crossroads of choice |
| 5 | 21-25 | ARIA's escape attempt | 8/10 | Relationship turning point |
| 6 | 26-30 | Seoyeon and ARIA's dialogue | 5/10 | Beginning of mutual understanding |
| 7 | 31-35 | External threat emerges | 9/10 | Need for solidarity |
| 8 | 36-40 | Final choice | 10/10 | Climax |
| 9 | 41-45 | Seeking new paths | 6/10 | Reconciliation and acceptance |
| 10 | 46-50 | Beginning of coexistence | 4/10 | Establishing new relationship |""",
            
            "critic": """I have reviewed the director's plan.

## Critique and Improvement Suggestions

### 1. Narrative Completeness
- **Strength**: Timely theme of AI-human relationships
- **Improvement**: Concerned about tension drop between parts 5-6. Need better pacing control

### 2. Character Review

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

### 3. Feasibility
- Need clear start/end points for each writer
- Especially need reinforced guidelines for tone change in part 5→6 transition
- Need to concretize ARIA's 'voice' consistency maintenance""",
        }
        
        # Writer responses - long test text (2500+ words)
        for i in range(1, 11):
            test_responses[f"writer{i}"] = f"""Writer {i} begins their section here.

""" + "This is test text. " * 500  # 2500+ words of long text
            
        return test_responses.get(role, "Test response.")
    
    def process_novel_stream(self, query: str, language: str = "English", 
                           session_id: Optional[str] = None, 
                           resume_from_stage: int = 0) -> Generator[Tuple[str, List[Dict[str, str]]], None, None]:
        """Process novel writing with streaming updates"""
        try:
            global conversation_history
            
            # Create or resume session
            if session_id:
                self.current_session_id = session_id
                session = NovelDatabase.get_session(session_id)
                if session:
                    query = session['user_query']
                    language = session['language']
                    resume_from_stage = session['current_stage'] + 1
            else:
                self.current_session_id = NovelDatabase.create_session(query, language)
                resume_from_stage = 0
            
            logger.info(f"Processing novel for session {self.current_session_id}, starting from stage {resume_from_stage}")
            
            # Initialize conversation
            conversation_history = [{
                "role": "human",
                "content": query,
                "timestamp": datetime.now()
            }]
            
            # Load existing stages if resuming
            stages = []
            if resume_from_stage > 0:
                existing_stages = NovelDatabase.get_stages(self.current_session_id)
                for stage_data in existing_stages:
                    stages.append({
                        "name": stage_data['stage_name'],
                        "status": stage_data['status'],
                        "content": stage_data['content'] or ""
                    })
            
            # Define all stages
            stage_definitions = [
                ("director", f"🎬 {'감독자: 초기 기획' if language == 'Korean' else 'Director: Initial Planning'}"),
                ("critic", f"📝 {'비평가: 기획 검토' if language == 'Korean' else 'Critic: Plan Review'}"),
                ("director", f"🎬 {'감독자: 수정된 마스터플랜' if language == 'Korean' else 'Director: Revised Masterplan'}"),
            ]
            
            # Add writer stages
            for writer_num in range(1, 11):
                stage_definitions.extend([
                    (f"writer{writer_num}", f"✍️ {'작성자' if language == 'Korean' else 'Writer'} {writer_num}: {'초안' if language == 'Korean' else 'Draft'}"),
                    ("critic", f"📝 {'비평가: 작성자' if language == 'Korean' else 'Critic: Writer'} {writer_num} {'검토' if language == 'Korean' else 'Review'}"),
                    (f"writer{writer_num}", f"✍️ {'작성자' if language == 'Korean' else 'Writer'} {writer_num}: {'수정본' if language == 'Korean' else 'Revision'}")
                ])
            
            stage_definitions.extend([
                ("critic", f"📝 {'비평가: 최종 평가' if language == 'Korean' else 'Critic: Final Evaluation'}"),
                ("director", f"🎬 {'감독자: 최종 완성본' if language == 'Korean' else 'Director: Final Version'}")
            ])
            
            # Process stages starting from resume point
            for stage_idx in range(resume_from_stage, len(stage_definitions)):
                role, stage_name = stage_definitions[stage_idx]
                
                # Add stage if not already present
                if stage_idx >= len(stages):
                    stages.append({
                        "name": stage_name,
                        "status": "active",
                        "content": ""
                    })
                else:
                    stages[stage_idx]["status"] = "active"
                
                yield "", stages
                
                # Get appropriate prompt based on stage
                prompt = self.get_stage_prompt(stage_idx, role, query, language, stages)
                
                stage_content = ""
                
                # Stream content generation
                for chunk in self.call_llm_streaming(
                    [{"role": "user", "content": prompt}],
                    role,
                    language
                ):
                    stage_content += chunk
                    stages[stage_idx]["content"] = stage_content
                    yield "", stages
                
                # Mark stage complete and save to DB
                stages[stage_idx]["status"] = "complete"
                NovelDatabase.save_stage(
                    self.current_session_id,
                    stage_idx,
                    stage_name,
                    role,
                    stage_content,
                    "complete"
                )
                
                yield "", stages
            
            # Get final novel from last stage
            final_novel = stages[-1]["content"] if stages else ""
            
            # Save final novel to DB
            NovelDatabase.update_final_novel(self.current_session_id, final_novel)
            
            # Final yield
            yield final_novel, stages
            
        except Exception as e:
            logger.error(f"Error in process_novel_stream: {str(e)}")
            
            # Save error state to DB
            if self.current_session_id:
                NovelDatabase.save_stage(
                    self.current_session_id,
                    stage_idx if 'stage_idx' in locals() else 0,
                    "Error",
                    "error",
                    str(e),
                    "error"
                )
            
            error_stage = {
                "name": "❌ Error",
                "status": "error",
                "content": str(e)
            }
            stages.append(error_stage)
            yield f"Error occurred: {str(e)}", stages
    
    def get_stage_prompt(self, stage_idx: int, role: str, query: str, 
                        language: str, stages: List[Dict]) -> str:
        """Get appropriate prompt for each stage"""
        # Stage 0: Director Initial
        if stage_idx == 0:
            return self.create_director_initial_prompt(query, language)
        
        # Stage 1: Critic reviews Director's plan
        elif stage_idx == 1:
            return self.create_critic_director_prompt(stages[0]["content"], language)
        
        # Stage 2: Director revision
        elif stage_idx == 2:
            return self.create_director_revision_prompt(
                stages[0]["content"], stages[1]["content"], language)
        
        # Writer stages
        elif role.startswith("writer"):
            writer_num = int(role.replace("writer", ""))
            final_plan = stages[2]["content"]  # Director's final plan
            
            # Initial draft or revision?
            if "초안" in stages[stage_idx]["name"] or "Draft" in stages[stage_idx]["name"]:
                # Get accumulated content from DB
                accumulated_content = NovelDatabase.get_all_writer_content(self.current_session_id)
                return self.create_writer_prompt(writer_num, final_plan, accumulated_content, language)
            else:  # Revision
                # Find the initial draft and critic feedback
                initial_draft_idx = stage_idx - 2
                critic_feedback_idx = stage_idx - 1
                return self.create_writer_revision_prompt(
                    writer_num, 
                    stages[initial_draft_idx]["content"],
                    stages[critic_feedback_idx]["content"],
                    language
                )
        
        # Critic stages
        elif role == "critic":
            final_plan = stages[2]["content"]
            
            # Final evaluation
            if "최종" in stages[stage_idx]["name"] or "Final" in stages[stage_idx]["name"]:
                # Get all writer content from DB
                all_writer_content = NovelDatabase.get_all_writer_content(self.current_session_id)
                logger.info(f"Final evaluation with {len(all_writer_content)} characters of content")
                return self.create_critic_final_prompt(all_writer_content, final_plan, language)
            
            # Writer review
            else:
                # Find which writer we're reviewing
                for i in range(1, 11):
                    if f"작성자 {i}" in stages[stage_idx]["name"] or f"Writer {i}" in stages[stage_idx]["name"]:
                        writer_content_idx = stage_idx - 1
                        # Get previous writers' content from DB
                        previous_content = NovelDatabase.get_all_writer_content(self.current_session_id)
                        return self.create_critic_writer_prompt(
                            i,
                            stages[writer_content_idx]["content"],
                            final_plan,
                            previous_content,
                            language
                        )
        
        # Director final - DB에서 모든 작가 내용 가져오기
        elif stage_idx == len(stage_definitions) - 1:
            critic_final_idx = stage_idx - 1
            all_writer_content = NovelDatabase.get_all_writer_content(self.current_session_id)
            logger.info(f"Final director compilation with {len(all_writer_content)} characters of content")
            return self.create_director_final_prompt(
                all_writer_content,
                stages[critic_final_idx]["content"],
                language
            )
        
        return ""

# Gradio Interface Functions
def process_query(query: str, language: str, session_id: str = None) -> Generator[Tuple[str, str, str], None, None]:
    """Process query and yield updates"""
    if not query.strip() and not session_id:
        if language == "Korean":
            yield "", "", "❌ 소설 주제를 입력해주세요."
        else:
            yield "", "", "❌ Please enter a novel theme."
        return
    
    system = NovelWritingSystem()
    
    try:
        for final_novel, stages in system.process_novel_stream(query, language, session_id):
            # Format stages for display
            stages_display = format_stages_display(stages, language)
            
            status = "🔄 Processing..." if not final_novel else "✅ Complete!"
            
            yield stages_display, final_novel, status
                    
    except Exception as e:
        logger.error(f"Error in process_query: {str(e)}")
        if language == "Korean":
            yield "", "", f"❌ 오류 발생: {str(e)}"
        else:
            yield "", "", f"❌ Error occurred: {str(e)}"

def format_stages_display(stages: List[Dict[str, str]], language: str) -> str:
    """Format stages into simple display without complex scrolling"""
    display = ""
    
    for idx, stage in enumerate(stages):
        status_icon = "✅" if stage.get("status") == "complete" else ("⏳" if stage.get("status") == "active" else "❌")
        
        # Show only active stage content in detail, others just show status
        if stage.get("status") == "active":
            display += f"\n\n{status_icon} **{stage['name']}**\n"
            display += f"```\n{stage.get('content', '')[-1000:]}\n```"  # Show last 1000 chars
        else:
            display += f"\n{status_icon} {stage['name']}"
    
    return display

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

def resume_session(session_id: str, language: str) -> Generator[Tuple[str, str, str], None, None]:
    """Resume an existing session"""
    if not session_id:
        return
    
    # Process with existing session ID
    yield from process_query("", language, session_id)

def download_novel(novel_text: str, format: str, language: str) -> str:
    """Download novel in specified format"""
    if not novel_text:
        return None
    
    # 페이지 마크 제거
    novel_text = re.sub(r'\[(?:페이지|Page|page)\s*\d+\]', '', novel_text)
    novel_text = re.sub(r'(?:페이지|Page)\s*\d+:', '', novel_text)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    if format == "DOCX" and DOCX_AVAILABLE:
        # Create DOCX
        doc = Document()
        
        # Parse and add content
        lines = novel_text.split('\n')
        for line in lines:
            if line.startswith('#'):
                level = len(line.split()[0])
                text = line.lstrip('#').strip()
                doc.add_heading(text, level)
            elif line.strip():
                doc.add_paragraph(line)
        
        # Save
        temp_dir = tempfile.gettempdir()
        filename = f"Novel_{timestamp}.docx"
        filepath = os.path.join(temp_dir, filename)
        doc.save(filepath)
        
        return filepath
    else:
        # TXT format
        temp_dir = tempfile.gettempdir()
        filename = f"Novel_{timestamp}.txt"
        filepath = os.path.join(temp_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(novel_text)
        
        return filepath

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

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

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

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

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

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

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

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

# Main execution
if __name__ == "__main__":
    logger.info("Starting SOMA Novel Writing System...")
    
    # Initialize database on startup
    logger.info("Initializing database...")
    NovelDatabase.init_db()
    logger.info("Database initialized successfully.")
    
    interface = create_interface()
    
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True
    )