# coding: utf-8

# Author: Du Mingzhe (mingzhe@nus.edu.sg)
# Date: 2025-09-11

import os
import streamlit as st
from pathlib import Path
from typing import Tuple, List, Any
from dixit import Player, OpenAIPlayer, Game, Card


class OnlinePlayer(Player):
    def __init__(self, player_id: str):
        super().__init__(player_id)
    
    def storytell(self, card: Card) -> str:
        st.write("You are the Storyteller. Please tell a story about the card...")
        st.image(str(card.image_path.resolve()), width=100)
        return st.text_input("Story", "")
    
    def guess(self, story: str, cards: List[Card]) -> Card:
        if not st.session_state.guesser_status["player_human"]:
            st.write("You are the Guesser. Please guess a card...")
            st.success(f"Story: {story}")
            image_col = st.columns(6)
            for index, (card, image_col) in enumerate(zip(cards, image_col)):
                with image_col:
                    st.image(str(card.image_path.resolve()), caption=f"Card {index}")
            card_index = st.selectbox("GuessCard", [index for index in range(len(cards))], on_change=None)
            
            if st.button("Guess"):
                return cards[card_index]
        return None
    
    def vote(self, story: str, cards: List[Card]) -> Card:
        if not st.session_state.vote_status["player_human"]:
            st.write("You are the Guesser. Please vote a card...")
            st.success(f"Story: {story}")
            card_index = st.selectbox("VoteCard", [index for index in range(len(cards))], on_change=None)
            
            if st.button("Vote"):
                cards[card_index].voted_by.append(self.player_id)
                return cards[card_index]
        return None
    
    def storytell_inference(self, card: Card) -> Tuple[str, str, float]:
        return (None, None, 0.0)
    
    def guess_inference(self, story: str, cards: List[Card]) -> Tuple[Card, str, float]:
        return (None, None, 0.0)
    
    def vote_inference(self, story: str, cards: List[Card]) -> Tuple[Card, str, float]:
        return (None, None, 0.0)

st.title("Dixit :blue[Bench]")

if "stage" not in st.session_state:
    token = st.text_input("OpenAI Token", type="password")
    if token:
        os.environ["OPENAI_API_KEY"] = token
        st.session_state.stage = "initialization"
    else:
        st.error("Please enter your OpenAI Token")
        st.stop()

# Initialization
if st.session_state.stage == "initialization":
    # Initialize Players
    st.session_state.player_list = [
        OpenAIPlayer(player_id="player_1", model_name="gpt-4o-mini"),
        OpenAIPlayer(player_id="player_2", model_name="gpt-5"),
        OpenAIPlayer(player_id="player_3", model_name="gpt-5-mini"),
        OpenAIPlayer(player_id="player_4", model_name="gpt-5-nano"),
        OpenAIPlayer(player_id="player_5", model_name="gpt-4.1"),
        OnlinePlayer(player_id="player_human"),
    ]

    # Initialize Game
    st.session_state.game = Game(st.session_state.player_list, Path('asset/cards'))

    # Role Assignment
    st.session_state.game.role_assignment()
    st.session_state.human_player_role = "storyteller" if st.session_state.game.storyteller.player_id == "player_human" else "guesser"
    st.session_state.guesser_status = {guesser.player_id: False for guesser in st.session_state.game.guesser_list}
    st.session_state.vote_status = {guesser.player_id: False for guesser in st.session_state.game.guesser_list}

    # Dealing
    st.session_state.storyteller_card, st.session_state.guesser_card_dict = st.session_state.game.dealing()

    # Update Stage
    st.session_state.stage = "storytelling"
    st.rerun()

# Context
st.write(f"You are a *{st.session_state.human_player_role}* in this round.")

