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
import traceback

# --- 로깅 설정 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- 환경 변수 및 상수 ---
FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "")
BRAVE_SEARCH_API_KEY = os.getenv("BRAVE_SEARCH_API_KEY", "")
API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
MODEL_ID = "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507"
BRAVE_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search"
DB_PATH = "screenplay_sessions_korean.db"

# 시나리오 길이 설정
SCREENPLAY_LENGTHS = {
   "영화": {"pages": 120, "description": "장편 영화 (110-130페이지)", "min_pages": 110},
   "드라마": {"pages": 60, "description": "TV 드라마 (55-65페이지)", "min_pages": 55},
   "웹드라마": {"pages": 50, "description": "웹/OTT 시리즈 (45-55페이지)", "min_pages": 45},
   "단편": {"pages": 20, "description": "단편 영화 (15-25페이지)", "min_pages": 15}
}

# 환경 검증
if not FIREWORKS_API_KEY:
   logger.error("FIREWORKS_API_KEY가 설정되지 않았습니다.")
   FIREWORKS_API_KEY = "dummy_token_for_testing"

# 글로벌 변수
db_lock = threading.Lock()

# 전문가 역할 정의
EXPERT_ROLES = {
    "프로듀서": {
        "emoji": "🎬",
        "description": "상업성과 시장성 분석",
        "focus": ["타겟 관객", "제작 가능성", "예산 규모", "마케팅 포인트"],
        "personality": "실용적이고 시장 지향적"
    },
    "스토리작가": {
        "emoji": "📖",
        "description": "내러티브 구조와 플롯 개발",
        "focus": ["3막 구조", "플롯 포인트", "서사 아크", "테마"],
        "personality": "창의적이고 구조적"
    },
    "캐릭터디자이너": {
        "emoji": "👥",
        "description": "인물 창조와 관계 설계",
        "focus": ["캐릭터 아크", "동기부여", "관계 역학", "대화 스타일"],
        "personality": "심리학적이고 공감적"
    },
    "감독": {
        "emoji": "🎭",
        "description": "비주얼 스토리텔링과 연출",
        "focus": ["시각적 구성", "카메라 워크", "미장센", "리듬과 페이싱"],
        "personality": "비주얼 중심적이고 예술적"
    },
    "고증전문가": {
        "emoji": "🔎",
        "description": "사실 확인과 고증",
        "focus": ["역사적 정확성", "과학적 타당성", "문화적 적절성", "현실성"],
        "personality": "정확하고 세심한"
    },
    "비평가": {
        "emoji": "🔍",
        "description": "객관적 분석과 개선점 제시",
        "focus": ["논리적 일관성", "감정적 임팩트", "원작 충실도", "완성도"],
        "personality": "분석적이고 비판적"
    },
    "편집자": {
        "emoji": "✂️",
        "description": "페이싱과 구조 최적화",
        "focus": ["씬 전환", "리듬", "긴장감 조절", "불필요한 부분 제거"],
        "personality": "정밀하고 효율적"
    },
    "대화전문가": {
        "emoji": "💬",
        "description": "대사와 서브텍스트 강화",
        "focus": ["자연스러운 대화", "캐릭터 보이스", "서브텍스트", "감정 전달"],
        "personality": "언어적이고 뉘앙스 중심"
    }
}

# 기획 단계 전문가
PLANNING_STAGES = [
    ("프로듀서", "producer", "🎬 프로듀서: 핵심 컨셉 및 시장성 분석"),
    ("스토리작가", "story_writer", "📖 스토리 작가: 시놉시스 및 3막 구조"),
    ("캐릭터디자이너", "character_designer", "👥 캐릭터 디자이너: 인물 프로필 및 관계도"),
    ("감독", "director", "🎭 감독: 비주얼 컨셉 및 연출 방향"),
    ("비평가", "critic", "🔍 비평가: 기획안 종합 검토 및 개선점"),
]

# 막별 작성 단계
ACT_WRITING_STAGES = {
    "1막": [
        ("스토리작가", "초고", "✍️ 스토리작가: 1막 초고 작성"),
        ("고증전문가", "검증", "🔎 고증전문가: 사실 확인 및 검증"),
        ("편집자", "편집", "✂️ 편집자: 구조 및 페이싱 조정"),
        ("감독", "연출", "🎭 감독: 비주얼 강화 및 연출 노트"),
        ("대화전문가", "대사", "💬 대화전문가: 대사 개선"),
        ("비평가", "검토", "🔍 비평가: 종합 검토 및 평가"),
        ("스토리작가", "완성", "✅ 스토리작가: 1막 최종 완성")
    ],
    "2막A": [
        ("스토리작가", "초고", "✍️ 스토리작가: 2막A 초고 작성"),
        ("고증전문가", "검증", "🔎 고증전문가: 사실 확인 및 검증"),
        ("편집자", "편집", "✂️ 편집자: 구조 및 페이싱 조정"),
        ("감독", "연출", "🎭 감독: 비주얼 강화 및 연출 노트"),
        ("대화전문가", "대사", "💬 대화전문가: 대사 개선"),
        ("비평가", "검토", "🔍 비평가: 종합 검토 및 평가"),
        ("스토리작가", "완성", "✅ 스토리작가: 2막A 최종 완성")
    ],
    "2막B": [
        ("스토리작가", "초고", "✍️ 스토리작가: 2막B 초고 작성"),
        ("고증전문가", "검증", "🔎 고증전문가: 사실 확인 및 검증"),
        ("편집자", "편집", "✂️ 편집자: 구조 및 페이싱 조정"),
        ("감독", "연출", "🎭 감독: 비주얼 강화 및 연출 노트"),
        ("대화전문가", "대사", "💬 대화전문가: 대사 개선"),
        ("비평가", "검토", "🔍 비평가: 종합 검토 및 평가"),
        ("스토리작가", "완성", "✅ 스토리작가: 2막B 최종 완성")
    ],
    "3막": [
        ("스토리작가", "초고", "✍️ 스토리작가: 3막 초고 작성"),
        ("고증전문가", "검증", "🔎 고증전문가: 사실 확인 및 검증"),
        ("편집자", "편집", "✂️ 편집자: 구조 및 페이싱 조정"),
        ("감독", "연출", "🎭 감독: 비주얼 강화 및 연출 노트"),
        ("대화전문가", "대사", "💬 대화전문가: 대사 개선"),
        ("비평가", "검토", "🔍 비평가: 종합 검토 및 평가"),
        ("스토리작가", "완성", "✅ 스토리작가: 3막 최종 완성")
    ]
}

# 웹 검색 클래스
class WebSearcher:
    """사실 확인과 고증을 위한 웹 검색"""
    def __init__(self):
        self.api_key = BRAVE_SEARCH_API_KEY
        self.enabled = bool(self.api_key)
    
    def search(self, query: str, count: int = 3) -> List[Dict]:
        """웹 검색 수행"""
        if not self.enabled:
            return []
        
        headers = {
            "Accept": "application/json",
            "X-Subscription-Token": self.api_key
        }
        
        params = {
            "q": query,
            "count": count,
            "search_lang": "ko",
            "safesearch": "moderate"
        }
        
        try:
            response = requests.get(BRAVE_SEARCH_URL, headers=headers, params=params, timeout=10)
            if response.status_code == 200:
                results = response.json().get("web", {}).get("results", [])
                return results
            else:
                logger.error(f"Search API error: {response.status_code}")
                return []
        except Exception as e:
            logger.error(f"Search error: {e}")
            return []
    
    def verify_facts(self, content: str, context: str) -> Dict:
        """내용의 사실 확인"""
        verification_results = {
            "verified": [],
            "needs_correction": [],
            "suggestions": []
        }
        
        # 주요 키워드 추출
        keywords = self._extract_keywords(content)
        
        for keyword in keywords[:3]:  # 상위 3개만 검색
            search_query = f"{keyword} {context} 사실 확인"
            results = self.search(search_query, count=2)
            
            if results:
                verification_results["verified"].append({
                    "keyword": keyword,
                    "sources": [r["title"] for r in results]
                })
        
        return verification_results
    
    def _extract_keywords(self, content: str) -> List[str]:
        """주요 키워드 추출"""
        # 간단한 키워드 추출 (실제로는 더 정교한 방법 필요)
        words = content.split()
        keywords = [w for w in words if len(w) > 4][:5]
        return keywords

