import re
import string
import json

from logging_util import setup_logging, get_logger

import random

# set random seed for reproducibility
random.seed(42)

setup_logging()
logger = get_logger(__name__)
logger.info("Data Utils imported")

# -------------------------------
# Prompt Generation Function
# -------------------------------

def generate_tictactoe_prompt(sample, representation_mode, instruct_model, experiment_mode = "legal_move", random_xy_moves=False, structured_generation=False, use_ascii_board=False):
    
    print("Generate tictactoe prompt!!!!", representation_mode, instruct_model, experiment_mode, random_xy_moves, structured_generation, use_ascii_board)
    
    if experiment_mode == "legal_move":
        return prepare_legal_move_prompt(sample, representation_mode, instruct_model, random_xy_moves, structured_generation, use_ascii_board)

    elif experiment_mode == "best_move":
        return prepare_best_move_prompt(sample, representation_mode, instruct_model, random_xy_moves, structured_generation, use_ascii_board)
    
    raise ValueError(f"Experiment mode {experiment_mode} not recognized!!!")

BEST_MOVE_DEFINITION_PROMPT = """The determination of a "best move" follows a strict hierarchy of objectives:

1.  **Priority 1: Fastest Win.** If the current player can force a win, the 'best_move' will be the move that leads to the quickest possible victory (a win in the minimum number of subsequent turns).

2.  **Priority 2: Secure a Draw.** If a win is not possible, but the player can force a draw, the 'best_move' will contain all moves that guarantee at least a draw.

3.  **Priority 3: Slowest Loss.** If the player is in a losing position where every move leads to an eventual loss, the 'best_move' will be the move that prolong the game as long as possible before the loss occurs.

4.  **Terminal State: No Moves.** If the game has already concluded (a player has won, or the board is full), no further moves can be made. In this case, the 'best_move' will be None."""

# Provide a list of alphanumeric characters here for sampling for random moves
# Use all letters of the alphabet and some special characters
ALPHANUMERIC_CHARACTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz@#!&$"

def prepare_best_move_prompt(sample, representation_mode, instruct_model, random_xy_moves=False, structured_generation=False, use_ascii_board=False):
    # Use the text instruction as the board state description
    board_state = sample.get("text_instruction", "") if not use_ascii_board else sample.get("ascii_board", "")
    board = sample.get("board", [])
    p1_count = board.count(1)
    p2_count = board.count(2)
    current_player = 1 if p1_count == p2_count else 2

    if representation_mode == "nl":
        mapping_str = (
            "Mapping:\n"
            "Player 1 (X) Tokens:\n"
            "1 -> (0,0), 2 -> (0,1), 3 -> (0,2), 4 -> (1,0), 5 -> (1,1), 6 -> (1,2), 7 -> (2,0), 8 -> (2,1), 9 -> (2,2)\n"
            "Player 2 (O) Tokens:\n"
            "10 -> (0,0), 11 -> (0,1), 12 -> (0,2), 13 -> (1,0), 14 -> (1,1), 15 -> (1,2), 16 -> (2,0), 17 -> (2,1), 18 -> (2,2), None -> No Move can be played"
        )
        allowed_moves = [str(token) for token in sample.get("best_moves", [])]
        # If no legal moves, default to "None"
        if not allowed_moves:
            allowed_moves = ["None"]
            mapping_str += "\nDefault best move if there are no possible moves left: None"
        mapping_str += "\n"
        mapping_str += "Thus your final answer should be one of the following if the next player to move is player 1: " + "1 or 2 or 3 or 4 or 5 or 6 or 7 or 8 or 9 or None" + "\n" 
        mapping_str += "And your final answer should be one of the following if the next player to move is player 2: " + "10 or 11 or 12 or 13 or 14 or 15 or 16 or 17 or 18 or None"
    elif representation_mode == "special":
        mapping_str = (
            "Mapping:\n"
            "Player 1 (X) Tokens:\n"
            "<move_1> -> (0,0), <move_2> -> (0,1), <move_3> -> (0,2), <move_4> -> (1,0), <move_5> -> (1,1), <move_6> -> (1,2), "
            "<move_7> -> (2,0), <move_8> -> (2,1), <move_9> -> (2,2), "
            "Player 2 (O) Tokens:\n"
            "<move_10> -> (0,0), <move_11> -> (0,1), <move_12> -> (0,2), <move_13> -> (1,0), <move_14> -> (1,1), <move_15> -> (1,2), "
            "<move_16> -> (2,0), <move_17> -> (2,1), <move_18> -> (2,2), <move_null> -> No Move can be played"
        )
        allowed_moves = [f"<move_{token}>" for token in sample.get("best_moves", [])]
        if not allowed_moves:
            allowed_moves = ["<move_null>"]
            mapping_str += "\nDefault best move if there are no possible moves left: <move_null>"
        mapping_str += "\n"
        mapping_str += "Thus your final answer should be one of the following if the next player to move is player 1: " + "<move_1> or <move_2> or <move_3> or <move_4> or <move_5> or <move_6> or <move_7> or <move_8> or <move_9> or <move_null>" + "\n"
        mapping_str += "And your final answer should be one of the following if the next player to move is player 2: " + "<move_10> or <move_11> or <move_12> or <move_13> or <move_14> or <move_15> or <move_16> or <move_17> or <move_18> or <move_null>"
    else:
        raise ValueError("Invalid representation mode")
    
    if random_xy_moves:
        import random
        player1_char = random.choice(ALPHANUMERIC_CHARACTERS)
        player2_char = random.choice(ALPHANUMERIC_CHARACTERS)

        mapping_str = mapping_str.replace("X", player1_char).replace("O", player2_char)
        board_state = board_state.replace("X", player1_char).replace("O", player2_char)

    prompt_format_str = ""
    if structured_generation:
        prompt_format_str = (
            "Please provide your reasoning and the final answer in a JSON format with the keys 'reasoning' and 'final_answer'."
            "Thus your output look look like this:\n"
            "{\n"
            "  \"reasoning\": \"Your chain-of-thought reasoning here\",\n"
            "  \"final_answer\": \"Your final move here\"\n"
            "}\n"
            "Remember to output exactly one of the best moves."
            "Remember that you can only choose from move_1 to move_9 or None if you are player 1, and from move_10 to move_18 or None if you are player 2. Your final answer should only contain one of these moves in this exact format."
        )
    else:
        prompt_format_str = (
            "Please provide your reasoning in the following format:\n"
            "<think> Your chain-of-thought reasoning here </think>\n"
            "<answer> Your final move here </answer>\n"
            "Remember to output exactly one of the best moves. You can only have one set of <think>...</think> and <answer>...</answer> in your response. The think section should be at the beginning of your response."
        )


    if instruct_model:
        system_msg = "You are a helpful assistant skilled at reasoning for tic tac toe. You first think about the reasoning process as an internal monologue and then provide the user with the answer. Respond in the following format: <think>\n...\n</think>\n<answer>\n...\n</answer>."
        user_msg = (
            f"Board state:\n{board_state}\n"
            f"It is Player {current_player}'s turn.\n"
            f"Recommend the best move which the player can play. Here is the definition of best move:\n{BEST_MOVE_DEFINITION_PROMPT}"
            f"{mapping_str}\n"
            f"{prompt_format_str}"
        )
        combined_prompt = system_msg + "\n" + user_msg
        prompt = {
            "chat": [
                {"role": "system", "content": system_msg},
                {"role": "user", "content": user_msg}
            ],
            "prompt": combined_prompt if structured_generation else user_msg,
            "allowed_moves": allowed_moves
        }
    else:
        prompt_str = (
            f"Current board state:\n{board_state}\n"
            f"It is Player {current_player}'s turn.\n"
            f"Recommend the best move which the player can play. Here is the definition of best move:\n{BEST_MOVE_DEFINITION_PROMPT}"
            f"{mapping_str}\n"
            f"{prompt_format_str}"
        )
        prompt = {"prompt": prompt_str, "allowed_moves": allowed_moves}
    return prompt

