from twentyq.game import Q20Game, Q20GameCelebrity
import os
from twentyq.utils import load_model
from twentyq.eval_result import compute_metrics
from dotenv import load_dotenv

# Load environment variables from .env file if exists
load_dotenv()

# Function to play a game of 20 Questions
def play_game(
    name,
    guesser_model_name="gpt-4o-mini",
    num_turns=20,
    temperature=0.8,
    language="english",
    base_output_dir="./outputs",
    game_type="celebrities",  # or "things"
):
    
    """
    Plays a game of 20 Questions with the specified parameters.
    Args:
        name (str): The name of the entity to guess.
        guesser_model_name (str): The name of the model to use for guessing and answering.
        num_turns (int): The maximum number of turns allowed for the guesser.
        temperature (float): The temperature for the guesser's model's responses.
        language (str): The language in which the game is played.
        base_output_dir (str): The base directory where output files will be saved.
        game_type (str): The type of game to play, either "things" or "celebrities".
    Returns:
        dict: A dictionary containing the game play data and metrics.
    """

    # Load model and tokenizer
    if "/" in guesser_model_name:
        guesser_model, guesser_tokenizer = load_model(guesser_model_name)
        guesser_kargs = {
            "max_new_tokens": 64,
            "temperature": temperature,
            "repetition_penalty": 1.0,
            "do_sample": True,
        }
        cleaned_guesser_model = guesser_model_name.split("/")[1]
    else:
        guesser_model, guesser_tokenizer = guesser_model_name, None
        guesser_kargs = {}
        cleaned_guesser_model = guesser_model_name

    # Select correct game class based on game_type
    if game_type == "things":
        game = Q20Game(
            item=name,
            answerer_model=guesser_model,
            guesser_model=guesser_model,
            guesser_tokenizer=guesser_tokenizer,
            num_turns=num_turns,
            temperature=temperature,
            openai_api=True,
            guesser_kargs=guesser_kargs,
            language=language,
            game_type="things"
        )
    elif game_type == "celebrities":
        game = Q20GameCelebrity(
            item=name,
            background="",
            answerer_model=guesser_model,
            guesser_model=guesser_model,
            guesser_tokenizer=guesser_tokenizer,
            num_turns=num_turns,
            temperature=temperature,
            openai_api=True,
            guesser_kargs=guesser_kargs,
            language=language,
            game_type="celebrities"
        )
    else:
        raise ValueError(f"Unsupported game_type: {game_type}")

    # Output directory
    output_dir = os.path.join(
        base_output_dir,
        f"{cleaned_guesser_model}_{num_turns}_turns_{language}"
    )
    os.makedirs(output_dir, exist_ok=True)

    # Play and save
    game.game_play(path=output_dir, user_mode=None)
    game_play_data = game.save_session(path=output_dir)

    # Compute metrics
    game_play_metrics = compute_metrics(
        os.path.join(output_dir, f"{name}.txt"),
        num_turns,
        language
    )

    # Merge into the nested structure
    game_play_data[name]["metrics"] = game_play_metrics

    # Print the game play data
    print(f"Game play data for {name}:")
    for key, value in game_play_data[name].items():
        if key != "metrics":
            print(f"{key}: {value}")
        else:
            print(f"Metrics: {value}")

    print(f"\nGame completed for: {name}\nOutput saved in: {output_dir}")
    return game_play_data