# 다이어그램 생성 클래스
class DiagramGenerator:
    """캐릭터 관계도와 갈등 구조 다이어그램 생성"""
    
    @staticmethod
    def create_character_relationship_diagram(characters: Dict) -> str:
        """캐릭터 관계도 생성 (Mermaid 형식)"""
        diagram = """
```mermaid
graph TB
    classDef protagonist fill:#f9d71c,stroke:#333,stroke-width:2px
    classDef antagonist fill:#ff6b6b,stroke:#333,stroke-width:2px
    classDef supporting fill:#95e1d3,stroke:#333,stroke-width:2px
"""
        
        # 주인공 노드
        if "주인공" in characters:
            protagonist = characters["주인공"]
            diagram += f"\n    P[주인공<br/>{protagonist.get('name', '미정')}]:::protagonist"
        
        # 적대자 노드
        if "적대자" in characters:
            antagonist = characters["적대자"]
            diagram += f"\n    A[적대자<br/>{antagonist.get('name', '미정')}]:::antagonist"
            diagram += f"\n    P -.갈등.- A"
        
        # 조연 노드들
        supporting_count = 0
        for key, char in characters.items():
            if key not in ["주인공", "적대자"] and supporting_count < 5:
                supporting_count += 1
                char_id = f"S{supporting_count}"
                diagram += f"\n    {char_id}[{key}<br/>{char.get('name', '미정')}]:::supporting"
                
                # 관계 연결
                if char.get("relationship_to_protagonist"):
                    diagram += f"\n    P --{char['relationship_to_protagonist']}--> {char_id}"
        
        diagram += "\n```"
        return diagram
    
    @staticmethod
    def create_conflict_structure_diagram(conflicts: List[Dict]) -> str:
        """갈등 구조 다이어그램 생성"""
        diagram = """
```mermaid
flowchart LR
    classDef conflict fill:#ff9999,stroke:#333,stroke-width:2px
    classDef resolution fill:#99ff99,stroke:#333,stroke-width:2px
    
    Start([시작]) --> C1
"""
        
        for i, conflict in enumerate(conflicts[:5], 1):
            conflict_name = conflict.get("name", f"갈등{i}")
            resolution = conflict.get("resolution", "해결")
            
            diagram += f"\n    C{i}[{conflict_name}]:::conflict"
            if i < len(conflicts):
                diagram += f" --> C{i+1}"
            else:
                diagram += f" --> End([{resolution}]):::resolution"
        
        diagram += "\n```"
        return diagram
    
    @staticmethod
    def create_story_arc_diagram(acts: Dict) -> str:
        """스토리 아크 다이어그램 생성"""
        diagram = """
```mermaid
graph LR
    subgraph "1막 - 설정"
        A1[일상] --> A2[사건 발생]
        A2 --> A3[새로운 세계]
    end
    
    subgraph "2막A - 상승"
        B1[도전] --> B2[성장]
        B2 --> B3[중간점]
    end
    
    subgraph "2막B - 하강"
        C1[위기] --> C2[절망]
        C2 --> C3[각성]
    end
    
    subgraph "3막 - 해결"
        D1[최종 대결] --> D2[클라이맥스]
        D2 --> D3[새로운 일상]
    end
    
    A3 --> B1
    B3 --> C1
    C3 --> D1
    
    style A1 fill:#e1f5fe
    style D3 fill:#c8e6c9
```
"""
        return diagram

# 데이터 클래스
@dataclass
class ExpertFeedback:
    """전문가 피드백"""
    role: str
    stage: str
    feedback: str
    suggestions: List[str]
    score: float
    timestamp: datetime = field(default_factory=datetime.now)

@dataclass
class ActProgress:
    """막별 진행 상황"""
    act_name: str
    current_stage: int
    total_stages: int
    current_expert: str
    status: str  # "ready", "in_progress", "complete"
    content: str = ""
    expert_feedbacks: List[ExpertFeedback] = field(default_factory=list)
    
    @property
    def progress_percentage(self):
        if self.total_stages == 0:
            return 0
        return (self.current_stage / self.total_stages) * 100

