import gradio as gr
import pixeltable as pxt
import numpy as np
from datetime import datetime
from pixeltable.functions.huggingface import sentence_transformer
from pixeltable.functions import openai
import os
import getpass
import re
import random

# Set up OpenAI API key
if 'OPENAI_API_KEY' not in os.environ:
    os.environ['OPENAI_API_KEY'] = getpass.getpass('Enter your OpenAI API key: ')

# Initialize Pixeltable
pxt.drop_dir('ai_rpg', force=True)
pxt.create_dir('ai_rpg')

# 일반 함수로 변경 (UDF가 아님)
def initialize_stats(genre: str) -> str:
    """Initialize player stats based on the selected genre"""
    base_stats = {
        "🧙‍♂️ Fantasy": "체력: 100, 마나: 80, 힘: 7, 지능: 8, 민첩: 6, 소지금: 50골드",
        "🚀 Sci-Fi": "체력: 100, 에너지: 90, 기술력: 8, 지능: 9, 민첩: 6, 크레딧: 500",
        "👻 Horror": "체력: 80, 정신력: 100, 힘: 6, 지능: 7, 민첩: 8, 소지품: 손전등, 기본 약품",
        "🔍 Mystery": "체력: 90, 집중력: 100, 관찰력: 9, 지능: 8, 카리스마: 7, 단서: 0",
        "🌋 Post-Apocalyptic": "체력: 95, 방사능 저항: 75, 힘: 8, 생존력: 9, 물자: 제한됨",
        "🤖 Cyberpunk": "체력: 90, 사이버웨어: 85%, 해킹: 8, 거리 신용도: 6, 엣지: 7, 누엔: 1000",
        "⚙️ Steampunk": "체력: 95, 증기력: 85, 기계공학: 8, 예술성: 7, 사교성: 6, 실링: 200"
    }
    
    if genre in base_stats:
        return base_stats[genre]
    else:
        # Default stats if genre not found
        return "체력: 100, 에너지: 100, 힘: 7, 지능: 7, 민첩: 7, 소지금: 100"

@pxt.udf
def generate_random_event(turn_number: int) -> str:
    """Generate a random event based on turn number"""
    if turn_number % 3 == 0 and turn_number > 0:  # Every 3rd turn
        events = [
            "갑자기 부근에서 이상한 소리가 들립니다",
            "낯선 여행자가 당신을 바라보고 있습니다",
            "지면이 미세하게 진동하기 시작합니다",
            "주머니에서 무언가가 빛납니다",
            "멀리서 무언가가 당신을 향해 다가오고 있습니다",
            "갑자기 날씨가 변하기 시작합니다",
            "주변에 숨겨진 통로를 발견합니다"
        ]
        return random.choice(events)
    return ""

@pxt.udf
def generate_messages(genre: str, player_name: str, initial_scenario: str, player_input: str, turn_number: int, stats: str) -> list[dict]:
    return [
        {
            'role': 'system',
            'content': f"""반드시 한국어(한글)로 작성하라. You are the game master for a {genre} RPG. The player's name is {player_name}.
            
            관리해야 할 플레이어 스탯: {stats}
            
            당신은 플레이어의 선택에 따라 스토리를 생생하게 전개하는 게임 마스터입니다. 
            상세한 설명과 감각적인 묘사를 통해 플레이어가 게임 속 세계에 몰입할 수 있도록 하세요.
            
            플레이어의 선택에 따라 스탯이 변하는 경우 이를 스토리에 반영하세요. 
            위험한 상황, 도전, 보상, 우연한 만남이 포함된 흥미로운 스토리를 만드세요.
            
            Provide your response in three clearly separated sections using exactly this format:

            📜 **STORY**: [Your engaging narrative response to the player's action with vivid descriptions]

            📊 **STATS UPDATE**: [Brief update on any changes to player stats based on their actions]

            🎯 **OPTIONS**:
            1. [A dialogue option with potential consequences]
            2. [An action they could take with different outcomes]
            3. [A unique or unexpected choice that might lead to adventure]
            4. [A risky but potentially rewarding option]"""
        },
        {
            'role': 'user',
            'content': f"Current scenario: {initial_scenario}\n"
                      f"Player's action: {player_input}\n"
                      f"Turn number: {turn_number}\n"
                      f"Current player stats: {stats}\n\n"
                      "Provide the story response, stats update, and options:"
        }
    ]