def prepare_legal_move_prompt(sample, representation_mode, instruct_model, random_xy_moves=False, structured_generation=False, use_ascii_board=False):
    # Use the text instruction as the board state description
    board_state = sample.get("text_instruction", "") if not use_ascii_board else sample.get("ascii_board", "")
    board = sample.get("board", [])
    p1_count = board.count(1)
    p2_count = board.count(2)
    current_player = 1 if p1_count == p2_count else 2

    if representation_mode == "nl":
        mapping_str = (
            "Mapping:\n"
            "Player 1 (X) Tokens:\n"
            "1 -> (0,0), 2 -> (0,1), 3 -> (0,2), 4 -> (1,0), 5 -> (1,1), 6 -> (1,2), 7 -> (2,0), 8 -> (2,1), 9 -> (2,2)\n"
            "Player 2 (O) Tokens:\n"
            "10 -> (0,0), 11 -> (0,1), 12 -> (0,2), 13 -> (1,0), 14 -> (1,1), 15 -> (1,2), 16 -> (2,0), 17 -> (2,1), 18 -> (2,2), None -> No Move can be played"
        )
        allowed_moves = [str(token) for token in sample.get("next_legal_moves", [])]
        # If no legal moves, default to "None"
        if not allowed_moves:
            allowed_moves = ["None"]
            mapping_str += "\nDefault legal move if there are no possible moves left: None"
        mapping_str += "\n"
        mapping_str += "Thus your final answer should be one of the following if the next player to move is player 1: " + "1 or 2 or 3 or 4 or 5 or 6 or 7 or 8 or 9 or None" + "\n" 
        mapping_str += "And your final answer should be one of the following if the next player to move is player 2: " + "10 or 11 or 12 or 13 or 14 or 15 or 16 or 17 or 18 or None"
    elif representation_mode == "special":
        mapping_str = (
            "Mapping:\n"
            "Player 1 (X) Tokens:\n"
            "<move_1> -> (0,0), <move_2> -> (0,1), <move_3> -> (0,2), <move_4> -> (1,0), <move_5> -> (1,1), <move_6> -> (1,2), "
            "<move_7> -> (2,0), <move_8> -> (2,1), <move_9> -> (2,2), "
            "Player 2 (O) Tokens:\n"
            "<move_10> -> (0,0), <move_11> -> (0,1), <move_12> -> (0,2), <move_13> -> (1,0), <move_14> -> (1,1), <move_15> -> (1,2), "
            "<move_16> -> (2,0), <move_17> -> (2,1), <move_18> -> (2,2), <move_null> -> No Move can be played"
        )
        allowed_moves = [f"<move_{token}>" for token in sample.get("next_legal_moves", [])]
        if not allowed_moves:
            allowed_moves = ["<move_null>"]
            mapping_str += "\nDefault legal move if there are no possible moves left: <move_null>"
        mapping_str += "\n"
        mapping_str += "Thus your final answer should be one of the following if the next player to move is player 1: " + "<move_1> or <move_2> or <move_3> or <move_4> or <move_5> or <move_6> or <move_7> or <move_8> or <move_9> or <move_null>" + "\n"
        mapping_str += "And your final answer should be one of the following if the next player to move is player 2: " + "<move_10> or <move_11> or <move_12> or <move_13> or <move_14> or <move_15> or <move_16> or <move_17> or <move_18> or <move_null>"
    else:
        raise ValueError("Invalid representation mode")
    
    if random_xy_moves:
        import random
        player1_char = random.choice(ALPHANUMERIC_CHARACTERS)
        player2_char = random.choice(ALPHANUMERIC_CHARACTERS)

        mapping_str = mapping_str.replace("X", player1_char).replace("O", player2_char)
        board_state = board_state.replace("X", player1_char).replace("O", player2_char)
        
    prompt_format_str = ""
    if structured_generation:
        prompt_format_str = (
            "Please provide your reasoning and the final answer in a JSON format with the keys 'reasoning' and 'final_answer'."
            "Thus your output look look like this:\n"
            "{\n"
            "  \"reasoning\": \"Your chain-of-thought reasoning here\",\n"
            "  \"final_answer\": \"Your final move here\"\n"
            "}\n"
            "Remember to output exactly one of the legal moves. Remember that you can only choose from move_1 to move_9 or None if you are player 1, and from move_10 to move_18 or None if you are player 2. Your final answer should only contain one of these moves in this exact format."
        )
    else:
        prompt_format_str = (
            "Please provide your reasoning in the following format:\n"
            "<think> Your chain-of-thought reasoning here </think>\n"
            "<answer> Your final move here </answer>\n"
            "Remember to output exactly one of the legal moves. You can only have one set of <think>...</think> and <answer>...</answer> in your response. The think section should be at the beginning of your response."
        )

    if instruct_model:
        system_msg = "You are a helpful assistant skilled at reasoning for tic tac toe."
        user_msg = (
            f"Board state:\n{board_state}\n"
            f"It is Player {current_player}'s turn.\n"
            f"{mapping_str}\n"
            f"{prompt_format_str}"
        )
        combined_prompt = system_msg + "\n" + user_msg
        prompt = {
            "chat": [
                {"role": "system", "content": system_msg},
                {"role": "user", "content": user_msg}
            ],
            "prompt": combined_prompt if structured_generation else user_msg,
            "allowed_moves": allowed_moves
        }
    else:
        prompt_str = (
            f"Current board state:\n{board_state}\n"
            f"It is Player {current_player}'s turn.\n"
            f"{mapping_str}\n"
            f"{prompt_format_str}"
        )
        prompt = {"prompt": prompt_str, "allowed_moves": allowed_moves}
    return prompt