# 데이터베이스 클래스
class ScreenplayDatabase:
    @staticmethod
    def init_db():
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            cursor = conn.cursor()
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS screenplay_sessions (
                    session_id TEXT PRIMARY KEY,
                    user_query TEXT NOT NULL,
                    screenplay_type TEXT NOT NULL,
                    genre TEXT NOT NULL,
                    target_pages INTEGER,
                    title TEXT,
                    logline TEXT,
                    planning_data TEXT,
                    character_diagram TEXT,
                    conflict_diagram TEXT,
                    act1_content TEXT,
                    act2a_content TEXT,
                    act2b_content TEXT,
                    act3_content TEXT,
                    expert_feedbacks TEXT,
                    created_at TEXT DEFAULT (datetime('now')),
                    updated_at TEXT DEFAULT (datetime('now')),
                    status TEXT DEFAULT 'planning',
                    current_act TEXT DEFAULT '1막',
                    total_pages REAL DEFAULT 0
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS act_progress (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    act_name TEXT NOT NULL,
                    stage_num INTEGER,
                    expert_role TEXT,
                    content TEXT,
                    feedback TEXT,
                    score REAL,
                    created_at TEXT DEFAULT (datetime('now')),
                    FOREIGN KEY (session_id) REFERENCES screenplay_sessions(session_id)
                )
            ''')
            
            conn.commit()

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

    @staticmethod
    def create_session(user_query: str, screenplay_type: str, genre: str) -> str:
        session_id = hashlib.md5(f"{user_query}{screenplay_type}{datetime.now()}".encode()).hexdigest()
        target_pages = SCREENPLAY_LENGTHS[screenplay_type]["pages"]
        
        with ScreenplayDatabase.get_db() as conn:
            conn.cursor().execute(
                '''INSERT INTO screenplay_sessions 
                   (session_id, user_query, screenplay_type, genre, target_pages) 
                   VALUES (?, ?, ?, ?, ?)''',
                (session_id, user_query, screenplay_type, genre, target_pages)
            )
            conn.commit()
        return session_id

    @staticmethod
    def save_planning_data(session_id: str, planning_data: Dict):
        with ScreenplayDatabase.get_db() as conn:
            conn.cursor().execute(
                '''UPDATE screenplay_sessions 
                   SET planning_data = ?, status = 'planned', updated_at = datetime('now')
                   WHERE session_id = ?''',
                (json.dumps(planning_data, ensure_ascii=False), session_id)
            )
            conn.commit()

    @staticmethod
    def save_diagrams(session_id: str, character_diagram: str, conflict_diagram: str):
        """다이어그램 저장"""
        with ScreenplayDatabase.get_db() as conn:
            conn.cursor().execute(
                '''UPDATE screenplay_sessions 
                   SET character_diagram = ?, conflict_diagram = ?, updated_at = datetime('now')
                   WHERE session_id = ?''',
                (character_diagram, conflict_diagram, session_id)
            )
            conn.commit()

    @staticmethod
    def save_act_content(session_id: str, act_name: str, content: str):
        """막별 콘텐츠 저장"""
        act_column = {
            "1막": "act1_content",
            "2막A": "act2a_content",
            "2막B": "act2b_content",
            "3막": "act3_content"
        }.get(act_name, "act1_content")
        
        with ScreenplayDatabase.get_db() as conn:
            conn.cursor().execute(
                f'''UPDATE screenplay_sessions 
                    SET {act_column} = ?, current_act = ?, updated_at = datetime('now')
                    WHERE session_id = ?''',
                (content, act_name, session_id)
            )
            conn.commit()

# 시나리오 생성 시스템
class ScreenplayGenerationSystem:
    def __init__(self):
        self.api_key = FIREWORKS_API_KEY
        self.api_url = API_URL
        self.model_id = MODEL_ID
        self.current_session_id = None
        self.original_query = ""
        self.planning_data = {}
        self.web_searcher = WebSearcher()
        self.diagram_generator = DiagramGenerator()
        self.expert_feedbacks = []
        ScreenplayDatabase.init_db()

    def create_headers(self):
        if not self.api_key or self.api_key == "dummy_token_for_testing":
            raise ValueError("유효한 FIREWORKS_API_KEY가 필요합니다")
        
        return {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }


    
    def call_llm_streaming(self, messages: List[Dict[str, str]], max_tokens: int = 8000) -> Generator[str, None, None]:
        """LLM 호출 - 버퍼 처리 개선"""
        try:
            # 시스템 메시지 강화
            if messages and messages[0].get("role") == "system":
                messages[0]["content"] = messages[0]["content"] + """
    
    【절대 준수 사항】
    1. 모든 응답을 완전한 한국어로 작성하세요.
    2. 시나리오 포맷 용어(INT. EXT. CUT TO 등)는 영어로 유지하세요.
    3. 캐릭터명과 대사는 한국어로 작성하세요.
    4. 미완성 문장(...) 없이 완전하게 작성하세요.
    5. 영화 시나리오 포맷을 정확히 준수하세요.
    """
            
            payload = {
                "model": self.model_id,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": 0.7,  # 약간 높임
                "top_p": 0.9,
                "top_k": 40,
                "presence_penalty": 0.3,
                "frequency_penalty": 0.3,
                "stream": True
            }
            
            headers = self.create_headers()
            
            response = requests.post(
                self.api_url, 
                headers=headers, 
                json=payload, 
                stream=True, 
                timeout=300
            )
        
            if response.status_code != 200:
                yield f"❌ API 오류: {response.status_code}"
                return
            
            buffer = ""
            total_output = ""  # 전체 출력 추적
            
            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]":
                        # 남은 버퍼 모두 출력
                        if buffer:
                            yield buffer
                        break
                    
                    data = json.loads(data_str)
                    if "choices" in data and len(data["choices"]) > 0:
                        content = data["choices"][0].get("delta", {}).get("content", "")
                        if content:
                            buffer += content
                            total_output += content
                            
                            # 버퍼가 충분히 쌓이면 yield (크기 늘림)
                            if len(buffer) >= 500 or '\n\n' in buffer:
                                yield buffer
                                buffer = ""
                                


                except Exception as e:
                    logger.error(f"Line processing error: {e}")
                    continue
            
            # 남은 버퍼 처리
            if buffer:
                yield buffer
                
            # 전체 출력 길이 로깅
            logger.info(f"Total LLM output length: {len(total_output)} characters, {len(total_output.split())} words")
            
        except Exception as e:
            logger.error(f"LLM streaming error: {e}")
            yield f"❌ 오류: {str(e)}"

    def generate_planning(self, query: str, screenplay_type: str, genre: str) -> Generator[Tuple[str, float, Dict, List, str], None, None]:
        """기획안 생성 with 다이어그램"""
        try:
            self.original_query = query
            self.current_session_id = ScreenplayDatabase.create_session(query, screenplay_type, genre)
            
            planning_content = {}
            self.expert_feedbacks = []
            total_stages = len(PLANNING_STAGES)
            character_diagram = ""
            conflict_diagram = ""
            
            for idx, (role, stage_key, stage_desc) in enumerate(PLANNING_STAGES):
                try:
                    progress = ((idx + 1) / total_stages) * 100
                    
                    # 상태 업데이트
                    yield f"🔄 {stage_desc} 진행 중...", progress, planning_content, self.expert_feedbacks, ""
                    
                    # 전문가별 프롬프트 생성
                    prompt = self._get_expert_prompt(role, stage_key, query, planning_content, genre, screenplay_type)
                    
                    # LLM 호출
                    messages = [
                        {"role": "system", "content": f"당신은 {role}입니다. {EXPERT_ROLES[role]['description']}"},
                        {"role": "user", "content": prompt}
                    ]
                    
                    content = ""
                    for chunk in self.call_llm_streaming(messages):
                        if chunk and not chunk.startswith("❌"):
                            content += chunk
                            planning_content[f"{role}_{stage_key}"] = content
                            # 중간 상태 업데이트
                            yield f"✏️ {stage_desc} 작성 중...", progress, planning_content, self.expert_feedbacks, ""
                    
                    # 캐릭터 디자이너 단계에서 다이어그램 생성
                    if role == "캐릭터디자이너" and content:
                        characters = self._extract_characters_from_content(content)
                        character_diagram = self.diagram_generator.create_character_relationship_diagram(characters)
                        conflicts = self._extract_conflicts_from_content(content)
                        conflict_diagram = self.diagram_generator.create_conflict_structure_diagram(conflicts)
                        
                        # 다이어그램 저장
                        ScreenplayDatabase.save_diagrams(self.current_session_id, character_diagram, conflict_diagram)
                    
                    # 전문가 피드백 생성
                    feedback = ExpertFeedback(
                        role=role,
                        stage=stage_key,
                        feedback=content[:500] if content else "작성 중...",
                        suggestions=self._extract_suggestions(content),
                        score=85.0 + random.uniform(-5, 10)
                    )
                    self.expert_feedbacks.append(feedback)
                    
                    time.sleep(0.5)  # API 제한 고려
                    
                except GeneratorExit:
                    logger.warning(f"Generator exit at stage {idx}")
                    break
                except Exception as stage_error:
                    logger.error(f"Error in stage {idx}: {stage_error}")
                    continue
            
            # 최종 저장
            ScreenplayDatabase.save_planning_data(self.current_session_id, planning_content)
            
            # 다이어그램 포함한 최종 결과
            final_diagrams = f"\n\n## 📊 캐릭터 관계도\n{character_diagram}\n\n## ⚔️ 갈등 구조\n{conflict_diagram}"
            
            yield "✅ 기획안 완성!", 100, planning_content, self.expert_feedbacks, final_diagrams
            
        except Exception as e:
            logger.error(f"Planning generation error: {e}\n{traceback.format_exc()}")
            yield f"❌ 오류 발생: {str(e)}", 0, {}, [], ""

    def _get_expert_prompt(self, role: str, stage: str, query: str, 
                          previous_content: Dict, genre: str, screenplay_type: str) -> str:
        """각 전문가별 프롬프트 생성"""
        
        expert = EXPERT_ROLES[role]
        focus_areas = ", ".join(expert["focus"])
        
        base_prompt = f"""
【당신의 역할】
당신은 {role}입니다. {expert['description']}
집중 영역: {focus_areas}

【원본 요청】
{query}

【장르】 {genre}
【형식】 {screenplay_type}
"""
        
        if role == "캐릭터디자이너":
            return base_prompt + """
【작성 요구사항】
1. 주인공 프로필
   - 이름, 나이, 직업
   - 성격과 특징
   - 목표와 동기
   - 성장 아크

2. 적대자 프로필
   - 이름과 역할
   - 주인공과의 갈등
   - 동기와 목표

3. 조연 캐릭터들 (3-5명)
   - 각자의 역할
   - 주인공과의 관계

4. 캐릭터 관계 설명
   - 인물간 관계
   - 주요 갈등 구조

구체적으로 작성하세요."""
        
        # 다른 역할들의 프롬프트...
        return base_prompt + "\n구체적이고 전문적인 분석을 제공하세요."

    def _extract_characters_from_content(self, content: str) -> Dict:
        """콘텐츠에서 캐릭터 정보 추출"""
        characters = {}
        
        # 간단한 파싱 (실제로는 더 정교한 방법 필요)
        if "주인공" in content:
            characters["주인공"] = {"name": "주인공", "role": "protagonist"}
        if "적대자" in content or "악역" in content:
            characters["적대자"] = {"name": "적대자", "role": "antagonist"}
        
        # 조연 추출
        for i in range(1, 4):
            characters[f"조연{i}"] = {"name": f"조연{i}", "relationship_to_protagonist": "동료"}
        
        return characters

    def _extract_conflicts_from_content(self, content: str) -> List[Dict]:
        """콘텐츠에서 갈등 구조 추출"""
        conflicts = [
            {"name": "초기 갈등", "type": "external"},
            {"name": "내적 갈등", "type": "internal"},
            {"name": "관계 갈등", "type": "relationship"},
            {"name": "최종 대결", "type": "climax", "resolution": "해결"}
        ]
        return conflicts

    def _extract_suggestions(self, content: str) -> List[str]:
        """제안사항 추출"""
        suggestions = []
        if content:
            lines = content.split('\n')
            for line in lines:
                if any(k in line for k in ['개선', '제안', '추천']):
                    suggestions.append(line.strip())
        return suggestions[:5]

    def generate_act(self, session_id: str, act_name: str, planning_data: Dict, 
                    previous_acts: Dict) -> Generator[Tuple[ActProgress, str], None, None]:
        """막별 시나리오 생성 - 개선된 버전"""
        
        # act_progress를 먼저 초기화
        act_progress = ActProgress(
            act_name=act_name,
            current_stage=0,
            total_stages=0,
            current_expert="",
            status="initializing"
        )
        
        try:
            self.current_session_id = session_id
            self.planning_data = planning_data
            
            # 세션 정보 복원
            session_data = ScreenplayGenerationSystem.get_session_data(session_id)

            if session_data:
                self.original_query = session_data.get('user_query', '')
            
            stages = ACT_WRITING_STAGES.get(act_name, [])
            
            # act_progress 업데이트
            act_progress.total_stages = len(stages)
            act_progress.status = "in_progress"
            
            act_content = ""
            
            for idx, (role, stage_type, description) in enumerate(stages):
                try:
                    act_progress.current_stage = idx + 1
                    act_progress.current_expert = role
                    
                    # 진행 상황 업데이트
                    yield act_progress, f"🔄 {description} 진행 중..."
                    
                    # 역할별 프롬프트 생성
                    prompt = self._create_act_prompt(role, act_name, act_content, planning_data, previous_acts, stage_type)
                    
                    # 한국어 및 영화 포맷 강조
                    system_message = f"""당신은 {role}입니다.
{EXPERT_ROLES[role]['description']}

【절대 규칙】
1. 모든 내용을 한국어로 작성하세요.
2. 영화 시나리오 포맷을 사용하세요 (연극 아님).
3. 기획안 내용을 100% 반영하세요.
4. 영어 사용 금지 (INT. EXT. 등 포맷 용어 제외).
5. 완전한 문장으로 작성 (... 사용 금지)."""
                    
                    messages = [
                        {"role": "system", "content": system_message},
                        {"role": "user", "content": prompt}
                    ]
                    
                    expert_output = ""
                    line_count = 0
                    
                    for chunk in self.call_llm_streaming(messages, max_tokens=14000):
                        if chunk and not chunk.startswith("❌"):
        # 영어 필터링 로직 완화
                            expert_output += chunk  # _clean_output 호출 제거 (나중에 한 번만)
                            line_count = len(expert_output.split('\n'))
        
        # 실시간 업데이트
                            yield act_progress, f"✏️ {role}: 작업 중... ({line_count}줄 작성)"

                    if stage_type in ["초고", "완성"]:
                        expert_output = self._clean_output(expert_output)  # 여기서만 정리
                    
                    # 웹 검색 고증
                    if role == "고증전문가" and self.web_searcher.enabled:
                        verification = self.web_searcher.verify_facts(act_content, act_name)
                        expert_output += f"\n\n【웹 검증 결과】\n{json.dumps(verification, ensure_ascii=False, indent=2)}"
                    
                    # 출력 검증 및 정리
                    if stage_type in ["초고", "완성"]:
                        expert_output = self._validate_and_clean_screenplay(expert_output)
                    
                    # 피드백 저장
                    feedback = ExpertFeedback(
                        role=role,
                        stage=f"{act_name}_{stage_type}",
                        feedback=f"{role} 작업 완료. {line_count}줄 작성.",
                        suggestions=self._extract_suggestions(expert_output),
                        score=85.0 + random.uniform(-5, 10)
                    )
                    act_progress.expert_feedbacks.append(feedback)
                    
                    # 최종 완성 단계에서만 콘텐츠 업데이트
                    if stage_type == "완성":
                        act_content = expert_output
                        act_progress.content = act_content
                    elif stage_type == "초고":
                        act_content = expert_output
                    
                    yield act_progress, f"✅ {role} 작업 완료 ({line_count}줄)"
                    
                    time.sleep(0.5)
                    
                except GeneratorExit:
                    logger.warning(f"Generator exit at act stage {idx}")
                    break
                except Exception as e:
                    logger.error(f"Error in act stage {idx}: {e}")
                    yield act_progress, f"⚠️ {role} 작업 중 오류 발생"
                    continue
            
            # 막 완성
            act_progress.status = "complete"
            ScreenplayDatabase.save_act_content(session_id, act_name, act_content)
            
            # 최종 페이지 수 계산
            page_count = len(act_content.split('\n')) / 58
            yield act_progress, f"✅ {act_name} 완성! (약 {page_count:.1f}페이지)"
            
        except Exception as e:
            logger.error(f"Act generation error: {e}\n{traceback.format_exc()}")
            act_progress.status = "error"
            yield act_progress, f"❌ 오류: {str(e)}"

    def _clean_output(self, text: str) -> str:
        """출력 텍스트 정리 - 개선된 버전"""
        
        # 빈 텍스트 체크
        if not text:
            return ""
        
        lines = text.split('\n')
        cleaned_lines = []
        
        for line in lines:
            # 원본 라인 유지 (strip하지 않음)
            original_line = line
            stripped = line.strip()
            
            # 빈 줄은 그대로 유지
            if not stripped:
                cleaned_lines.append(original_line)
                continue
            
            # 시나리오 포맷 라인은 무조건 유지
            # INT., EXT., CUT TO:, FADE 등이 포함된 라인
            format_keywords = ["INT.", "EXT.", "INT ", "EXT ", "CUT TO", "FADE", 
                              "DISSOLVE", "MONTAGE", "FLASHBACK", "THE END"]
            
            is_format_line = any(keyword in stripped.upper() for keyword in format_keywords)
            
            if is_format_line:
                # 포맷 라인은 그대로 유지
                cleaned_lines.append(original_line)
            else:
                # 영어 체크 (새로운 로직 사용)
                if not self._contains_english(stripped):
                    # 불완전한 ... 제거
                    if "..." in original_line:
                        original_line = original_line.replace("...", ".")
                    cleaned_lines.append(original_line)
                else:
                    # 실제 영어가 포함된 경우만 제거
                    logger.debug(f"Removing English line: {stripped[:50]}")
                    # 중요한 대사나 지시문이 아닌 경우만 제거
                    if len(stripped) < 5:  # 너무 짧은 줄은 유지
                        cleaned_lines.append(original_line)
        
        return '\n'.join(cleaned_lines)

    def _contains_english(self, text: str) -> bool:
        """영어 포함 여부 확인 - 시나리오 포맷 친화적 버전"""
        
        # 빈 문자열 체크
        if not text or not text.strip():
            return False
        
        # 시나리오에서 반드시 허용해야 하는 용어들 (대폭 확장)
        allowed_terms = [
            # 씬 헤더
            "INT", "EXT", "INT.", "EXT.", 
            # 트랜지션
            "CUT", "TO", "FADE", "IN", "OUT", "DISSOLVE", "CONT'D",
            # 카메라 용어
            "O.S.", "V.O.", "POV", "CLOSE", "UP", "WIDE", "SHOT",
            "ESTABLISHING", "MONTAGE", "FLASHBACK", "DREAM", "SEQUENCE",
            # 기술 용어
            "CG", "SFX", "VFX", "BGM", "OST", "FX",
            # 일반 약어
            "vs", "ex", "etc", "OK", "NG", "VIP", "CEO", "AI", "IT",
            "DJ", "MC", "PD", "VJ", "PC", "TV", "DVD", "USB", "CD",
            # 알파벳 단독 문자 (A, B, C 등)
            "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M",
            "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z",
            # 씬 관련
            "SCENE", "THE", "END", "TITLE", "SUPER",
            # 기타 허용 단어
            "sir", "Mr", "Ms", "Dr", "Prof"
        ]
        
        # 대소문자 구분 없이 체크
        test_text = text.upper()
        for term in allowed_terms:
            test_text = test_text.replace(term.upper(), "")
        
        # 남은 텍스트에서 실제 영어 단어 찾기
        import re
        # 6자 이상의 연속된 알파벳만 영어로 판단
        english_pattern = re.compile(r'[a-zA-Z]{6,}')
        matches = english_pattern.findall(test_text)
        
        if matches:
            # 추가 필터링: 알려진 한국어 로마자 표기는 제외
            korean_romanization = ['hyung', 'oppa', 'noona', 'dongsaeng', 'sunbae', 'hoobae']
            real_english = []
            for match in matches:
                if match.lower() not in korean_romanization:
                    real_english.append(match)
            
            if real_english:
                logger.debug(f"Real English words found: {real_english}")
                return True
        
        return False

    
    def _validate_and_clean_screenplay(self, content: str) -> str:
        """시나리오 검증 및 정리"""
        if not content:
            return ""
        
        lines = content.split('\n')
        cleaned = []
        
        for line in lines:
            # 빈 줄은 유지
            if not line.strip():
                cleaned.append(line)
                continue
            
            # 불완전한 부분 제거
            if "..." in line and not line.strip().endswith('.'):
                continue
            
            # 영어 체크 (포맷 제외)
            if self._contains_english(line):
                logger.warning(f"Removing English line: {line[:50]}")
                continue
            
            cleaned.append(line)
        
        return '\n'.join(cleaned)
    

        
    @staticmethod
    def get_session_data(session_id: str) -> Optional[Dict]:
        """세션 데이터 조회"""
        with ScreenplayDatabase.get_db() as conn:
            cursor = conn.cursor()
            row = cursor.execute(
                'SELECT * FROM screenplay_sessions WHERE session_id = ?',
                (session_id,)
            ).fetchone()
            if row:
                columns = [description[0] for description in cursor.description]
                return dict(zip(columns, row))
        return None
    
    def _create_act_prompt(self, role: str, act_name: str, current_content: str, 
                          planning_data: Dict, previous_acts: Dict, stage_type: str) -> str:
        """막 작성 프롬프트 생성 - 개선된 버전"""
        
        # 기획안에서 핵심 정보 추출
        title = ""
        genre = ""
        screenplay_type = ""
        logline = ""
        synopsis = ""
        characters = ""
        
        for key, value in planning_data.items():
            if "프로듀서" in key and value:
                # 제목과 로그라인 추출
                if "제목" in value:
                    title_match = re.search(r'제목[:\s]*([^\n]+)', value)
                    if title_match:
                        title = title_match.group(1).strip()
                if "로그라인" in value:
                    logline_match = re.search(r'로그라인[:\s]*([^\n]+)', value)
                    if logline_match:
                        logline = logline_match.group(1).strip()
            elif "스토리" in key and value:
                synopsis = value[:1000]
            elif "캐릭터" in key and value:
                characters = value[:1000]
        
        # 원본 요청 정보 강조
        core_info = f"""
【핵심 정보 - 절대 준수】
원본 요청: {self.original_query}
제목: {title}
로그라인: {logline}
장르: {genre if genre else '드라마'}
형식: {screenplay_type if screenplay_type else '영화'}

⚠️ 중요: 위 정보를 절대적으로 준수하세요. 다른 이야기로 바꾸지 마세요.
⚠️ 중요: 모든 내용을 한국어로만 작성하세요. 영어 사용 금지.
⚠️ 중요: 표준 시나리오 포맷을 정확히 준수하세요.
"""
        
        if role == "스토리작가":
            if stage_type == "초고":
                return f"""【{act_name} 초고 작성】
{core_info}

【기획안 내용】
시놉시스:
{synopsis}

캐릭터:
{characters}

【이전 막 내용】
{self._summarize_previous_acts(previous_acts)}

【작성 요구사항】
1. 반드시 한국어로만 작성
2. 표준 시나리오 포맷:
   INT. 장소 - 시간  (또는 EXT. 장소 - 시간)
   
   장면 설명은 현재형으로 작성.
   
   캐릭터명 (대문자)
   대사는 자연스러운 한국어로.

3. 분량: 최소 1000줄 이상
4. 각 씬: 5-7페이지 분량
5. 기획안 내용 100% 반영

절대 영어를 사용하지 마세요.
절대 다른 이야기로 바꾸지 마세요.
절대 연극이나 다른 형식으로 작성하지 마세요.

{act_name}을 완전한 시나리오 형식으로 작성하세요."""
            
            else:  # 완성
                return f"""【{act_name} 최종 완성】
{core_info}

이전 전문가들의 피드백을 반영하여 최종본을 작성하세요.

【필수 사항】
1. 한국어로만 작성
2. 영화 시나리오 포맷 준수
3. 기획안 스토리 유지
4. 1200줄 이상

깨진 부분이나 ... 같은 미완성 부분 없이 완전하게 작성하세요."""
        
        elif role == "고증전문가":
            return f"""【{act_name} 사실 확인】
{core_info}

다음 시나리오를 검토하세요:
{current_content[:2000]}...

【검토 사항】
1. 사실 오류 확인
2. 시대 고증
3. 영어 사용 부분을 한국어로 수정
4. 깨진 텍스트나 ... 부분 수정

모든 수정 사항을 한국어로 제시하세요."""
        
        elif role == "편집자":
            return f"""【{act_name} 편집】
{core_info}

【편집 지침】
1. 영어가 있다면 모두 한국어로 변경
2. 깨진 부분(...) 완성
3. 기획안과 일치하지 않는 내용 수정
4. 시나리오 포맷 교정

편집된 버전을 완전한 한국어로 제시하세요."""
        
        elif role == "감독":
            return f"""【{act_name} 연출 노트】
{core_info}

【연출 강화】
1. 시각적 디테일 추가
2. 카메라 앵글 제안
3. 분위기 묘사 강화

모든 연출 노트를 한국어로 작성하세요.
영화 시나리오임을 명확히 하세요."""
        
        elif role == "대화전문가":
            return f"""【{act_name} 대사 개선】
{core_info}

【대사 수정】
1. 모든 대사를 자연스러운 한국어로
2. 영어 대사가 있다면 한국어로 번역
3. 캐릭터별 말투 차별화
4. 어색한 번역투 제거

완전한 한국어 대사로 수정하세요."""
        
        elif role == "비평가":
            return f"""【{act_name} 검토】
{core_info}

【평가 항목】
1. 기획안과의 일치도 (최우선)
2. 한국어 사용 (영어 혼용 금지)
3. 시나리오 포맷 준수
4. 스토리 완성도

문제점과 개선 사항을 구체적으로 제시하세요."""
        
        return f"""【{role} 작업】
{core_info}

{act_name}을(를) 검토하고 개선하세요.
반드시 한국어로만 작성하고, 기획안 내용을 충실히 반영하세요."""

    def _summarize_planning(self, planning_data: Dict) -> str:
        """기획안 요약"""
        summary = ""
        for key, value in planning_data.items():
            if value:
                summary += f"[{key}]\n{value[:200]}...\n"
        return summary[:1000]

    def _summarize_previous_acts(self, previous_acts: Dict) -> str:
        """이전 막 요약"""
        summary = ""
        for act, content in previous_acts.items():
            if content:
                summary += f"[{act}]\n{content[:300]}...\n"
        return summary[:1000]


# UI 함수들
def create_act_progress_display(act_progress: ActProgress) -> str:
    """막별 진행 상황 표시"""
    if not act_progress:
        return ""
    
    html = f"""
    <div style="border: 2px solid #667eea; border-radius: 10px; padding: 15px; margin: 10px 0;">
        <h3>{act_progress.act_name} 진행 상황</h3>
        
        <div style="margin: 10px 0;">
            <div style="background: #e0e0e0; border-radius: 10px; height: 30px; position: relative;">
                <div style="background: linear-gradient(90deg, #667eea, #764ba2); 
                           height: 100%; border-radius: 10px; width: {act_progress.progress_percentage}%; 
                           transition: width 0.3s ease;">
                    <span style="position: absolute; left: 50%; transform: translateX(-50%); 
                                 color: white; line-height: 30px; font-weight: bold;">
                        {act_progress.progress_percentage:.0f}%
                    </span>
                </div>
            </div>
        </div>
        
        <div style="margin-top: 15px;">
            <strong>현재 작업:</strong> {EXPERT_ROLES.get(act_progress.current_expert, {}).get('emoji', '')} {act_progress.current_expert}
            <br>
            <strong>단계:</strong> {act_progress.current_stage} / {act_progress.total_stages}
            <br>
            <strong>상태:</strong> {act_progress.status}
        </div>
    </div>
    """
    
    return html

def format_planning_with_diagrams(planning_data: Dict, diagrams: str) -> str:
    """기획안과 다이어그램 포맷팅"""
    formatted = "# 📋 시나리오 기획안\n\n"
    
    for key, content in planning_data.items():
        if content:
            role = key.split('_')[0]
            emoji = EXPERT_ROLES.get(role, {}).get('emoji', '📝')
            formatted += f"## {emoji} {key}\n\n{content}\n\n---\n\n"
    
    if diagrams:
        formatted += diagrams
    
    return formatted

# Gradio 인터페이스
def create_interface():
    css = """
    .main-header {
        text-align: center;
        margin-bottom: 2rem;
        padding: 2.5rem;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 15px;
        color: white;
        box-shadow: 0 10px 30px rgba(0,0,0,0.3);
    }
    
    /* 시나리오 출력 스타일 */
    .screenplay-output {
        font-family: 'Courier New', monospace;
        font-size: 14px;
        line-height: 1.6;
        white-space: pre-wrap;
        background: #f8f9fa;
        padding: 20px;
        border-radius: 8px;
    }
    
    .screenplay-output h3 {
        color: #2c3e50;
        margin: 20px 0 10px 0;
        font-size: 16px;
        font-weight: bold;
    }
    
    .screenplay-output strong {
        color: #34495e;
        font-weight: bold;
    }

    .screenplay-output em {
        font-style: italic;
        color: #7f8c8d;
    }
    
    .act-button {
        min-height: 80px;
        font-size: 1.1rem;
        margin: 5px;
        border-radius: 10px;
    }
    
    .expert-card {
        background: white;
        border: 1px solid #e0e0e0;
        border-radius: 8px;
        padding: 10px;
        margin: 5px;
        display: inline-block;
    }
    
    .progress-container {
        background: #f8f9fa;
        border-radius: 10px;
        padding: 20px;
        margin: 15px 0;
    }
    """

    with gr.Blocks(theme=gr.themes.Soft(), css=css, title="AI 시나리오 작가") as interface:
        gr.HTML("""
        <div class="main-header">
            <h1>🎬 AI 시나리오 작가</h1>
            <p>전문가 협업 시스템 | 다이어그램 지원 | 웹 검증</p>
        </div>
        """)
        
        # 상태 변수
        current_session_id = gr.State(None)
        current_planning = gr.State({})
        act1_content = gr.State("")
        act2a_content = gr.State("")
        act2b_content = gr.State("")
        act3_content = gr.State("")
        
        with gr.Tabs():
            # 기획 탭
            with gr.Tab("📋 기획안 작성"):
                with gr.Row():
                    with gr.Column(scale=2):
                        query_input = gr.Textbox(
                            label="💡 시나리오 아이디어",
                            placeholder="구체적인 스토리를 입력하세요...",
                            lines=4
                        )
                    
                    with gr.Column(scale=1):
                        screenplay_type = gr.Radio(
                            choices=list(SCREENPLAY_LENGTHS.keys()),
                            value="영화",
                            label="📽️ 유형"
                        )
                        genre_select = gr.Dropdown(
                            choices=["액션", "스릴러", "드라마", "코미디", "공포", "SF", "로맨스", "판타지"],
                            value="드라마",
                            label="🎭 장르"
                        )
                
                planning_btn = gr.Button("📋 기획안 생성 (다이어그램 포함)", variant="primary", size="lg")
                
                # 진행 상태
                with gr.Row():
                    progress_bar = gr.Slider(0, 100, value=0, label="진행률", interactive=False)
                    status_text = gr.Textbox(label="상태", value="대기 중", interactive=False)
                
                # 기획안과 다이어그램
                planning_display = gr.Markdown("*기획안과 다이어그램이 여기에 표시됩니다...*")
            
            # 시나리오 작성 탭
            with gr.Tab("🎬 시나리오 작성"):
                gr.Markdown("""
                ### 📝 막별 시나리오 작성
                
                기획안 생성 후, 각 막을 순차적으로 작성하세요.
                각 막마다 7명의 전문가가 협업합니다:
                
                📖 스토리작가 → 🔎 고증전문가 → ✂️ 편집자 → 🎭 감독 → 💬 대화전문가 → 🔍 비평가 → ✅ 최종완성
                """)
                
                # 막별 작성 버튼들
                with gr.Row():
                    act1_btn = gr.Button(
                        "1️⃣ 1막 작성\n[설정 - 25%]\n약 30페이지", 
                        variant="primary",
                        elem_classes=["act-button"]
                    )
                    act2a_btn = gr.Button(
                        "2️⃣ 2막A 작성\n[상승 - 25%]\n약 30페이지", 
                        variant="secondary",
                        elem_classes=["act-button"],
                        interactive=False  # 1막 완성 후 활성화
                    )
                    act2b_btn = gr.Button(
                        "3️⃣ 2막B 작성\n[복잡화 - 25%]\n약 30페이지", 
                        variant="secondary",
                        elem_classes=["act-button"],
                        interactive=False  # 2막A 완성 후 활성화
                    )
                    act3_btn = gr.Button(
                        "4️⃣ 3막 작성\n[해결 - 25%]\n약 30페이지", 
                        variant="secondary",
                        elem_classes=["act-button"],
                        interactive=False  # 2막B 완성 후 활성화
                    )
                
                # 진행 상황 표시
                gr.Markdown("### 📊 작성 진행 상황")
                act_progress_display = gr.HTML(
                    value='<div class="progress-container">막을 선택하여 작성을 시작하세요...</div>'
                )
                act_status_text = gr.Textbox(
                    label="현재 상태",
                    value="기획안을 먼저 생성한 후 시나리오 작성을 시작하세요",
                    interactive=False
                )
                
                # 전문가 작업 현황
                gr.Markdown("### 👥 전문가 작업 현황")
                expert_status = gr.HTML("""
                <div style="padding: 10px;">
                    <span class="expert-card">📖 스토리작가</span>
                    <span class="expert-card">🔎 고증전문가</span>
                    <span class="expert-card">✂️ 편집자</span>
                    <span class="expert-card">🎭 감독</span>
                    <span class="expert-card">💬 대화전문가</span>
                    <span class="expert-card">🔍 비평가</span>
                </div>
                """)
                
                # 시나리오 출력
                gr.Markdown("### 📄 작성된 시나리오")
                with gr.Tabs():
                    with gr.Tab("1막"):
                        act1_output = gr.Markdown("*1막이 여기에 표시됩니다...*")
                    with gr.Tab("2막A"):
                        act2a_output = gr.Markdown("*2막A가 여기에 표시됩니다...*")
                    with gr.Tab("2막B"):
                        act2b_output = gr.Markdown("*2막B가 여기에 표시됩니다...*")
                    with gr.Tab("3막"):
                        act3_output = gr.Markdown("*3막이 여기에 표시됩니다...*")
                    with gr.Tab("전체 시나리오"):
                        full_output = gr.Markdown("*전체 시나리오가 여기에 표시됩니다...*")
                
                # 다운로드 버튼
                with gr.Row():
                    download_btn = gr.Button("💾 전체 시나리오 다운로드 (TXT)", variant="secondary")
                    download_file = gr.File(label="다운로드 파일", visible=False)
        
        # 이벤트 핸들러
        def handle_planning(query, s_type, genre):
            if not query:
                yield 0, "❌ 아이디어를 입력하세요", "*기획안이 여기에 표시됩니다...*", None, None
                return
            
            system = ScreenplayGenerationSystem()
            session_id = None
            
            for status, progress, planning_data, feedbacks, diagrams in system.generate_planning(query, s_type, genre):
                formatted = format_planning_with_diagrams(planning_data, diagrams)
                session_id = system.current_session_id
                yield progress, status, formatted, planning_data, session_id
        
        # 4. 이벤트 핸들러 수정 - act 내용 저장 개선
        def handle_act_writing(act_name, session_id, planning_data, previous_acts):
            """막별 시나리오 작성 - generator 래퍼"""
            if not session_id or not planning_data:
                yield "", "❌ 먼저 기획안을 생성해주세요", ""
                return
            
            try:
                system = ScreenplayGenerationSystem()
                final_content = ""  # 최종 콘텐츠 저장
                
                # Generator를 순회하며 yield
                for act_progress, status_msg in system.generate_act(
                    session_id, act_name, planning_data, previous_acts
                ):
                    progress_html = create_act_progress_display(act_progress)
                    screenplay_display = format_screenplay_display(act_progress.content)
                    final_content = act_progress.content  # 원본 콘텐츠 저장
                    
                    yield progress_html, status_msg, screenplay_display
                
                # 완료 후 원본 콘텐츠 반환
                if final_content:
                    yield progress_html, f"✅ {act_name} 완성 및 저장됨", screenplay_display
                    
            except Exception as e:
                logger.error(f"Act writing error: {e}")
                yield f"<div>❌ 오류 발생: {str(e)}</div>", f"오류: {str(e)}", ""

        def format_screenplay_display(content):
            """시나리오 표시 포맷팅 - 개선된 버전"""
            if not content:
                return "*작성 중...*"
    
            lines = content.split('\n')
            result = []
            scene_number = 0
    
            for line in lines:
                stripped = line.strip()
        
                if not stripped:
                    result.append("")
                    continue     
                    
        # 씬 헤더 감지 및 포맷팅 (INT. 또는 EXT.로 시작)
                if stripped.startswith(('INT.', 'EXT.', 'INT ', 'EXT ')):
                    scene_number += 1
                    # 씬 구분선 추가
                    result.append("\n" + "="*80)
                    result.append(f"### 씬 {scene_number}")
                    result.append(f"**{stripped}**")
                    result.append("-"*40)
        # 캐릭터명 감지 (대문자로만 된 짧은 줄)
                elif stripped.isupper() and len(stripped.split()) <= 3 and not any(c.isdigit() for c in stripped):
                    # 캐릭터명 강조
                    result.append(f"\n**{stripped}**")
        # 괄호로 된 지시문 (연기 지시)
                elif stripped.startswith('(') and stripped.endswith(')'):
                    result.append(f"*{stripped}*")
                # 트랜지션 (CUT TO:, FADE IN: 등)
                elif any(trans in stripped.upper() for trans in ['CUT TO:', 'FADE IN:', 'FADE OUT:', 'DISSOLVE TO:']):
                    result.append(f"\n_{stripped}_\n")
                # 일반 텍스트
                else:
                    result.append(line)
    
    # 최종 정리
            formatted = '\n'.join(result)
    
    # 연속된 빈 줄 제거
            while '\n\n\n' in formatted:
                formatted = formatted.replace('\n\n\n', '\n\n')
    
            return formatted

        # 2. combine_all_acts 함수 수정 - State 값 올바르게 처리
        def combine_all_acts(act1_state, act2a_state, act2b_state, act3_state):
            """모든 막 합치기 - State 값 처리 개선"""
            full = "# 🎬 전체 시나리오\n\n"
            
            # State 객체인 경우 value 속성 접근
            act1 = act1_state.value if hasattr(act1_state, 'value') else act1_state
            act2a = act2a_state.value if hasattr(act2a_state, 'value') else act2a_state
            act2b = act2b_state.value if hasattr(act2b_state, 'value') else act2b_state
            act3 = act3_state.value if hasattr(act3_state, 'value') else act3_state
            
            # 실제 콘텐츠가 있는지 확인
            has_content = False
            
            if act1 and act1 != "*1막이 여기에 표시됩니다...*" and act1 != "":
                # 마크다운 제거한 원본 텍스트 추출
                clean_act1 = act1.replace("**", "").replace("###", "").replace("="*80, "").replace("-"*40, "")
                full += "## 제1막 - 설정\n\n" + clean_act1 + "\n\n" + "="*80 + "\n\n"
                has_content = True
                
            if act2a and act2a != "*2막A가 여기에 표시됩니다...*" and act2a != "":
                clean_act2a = act2a.replace("**", "").replace("###", "").replace("="*80, "").replace("-"*40, "")
                full += "## 제2막A - 상승\n\n" + clean_act2a + "\n\n" + "="*80 + "\n\n"
                has_content = True
                
            if act2b and act2b != "*2막B가 여기에 표시됩니다...*" and act2b != "":
                clean_act2b = act2b.replace("**", "").replace("###", "").replace("="*80, "").replace("-"*40, "")
                full += "## 제2막B - 복잡화\n\n" + clean_act2b + "\n\n" + "="*80 + "\n\n"
                has_content = True
                
            if act3 and act3 != "*3막이 여기에 표시됩니다...*" and act3 != "":
                clean_act3 = act3.replace("**", "").replace("###", "").replace("="*80, "").replace("-"*40, "")
                full += "## 제3막 - 해결\n\n" + clean_act3 + "\n\n"
                has_content = True
            
            if not has_content:
                return "*전체 시나리오가 완성되면 여기에 표시됩니다...*"
            
            return full

        # 3. save_to_file 함수 수정 - 데이터베이스에서 직접 가져오기
        def save_to_file(full_screenplay, session_id):
            """시나리오를 파일로 저장 - 개선된 버전"""
            
            # 세션 ID가 있으면 DB에서 직접 가져오기
            if session_id:
                try:
                    with ScreenplayDatabase.get_db() as conn:
                        cursor = conn.cursor()
                        row = cursor.execute(
                            '''SELECT user_query, title, logline, genre, screenplay_type,
                                      act1_content, act2a_content, act2b_content, act3_content
                               FROM screenplay_sessions WHERE session_id = ?''',
                            (session_id,)
                        ).fetchone()
                        
                        if row:
                            # 파일 내용 구성
                            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                            filename = f"screenplay_{timestamp}.txt"
                            
                            with open(filename, 'w', encoding='utf-8') as f:
                                # 헤더 정보
                                f.write("="*80 + "\n")
                                f.write("시나리오\n")
                                f.write("="*80 + "\n\n")
                                
                                if row['title']:
                                    f.write(f"제목: {row['title']}\n")
                                if row['genre']:
                                    f.write(f"장르: {row['genre']}\n")
                                if row['screenplay_type']:
                                    f.write(f"형식: {row['screenplay_type']}\n")
                                if row['logline']:
                                    f.write(f"로그라인: {row['logline']}\n")
                                
                                f.write("\n" + "="*80 + "\n\n")
                                
                                # 각 막 내용 추가
                                if row['act1_content']:
                                    f.write("제1막 - 설정\n")
                                    f.write("="*80 + "\n\n")
                                    f.write(row['act1_content'])
                                    f.write("\n\n" + "="*80 + "\n\n")
                                
                                if row['act2a_content']:
                                    f.write("제2막A - 상승\n")
                                    f.write("="*80 + "\n\n")
                                    f.write(row['act2a_content'])
                                    f.write("\n\n" + "="*80 + "\n\n")
                                
                                if row['act2b_content']:
                                    f.write("제2막B - 복잡화\n")
                                    f.write("="*80 + "\n\n")
                                    f.write(row['act2b_content'])
                                    f.write("\n\n" + "="*80 + "\n\n")
                                
                                if row['act3_content']:
                                    f.write("제3막 - 해결\n")
                                    f.write("="*80 + "\n\n")
                                    f.write(row['act3_content'])
                                    f.write("\n\n")
                                
                                # 작성 정보
                                f.write("\n" + "="*80 + "\n")
                                f.write(f"작성일: {datetime.now().strftime('%Y년 %m월 %d일 %H:%M')}\n")
                                f.write("="*80 + "\n")
                            
                            return filename
                except Exception as e:
                    logger.error(f"DB에서 파일 저장 중 오류: {e}")
            
            # DB 접근 실패시 전달받은 내용으로 저장
            if not full_screenplay or full_screenplay == "*전체 시나리오가 완성되면 여기에 표시됩니다...*":
                return None
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"screenplay_{timestamp}.txt"
            
            with open(filename, 'w', encoding='utf-8') as f:
                # 모든 마크다운 및 포맷팅 제거
                clean_text = full_screenplay
                clean_text = clean_text.replace("**", "")
                clean_text = clean_text.replace("###", "")
                clean_text = clean_text.replace("#", "")
                clean_text = clean_text.replace("="*80, "-"*80)
                clean_text = clean_text.replace("-"*40, "")
                clean_text = clean_text.replace("_", "")
                clean_text = clean_text.replace("*", "")
                
                f.write(clean_text)
            
            return filename
        
        # 이벤트 연결
        planning_output = planning_btn.click(
            fn=handle_planning,
            inputs=[query_input, screenplay_type, genre_select],
            outputs=[progress_bar, status_text, planning_display, current_planning, current_session_id]
        )
        
        planning_output.then(
            # 기획안 생성 후 1막 버튼 활성화
            lambda: gr.update(interactive=True),
            outputs=[act1_btn]
        )
        
        # 1막 작성 - generator 처리를 위해 수정
        act1_output_event = act1_btn.click(
            fn=handle_act_writing,
            inputs=[
                gr.State("1막"),
                current_session_id,
                current_planning,
                gr.State({})
            ],
            outputs=[act_progress_display, act_status_text, act1_output]
        )
        
        act1_output_event.then(
            lambda x: (x, gr.update(interactive=True)),
            inputs=[act1_output],
            outputs=[act1_content, act2a_btn]
        )
        
        # 2막A 작성
        act2a_output_event = act2a_btn.click(
            fn=handle_act_writing,
            inputs=[
                gr.State("2막A"),
                current_session_id,
                current_planning,
                act1_content
            ],
            outputs=[act_progress_display, act_status_text, act2a_output]
        )
        
        act2a_output_event.then(
            lambda x: (x, gr.update(interactive=True)),
            inputs=[act2a_output],
            outputs=[act2a_content, act2b_btn]
        )
        
        # 2막B 작성
        act2b_output_event = act2b_btn.click(
            fn=handle_act_writing,
            inputs=[
                gr.State("2막B"),
                current_session_id,
                current_planning,
                gr.State(lambda: {"1막": act1_content.value, "2막A": act2a_content.value})
            ],
            outputs=[act_progress_display, act_status_text, act2b_output]
        )
        
        act2b_output_event.then(
            lambda x: (x, gr.update(interactive=True)),
            inputs=[act2b_output],
            outputs=[act2b_content, act3_btn]
        )
        
        # 3막 작성
        act3_output_event = act3_btn.click(
            fn=handle_act_writing,
            inputs=[
                gr.State("3막"),
                current_session_id,
                current_planning,
                gr.State(lambda: {"1막": act1_content.value, "2막A": act2a_content.value, "2막B": act2b_content.value})
            ],
            outputs=[act_progress_display, act_status_text, act3_output]
        )
        
        act3_output_event.then(
            lambda x: x,
            inputs=[act3_output],
            outputs=[act3_content]
        ).then(
            # 모든 막 완성 후 전체 시나리오 생성
            fn=combine_all_acts,
            inputs=[act1_content, act2a_content, act2b_content, act3_content],
            outputs=[full_output]
        )
        
        # 다운로드
        download_btn.click(
            fn=lambda full, sid: save_to_file(full, sid),
            inputs=[full_output, current_session_id],  # session_id 추가
            outputs=[download_file]
        ).then(
            lambda x: gr.update(visible=True, value=x) if x else gr.update(visible=False),
            inputs=[download_file],
            outputs=[download_file]
        )


    return interface

# 메인 실행
if __name__ == "__main__":
    logger.info("AI 시나리오 작가 시작...")
    ScreenplayDatabase.init_db()
    interface = create_interface()
    interface.launch(server_name="0.0.0.0", server_port=7860, share=False)