@pxt.udf
def get_story(response: str) -> str:
    """Extract just the story part from the response"""
    match = re.search(r'📜\s*\*\*STORY\*\*:\s*(.*?)(?=📊\s*\*\*STATS|$)', response, re.DOTALL)
    if match:
        return match.group(1).strip()
    parts = response.split("STATS UPDATE:")
    if len(parts) > 1:
        story_part = parts[0].replace("STORY:", "").replace("📜", "").replace("**STORY**:", "").strip()
        return story_part
    return response

@pxt.udf
def get_stats_update(response: str) -> str:
    """Extract the stats update from the response"""
    match = re.search(r'📊\s*\*\*STATS UPDATE\*\*:\s*(.*?)(?=🎯\s*\*\*OPTIONS\*\*|$)', response, re.DOTALL)
    if match:
        return match.group(1).strip()
    parts = response.split("STATS UPDATE:")
    if len(parts) > 1:
        stats_part = parts[1].split("OPTIONS:")[0].strip()
        return stats_part
    return "스탯 변화 없음"

@pxt.udf
def get_options(response: str) -> list[str]:
    """Extract the options from the response"""
    match = re.search(r'🎯\s*\*\*OPTIONS\*\*:\s*(.*?)(?=$)', response, re.DOTALL)
    if match:
        options_text = match.group(1)
        options = re.findall(r'\d+\.\s*(.*?)(?=\d+\.|$)', options_text, re.DOTALL)
        options = [opt.strip() for opt in options if opt.strip()]
        while len(options) < 4:
            options.append("다른 행동 시도...")
        return options[:4]
    
    parts = response.split("OPTIONS:")
    if len(parts) > 1:
        options = re.findall(r'\d+\.\s*(.*?)(?=\d+\.|$)', parts[1], re.DOTALL)
        options = [opt.strip() for opt in options if opt.strip()]
        while len(options) < 4:
            options.append("다른 행동 시도...")
        return options[:4]
    
    return ["계속하기...", "다른 행동 취하기", "뭔가 새로운 시도하기", "주변 탐색하기"]

# Create a single table for all game data
interactions = pxt.create_table(
    'ai_rpg.interactions',
    {
        'session_id': pxt.String,
        'player_name': pxt.String,
        'genre': pxt.String,
        'initial_scenario': pxt.String,
        'turn_number': pxt.Int,
        'player_input': pxt.String,
        'timestamp': pxt.Timestamp,
        'player_stats': pxt.String,
        'random_event': pxt.String
    }
)

# Add computed columns for AI responses
interactions.add_computed_column(messages=generate_messages(
    interactions.genre,
    interactions.player_name,
    interactions.initial_scenario,
    interactions.player_input,
    interactions.turn_number,
    interactions.player_stats
))

interactions.add_computed_column(ai_response=openai.chat_completions(
    messages=interactions.messages,
    model='gpt-4.1-mini',
    max_tokens=800,
    temperature=0.8
))

interactions.add_computed_column(full_response=interactions.ai_response.choices[0].message.content)
interactions.add_computed_column(story_text=get_story(interactions.full_response))
interactions.add_computed_column(stats_update=get_stats_update(interactions.full_response))
interactions.add_computed_column(options=get_options(interactions.full_response))