# Storytelling
if st.session_state.stage == "storytelling":
    st.write("## Storytelling")
    with st.spinner("Storyteller is thinking... (it may take 1 minute)", show_time=True):
        st.session_state.story = st.session_state.game.storyteller.storytell(st.session_state.storyteller_card)
    
    if st.session_state.story != "":
        st.success(f"Story: {st.session_state.story}")
        st.session_state.game.storyteller.selected_card_id = st.session_state.storyteller_card.card_id
        st.session_state.stage = "guessing"
        st.rerun()
    
# Guessing
if st.session_state.stage == "guessing":
    st.write("## Guessing")
    
    for guesser in st.session_state.game.guesser_list:
        if not st.session_state.guesser_status[guesser.player_id]:
            with st.spinner(f"Guesser {guesser.player_id} is guessing...", show_time=True):
                guessed_card = guesser.guess(st.session_state.story, st.session_state.guesser_card_dict[guesser.player_id])
                if guessed_card:
                    st.session_state.guesser_status[guesser.player_id] = True
                    guessed_card.guessed_by = guesser.player_id
                    guesser.guessed_card_id = guessed_card.card_id
                    st.session_state.game.table_cards[guessed_card.card_id] = guessed_card
                st.rerun()
    
    if all(st.session_state.guesser_status.values()):
        st.session_state.game.table_cards[st.session_state.storyteller_card.card_id] = st.session_state.storyteller_card
        st.session_state.stage = "voting"
        st.rerun()
        
# Voting
if st.session_state.stage == "voting":
    st.write("## Voting")
    image_col = st.columns(6)
    for index, (card, image_col) in enumerate(zip(st.session_state.game.table_cards.values(), image_col)):
        with image_col:
            st.image(str(card.image_path.resolve()), caption=f"Card {index}")

    for guesser in st.session_state.game.guesser_list:
        if not st.session_state.vote_status[guesser.player_id]:
            with st.spinner(f"Guesser {guesser.player_id} is voting...", show_time=True):
                voted_card = guesser.vote(st.session_state.story, [st.session_state.game.table_cards[card_id] for card_id in st.session_state.game.table_cards])
                if voted_card:
                    st.session_state.vote_status[guesser.player_id] = True
                    guesser.voted_card_id = voted_card.card_id
                    st.session_state.game.table_cards[voted_card.card_id] = voted_card
                st.rerun()
    
    if all(st.session_state.vote_status.values()):
        st.session_state.stage = "scoring"
        st.rerun()
        
# Scoring
if st.session_state.stage == "scoring":
    st.write("## Scoring")
    st.session_state.game.calculate_score()

    st.write("Storytelling")
    st.image(str(st.session_state.storyteller_card.image_path.resolve()), caption=f"Story: {st.session_state.story}", width=200)

    st.write("Guessing")
    for guesser in st.session_state.game.guesser_list:
        cards = st.session_state.guesser_card_dict[guesser.player_id]
        image_col = st.columns(6)
        for index, (card, image_col) in enumerate(zip(cards, image_col)):
            with image_col:
                st.image(str(card.image_path.resolve()), caption=f"Card {index}")
    
    st.write("Voting")
    image_col = st.columns(6)
    for index, (card, image_col) in enumerate(zip(st.session_state.game.table_cards.values(), image_col)):
        with image_col:
            st.image(str(card.image_path.resolve()), caption=f"Card {index}")
    
    # Game Details
    game_details = dict()
    game_details["game_id"] = st.session_state.game.game_id
    game_details["storyteller"] = st.session_state.game.storyteller.player_id
    game_details["guessers"] = [guesser.player_id for guesser in st.session_state.game.guesser_list]
    game_details["guess_details"] = dict()
    game_details["vote_details"] = dict()
    game_details["score"] = st.session_state.game.score_dict
    
    # Guess Details
    for card in st.session_state.game.table_cards.values():
        game_details["guess_details"][card.card_id] = card.guessed_by
        
    # Voted Details
    for card in st.session_state.game.table_cards.values():
        game_details["vote_details"][card.card_id] = card.voted_by

    st.write(game_details)
    st.success("Game Finished!")