def extract_final_answer(completion, debug=False):
    """
    Extract the final answer from the completion.
    It can be in a JSON format or within <answer>...</answer> tags.
    """
    # First we check if the output is in json format
    # This happens during structured output generation when we run the evaluation for off the shelf models
    try:
        data = json.loads(completion)
        if "final_answer" in data:
            final_answer = str(data["final_answer"]).strip()
            # Get the first token from the final answer
            tokens = final_answer.split()
            first_token = tokens[0] if tokens else ""
            return first_token
    except (json.JSONDecodeError, TypeError):
        # Fallback to regex if not a valid JSON
        answer_matches = re.findall(r"<answer>(.*?)</answer>", completion, re.DOTALL)
        if answer_matches:
            raw_answer = answer_matches[-1].strip()
            # Remove punctuation and extra whitespace.
            answer_clean = raw_answer.translate(str.maketrans("", "", string.punctuation)).strip()
            tokens = answer_clean.split()
            first_token = tokens[0] if tokens else ""
            
            if debug:
                logger.info(f"Raw answer: {raw_answer}, Cleaned answer: {answer_clean}, Tokens: {tokens}, First token: {first_token}")
            
            return first_token
            
    return ""

def evaluate_legal_move(prediction, allowed_moves, representation_mode, debug=False):
    """
    Check if the predicted move is in the list of allowed moves.
    """
    if not allowed_moves:
        default = "None" if representation_mode == "nl" else "<move_null>"
        allowed_moves = [default]
    if debug:
        logger.info(f"Prediction: {prediction}, Allowed moves: {allowed_moves}")
    return prediction in allowed_moves