class RPGGame:
    def __init__(self):
        self.current_session_id = None
        self.turn_number = 0
        self.current_stats = ""
    
    def start_game(self, player_name: str, genre: str, scenario: str) -> tuple[str, str, str, list[str]]:
        session_id = f"session_{datetime.now().strftime('%Y%m%d%H%M%S')}_{player_name}"
        self.current_session_id = session_id
        self.turn_number = 0
        
        # 함수를 직접 호출하여 결과를 저장
        initial_stats = initialize_stats(genre)
        self.current_stats = initial_stats
        
        interactions.insert([{
            'session_id': session_id,
            'player_name': player_name,
            'genre': genre,
            'initial_scenario': scenario,
            'turn_number': 0,
            'player_input': "Game starts",
            'timestamp': datetime.now(),
            'player_stats': initial_stats,  # 문자열 결과 저장
            'random_event': ""
        }])
        
        result = interactions.select(
            interactions.story_text,
            interactions.stats_update,
            interactions.options
        ).where(
            (interactions.session_id == session_id) & 
            (interactions.turn_number == 0)
        ).collect()
        
        return session_id, result['story_text'][0], result['stats_update'][0], result['options'][0]
    
    def process_action(self, action: str) -> tuple[str, str, list[str]]:
        if not self.current_session_id:
            return "게임 세션이 활성화되지 않았습니다. 새 게임을 시작하세요.", "스탯 없음", []
        
        self.turn_number += 1
        
        prev_turn = interactions.select(
            interactions.player_name,
            interactions.genre,
            interactions.initial_scenario,
            interactions.player_stats
        ).where(
            (interactions.session_id == self.current_session_id) &
            (interactions.turn_number == self.turn_number - 1)
        ).collect()
        
        self.current_stats = prev_turn['player_stats'][0]
        
        # 일반 함수로 변환했으므로 그냥 호출해서 사용
        random_event_val = ""
        if self.turn_number % 3 == 0 and self.turn_number > 0:
            events = [
                "갑자기 부근에서 이상한 소리가 들립니다",
                "낯선 여행자가 당신을 바라보고 있습니다",
                "지면이 미세하게 진동하기 시작합니다",
                "주머니에서 무언가가 빛납니다",
                "멀리서 무언가가 당신을 향해 다가오고 있습니다",
                "갑자기 날씨가 변하기 시작합니다",
                "주변에 숨겨진 통로를 발견합니다"
            ]
            random_event_val = random.choice(events)
        
        if random_event_val:
            action = f"{action} ({random_event_val})"
        
        interactions.insert([{
            'session_id': self.current_session_id,
            'player_name': prev_turn['player_name'][0],
            'genre': prev_turn['genre'][0],
            'initial_scenario': prev_turn['initial_scenario'][0],
            'turn_number': self.turn_number,
            'player_input': action,
            'timestamp': datetime.now(),
            'player_stats': self.current_stats,
            'random_event': random_event_val
        }])
        
        result = interactions.select(
            interactions.story_text,
            interactions.stats_update,
            interactions.options
        ).where(
            (interactions.session_id == self.current_session_id) & 
            (interactions.turn_number == self.turn_number)
        ).collect()
        
        # Update stats for next turn
        self.current_stats = result['stats_update'][0]
        
        return result['story_text'][0], result['stats_update'][0], result['options'][0]

