import asyncio
import streamlit as st
from dotenv import load_dotenv
import os
import random

from middle_earth_adventure.constants import ALL_NAMES, ALL_SKILLS, ALL_TYPES, TEXT_MODEL, AUDIO_MODEL
from middle_earth_adventure.game_core import GameCore
from middle_earth_adventure.prompts import IMAGE_PROMPT
from middle_earth_adventure.utils import are_all_options_are_filled, check_valid_player, pick_rand_index, pick_rand_items
from middle_earth_adventure.schemas import Player, TechSpecs

################ BACKEND CODE ################
load_dotenv()
key = os.environ.get("OPENAI_PERSONAL_KEY")

# state-variables initialization 
if "text_area_value" not in st.session_state:
    st.session_state.text_area_value = "Choose you character..."
if "player" not in st.session_state:
    st.session_state.player = None
if "tech_specs" not in st.session_state:
    st.session_state.tech_specs = TechSpecs(narrator_voice="nova", image_model="", game_lenght=0)
if "image" not in st.session_state:
    st.session_state.image = "resources/intro.jpg"
if "narrator_audio"not in st.session_state:
    st.session_state.narrator_audio = None
if "game"not in st.session_state:
    st.session_state.game = GameCore(api_key=key, text_model=TEXT_MODEL, tts_model=AUDIO_MODEL)
if "game_iteration" not in st.session_state:
    st.session_state.game_iteration = 0
if "rand" not in st.session_state:
    st.session_state.rand = random.random()

game = st.session_state.game

async def progress_game(text_to_write, selection=None, start=False):
    with st.spinner('Loading...'):
        # utils
        player = st.session_state.player
        tech_specs = st.session_state.tech_specs
        st.session_state.game_iteration += 1 # count game rounds
        check_valid_player(player=st.session_state.player)
        # write text
        st.write(text_to_write)
        # Chat completion
        if start:
            narration_txt = await game.start_adventure(player=player)
        elif st.session_state.game_iteration < tech_specs.game_lenght:
            narration_txt = await game.continue_adventure(player=player, selection=selection)
        elif st.session_state.game_iteration == tech_specs.game_lenght:
            narration_txt = await game.finish_adventure(player=player, selection=selection)
        else: 
            narration_txt = "Game has ended. Thanks for playing!"
        # update
        st.session_state.text_area_value = narration_txt # update

        # Text to Speech
        mp3_audio_bytes = await game.narrate_adventure_out_loud(narration_txt, tech_specs.narrator_voice)
        st.session_state.narrator_audio = mp3_audio_bytes # update
        
        # Text to image generation
        prompt = IMAGE_PROMPT.format(narration=narration_txt, response_format='b64_json',name=name, sex=sex, type=character_type)
        image_url = await game.generate_picture_of_the_adventure(prompt, tech_specs.image_model)
        st.session_state.image = image_url 
        
        # Re-run to update states
        st.rerun()

default_name = ALL_NAMES[pick_rand_index(ALL_NAMES)]
default_type = pick_rand_index(ALL_TYPES)
default_skills = pick_rand_items(ALL_SKILLS, 2)

################ USER INTERFACE (Streamlit) ################

# Title
st.title("Middle Earth Adventures")


# Character and Game Selection
with st.form("selection_form"):
    # Tech Specs
    with st.expander("Technical Specs", expanded=False):
        narrator_voice = st.radio("Narrator's Voice", ["nova", "echo"], index=0)
        image_model = st.radio("Image Model", ['dall-e-2', 'dall-e-3'], index=1)
        game_lenght = st.selectbox("Game Lenght (nr of conversation turns)", [5, 7, 10, 15, 20], index=2)

    # Character Selection
    with st.expander("Character Selection", expanded=True):
        name = st.text_input("Name", value=default_name)
        character_type = st.selectbox("Type", ALL_TYPES, index=default_type)
        sex = st.radio("Gender", ["she", "he"], index=0)
        skills = st.multiselect("Skills (pick 2)", ALL_SKILLS, max_selections=2, help="")
        submit_button = st.form_submit_button("Create Character", use_container_width=True)
    if submit_button:
        # write player
        player = Player(name=name, type=character_type, sex=sex, skills=skills)
        st.session_state.player = player
        # write tech-specs
        tech_specs = TechSpecs(narrator_voice=narrator_voice, image_model=image_model, 
                            game_lenght=game_lenght)
        st.session_state.tech_specs = tech_specs
        # start adventure
        message = f"You are {name}, {sex} is a {character_type}. Your are good at {' and '.join(skills)}"
        if are_all_options_are_filled(player, name, character_type, sex, skills):
            asyncio.run(progress_game(message, start=True))
        else:
            st.toast("Invalid character definition!")

# Story Image
st.image(st.session_state.image, use_column_width=True)

# Story text
st.markdown(f'{st.session_state.text_area_value}')

# Narrator's Audio
st.audio(st.session_state.narrator_audio, format="audio/mp3", start_time=0, loop=False, autoplay=False)

# Action Buttons
col1, col2, col3 = st.columns(3)
with col1:
    if st.button("A", use_container_width=True) and check_valid_player(st.session_state.player):
        asyncio.run(progress_game("You chose option A", selection="A"))
with col2:
    if st.button("B", use_container_width=True) and check_valid_player(st.session_state.player):
        asyncio.run(progress_game("You chose option B", selection="B"))
with col3:
    if st.button("C", use_container_width=True) and check_valid_player(st.session_state.player):
        asyncio.run(progress_game("You chose option C", selection="C"))