def create_interface():
    game = RPGGame()
    
    # Custom CSS for improved visuals
    custom_css = """
    .container {
        max-width: 1200px;
        margin: 0 auto;
    }
    
    .title-container {
        background: linear-gradient(135deg, #6e48aa 0%, #9c27b0 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin-bottom: 20px;
        text-align: center;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    }
    
    .story-container {
        background: #f8f9fa;
        border-left: 5px solid #9c27b0;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
        font-family: 'Noto Sans KR', sans-serif;
    }
    
    .stats-container {
        background: #e8f5e9;
        border-left: 5px solid #4caf50;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    
    .options-container {
        background: #e3f2fd;
        border-left: 5px solid #2196f3;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-bottom: 20px;
    }
    
    .action-button {
        background: linear-gradient(135deg, #6e48aa 0%, #9c27b0 100%);
        color: white;
        border: none;
        padding: 10px 20px;
        border-radius: 5px;
        cursor: pointer;
        transition: all 0.3s ease;
    }
    
    .action-button:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 10px rgba(0,0,0,0.2);
    }
    
    .history-container {
        background: #fff8e1;
        border-left: 5px solid #ffc107;
        padding: 15px;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        margin-top: 20px;
    }
    """
    
    with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
        gr.HTML(
            """
            <div class="title-container">
                <h1 style="margin-bottom: 0.5em; font-size: 2.5em;">🎲 AI RPG 어드벤처</h1>
                <p style="font-size: 1.2em;">Pixeltable과 OpenAI로 구현된 몰입형 롤플레잉 게임 경험!</p>
            </div>
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Accordion("🎯 이 앱은 무엇인가요?", open=False):
                    gr.HTML(
                        """
                        <div style="padding: 15px;">
                            <h3>AI RPG 어드벤처는 다음 기능을 보여줍니다:</h3>
                            <ul style="list-style-type: none; padding-left: 5px;">
                                <li>🎮 <b>동적 스토리텔링:</b> AI가 생성하는 몰입형 이야기 경험</li>
                                <li>🔄 <b>게임 상태 관리:</b> Pixeltable로 게임 상태와 기록 추적</li>
                                <li>💭 <b>컨텍스트 기반 선택지:</b> 플레이어 행동에 따른 맞춤형 옵션</li>
                                <li>🤖 <b>AI 스토리텔링:</b> 생생한 내러티브 생성</li>
                                <li>📊 <b>캐릭터 상태 추적:</b> 게임 진행에 따른 스탯 변화</li>
                            </ul>
                        </div>
                        """
                    )
                
                with gr.Accordion("🎨 캐릭터 생성", open=True):
                    player_name = gr.Textbox(
                        label="👤 캐릭터 이름",
                        placeholder="당신의 캐릭터 이름을 입력하세요...",
                        container=False
                    )
                    genre = gr.Dropdown(
                        choices=[
                            "🧙‍♂️ Fantasy",
                            "🚀 Sci-Fi",
                            "👻 Horror",
                            "🔍 Mystery",
                            "🌋 Post-Apocalyptic",
                            "🤖 Cyberpunk",
                            "⚙️ Steampunk"
                        ],
                        label="🎭 장르 선택",
                        container=False,
                        value="🧙‍♂️ Fantasy"
                    )
                    scenario = gr.Textbox(
                        label="📖 시작 시나리오",
                        lines=3,
                        placeholder="초기 설정과 상황을 설명하세요...",
                        container=False
                    )
                    start_button = gr.Button("🎮 모험 시작!", variant="primary")
            
            with gr.Column(scale=2):
                story_display = gr.Markdown(
                    label="📜 스토리",
                    value="<div class='story-container'>모험을 시작하려면 캐릭터를 생성하고 '모험 시작!' 버튼을 클릭하세요.</div>",
                    show_label=False
                )
                
                stats_display = gr.Markdown(
                    label="📊 캐릭터 스탯",
                    value="<div class='stats-container'>모험을 시작하면 캐릭터 스탯이 이곳에 표시됩니다.</div>",
                    show_label=False
                )
                
                gr.HTML("<div class='options-container'><h3>🎯 다음 행동 선택</h3></div>")
                
                action_input = gr.Radio(
                    choices=[],
                    label="",
                    interactive=True
                )
                submit_action = gr.Button("⚡ 행동 실행", variant="primary")
        
        with gr.Row():
            with gr.Column():
                gr.HTML("<div class='history-container'><h3>💫 어드벤처 예시</h3></div>")
                gr.Examples(
                    examples=[
                        ["이순신", "🧙‍♂️ Fantasy", "당신은 잊혀진 신비한 숲의 가장자리에서 눈을 뜹니다. 멀리서 성의 첨탑이 보이고, 당신의 머리속에는 왕국을 위협하는 고대 마법에 관한 단서만이 남아있습니다. 갑자기 숲에서 이상한 빛이 보입니다..."],
                        ["김지영", "🚀 Sci-Fi", "우주선 '호라이즌'의 항해사로서, 당신은 미지의 행성 탐사 중 비상 알람에 깨어납니다. 선장과 연락이 두절되었고, 생명 유지 시스템이 점차 실패하고 있습니다. 조용한 선내에서 이상한 발걸음 소리가 들립니다..."],
                        ["일론 머스크", "🤖 Cyberpunk", "2077년 서울, 당신은 뉴럴링크 인더스트리의 CEO입니다. 당신의 최신 뇌-컴퓨터 인터페이스 기술이 사용자들에게 예상치 못한 능력을 부여하기 시작했습니다 - 그들의 집안 식물과 교감할 수 있게 된 것입니다. 주요 투자자 프레젠테이션을 준비하는 동안, AI 비서가 테스트 참가자들이 불가사의한 '식물 혁명'을 조직하고 있다는 보고를 전합니다..."],
                        ["고든 램지", "🌋 Post-Apocalyptic", "당신은 뉴 서울의 마지막 마스터 셰프로, 옛 럭셔리 호텔 폐허에서 지하 레스토랑을 운영하고 있습니다. 당신의 시그니처 요리는 위험한 방사능 구역에서만 자라는 희귀 버섯이 필요합니다. 오늘 밤의 비밀 모임을 준비하던 중, 정찰병이 주변 지역의 경쟁 요리사 갱단에 관한 불길한 소식을 가지고 돌아옵니다..."],
                        ["유정호", "👻 Horror", "당신은 친구의 초대로 삼림 속 외딴 별장에 주말을 보내러 왔습니다. 첫날 밤, 창밖으로 보이는 기이한 빛에 이끌려 숲으로 들어가게 됩니다. 돌아오는 길을 찾으려 하지만, 별장이 보이지 않고 낯선 안개가 점점 짙어집니다. 멀리서 누군가—아니, 무언가가 당신을 부르는 소리가 들립니다..."],
                        ["박지훈", "🔍 Mystery", "신임 형사로서 당신의 첫 사건은 도시 최고의 기술 기업 CEO의 의문의 실종입니다. 그의 사무실에는 혈흔이 없고, 유일한 단서는 책상 위에 놓인 암호화된 메모와 꺼져있는 그의 최첨단 AI 비서뿐입니다. 조사를 시작하자마자, 당신은 CEO가 마지막으로 작업하던 비밀 프로젝트에 관한 이야기를 듣게 됩니다..."],
                        ["이민수", "⚙️ Steampunk", "증기와 기어로 가득한 뉴 조선에서, 당신은 혁신적인 비행선 설계자입니다. 당신의 최신 발명품 시연 중, 정부의 비밀 요원이 접근해 위험에 처한 황실 가족을 위한 비밀 임무를 제안합니다. 지하 단체들이 왕좌를 위협하고 있으며, 당신의 발명품이 왕가의 유일한 희망이라고 합니다..."]
                    ],
                    inputs=[player_name, genre, scenario]
                )
            
            with gr.Column():
                history_df = gr.Dataframe(
                    headers=["📅 턴", "🎯 플레이어 행동", "💬 게임 반응"],
                    label="📚 모험 역사",
                    wrap=True,
                    row_count=5,
                    col_count=(3, "fixed")
                )

        def start_new_game(name, genre_choice, scenario_text):
            if not name or not genre_choice or not scenario_text:
                return (
                    "<div class='story-container'>모든 필드를 입력한 후 시작하세요.</div>",
                    "<div class='stats-container'>스탯 정보 없음</div>",
                    [],
                    []
                )
            
            try:
                _, initial_story, initial_stats, initial_options = game.start_game(name, genre_choice, scenario_text)
                
                history_df = interactions.select(
                    turn=interactions.turn_number,
                    action=interactions.player_input,
                    response=interactions.story_text
                ).where(
                    interactions.session_id == game.current_session_id
                ).order_by(
                    interactions.turn_number
                ).collect().to_pandas()
                
                history_data = [
                    [str(row['turn']), row['action'], row['response']] 
                    for _, row in history_df.iterrows()
                ]
                
                story_html = f"<div class='story-container'>{initial_story}</div>"
                stats_html = f"<div class='stats-container'><h3>📊 캐릭터 상태</h3>{initial_stats}</div>"
                
                return story_html, stats_html, gr.Radio(choices=initial_options, interactive=True), history_data
            except Exception as e:
                return (
                    f"<div class='story-container'>게임 시작 오류: {str(e)}</div>",
                    "<div class='stats-container'>스탯 정보 없음</div>",
                    [],
                    []
                )

        def process_player_action(action_choice):
            try:
                if not action_choice:
                    return (
                        "<div class='story-container'>계속하려면 행동을 선택하세요.</div>",
                        "<div class='stats-container'>스탯 정보 없음</div>",
                        [],
                        []
                    )
                    
                story, stats, options = game.process_action(action_choice)
                
                history_df = interactions.select(
                    turn=interactions.turn_number,
                    action=interactions.player_input,
                    response=interactions.story_text
                ).where(
                    interactions.session_id == game.current_session_id
                ).order_by(
                    interactions.turn_number
                ).collect().to_pandas()
                
                history_data = [
                    [str(row['turn']), row['action'], row['response']] 
                    for _, row in history_df.iterrows()
                ]
                
                story_html = f"<div class='story-container'>{story}</div>"
                stats_html = f"<div class='stats-container'><h3>📊 캐릭터 상태</h3>{stats}</div>"
                
                return story_html, stats_html, gr.Radio(choices=options, interactive=True), history_data
            except Exception as e:
                return (
                    f"<div class='story-container'>오류: {str(e)}</div>",
                    "<div class='stats-container'>스탯 정보 없음</div>",
                    [],
                    []
                )
        
        start_button.click(
            start_new_game,
            inputs=[player_name, genre, scenario],
            outputs=[story_display, stats_display, action_input, history_df]
        )
        
        submit_action.click(
            process_player_action,
            inputs=[action_input],
            outputs=[story_display, stats_display, action_input, history_df]
        )
        
        gr.HTML("""
        <div style="text-align: center; margin-top: 30px; padding: 20px; background: #f5f5f5; border-radius: 10px;">
            <h3>🌟 AI RPG 어드벤처 - Pixeltable로 제작된 몰입형 롤플레잉 경험</h3>
            <p>자신만의 캐릭터를 만들고, 선택한 장르의 세계에서 모험을 즐기세요. 당신의 선택이 스토리를 형성합니다!</p>
        </div>
        """)

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()