import os
import tkinter as tk
import os
import random
import time
import json
from tqdm import tqdm

from modules.wires import SimpleWirePuzzle
from modules.button import ButtonPuzzle
from modules.keypad import KeyPadPuzzle
from modules.password import PasswordPuzzle
from modules.memory import MemoryPuzzle
from modules.sound import SoundPuzzle
from modules.maze import MazePuzzle
from modules.led import LedPuzzle
from modules.colour import ColourPuzzle
from modules.dog import DogPuzzle
from modules.who import WhoPuzzle
from modules.test import ExamplePuzzle

from agents.chatgpt_agent import ChatGPTAgent
from agents.gpt4v_agent import GPT4VAgent
from agents.gpt4o_agent import GPT4oAgent

# from agents.internVL import InternVLAgent
# from agents.qwenVL2 import QwenVLAgent
# from agents.llava import LLAVAAgent
from agents.humanAgent import HumanAgent
# from agents.random_agent import RandomAgent


from tkinter import font

import argparse

from agents.loaders import load_class_from_file
from agents.loaders import load_classes_from_config, load_one_agent_from_config


class GameManager:
    def __init__(self, master, args):

        self.MAX_STEPS = 20
        self.cur_step = 0
        self.use_sound = args.use_sound
        self.gui = args.gui
        self.save_folder = args.save_folder
        self.socket = None
        self.record_all = ""
        self.master = master
        self.master.title("Bomb")
        

        self.config = None
        if args.config:
            self.config = json.load(open(args.config))
        
        self.model_config = None
        if args.model_config:
            self.model_config = json.load(open(args.model_config))

        # Prevent resizing of the window
        self.master.resizable(False, False)
        # self.master.geometry("+0+0")
        # self.master.overrideredirect(True)
        self.master.attributes("-topmost", True)
        
        self.total_seconds = 6000

        # Create canvas
        self.width = 500
        self.height = 500

        chatbox_width = 76
        chatbox_height = 22

        self.canvas = tk.Canvas(
            self.master, width=self.width + 800, height=self.height + 80, bg="white")

        self.chat_display = tk.Text(
            self.master, height=chatbox_height, width=chatbox_width, state=tk.DISABLED, bg="black")

        # List of Puzzles in order they appear
        self.puzzles = []
        if self.config:
            self.serial_number = 123456 # TODO, make this part of puzzle config
            for entry in self.config:
                self.puzzles.append(globals()[list(entry.keys())[0]])
        else:
            self.serial_number = random.randint(100000, 999999)
            self.puzzles = [SimpleWirePuzzle, ExamplePuzzle]

        self.canvas.focus_set()
        self.draw()

        self.chat_display.place(x=self.width + 12, y= 12 + 25)
        self.chat_display.tag_config("message", foreground="#33ff00")
        self.chat_display.tag_config("solver_title", foreground="#3d80d9")
        self.chat_display.tag_config("expert_title", foreground="#d93d5a")

        self.wire_colors = []
        self.update_timer()

        if self.model_config:
            ExpertClass, SolverClass, self.model_config = load_classes_from_config(self.model_config)
            self.expert = ExpertClass(role="EXPERT")
            self.solver = SolverClass(role="SOLVER")
        else:
            self.expert = HumanAgent(role="EXPERT")
            self.solver = HumanAgent(role="SOLVER")

        self.p_bar = tqdm(range(len(self.puzzles)))

        self.num_puzzles = len(self.puzzles)
        self.cur_puzzle = 0
        self.puzzle = self.puzzles[self.cur_puzzle]
        os.makedirs(self.save_folder, exist_ok=True)
        
        self.setup_puzzle()

        # Create the user input entry
        self.entry = tk.Entry(self.master, width=80, bg="#AEAEAE", foreground="black")
        self.entry.place(x= self.width + 12 + 80, y = 500)
        self.entry.bind("<Return>", self.send_message_gui)
        self.entry.bind("<Escape>", self.on_escape)

        # If no GUI, run the loop between agents
        if not args.gui:
            self.send_message()

    # Initializes widgets in new puzzle and stores their location
    def setup_puzzle(self):
        self.puzzle = self.puzzles[self.cur_puzzle](self)
        self.save_dir = os.path.join(self.save_folder, f"{self.solver.__class__.__name__}_solver_{self.expert.__class__.__name__}_expert", self.puzzle.__class__.__name__)   
         
        run_num = 0
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir, exist_ok=True)
        elif len(os.listdir(self.save_dir)) > 0:
            runs = os.listdir(self.save_dir)
            run_numbers = [int(x.split("_")[-1]) for x in runs]
            run_num = 1 + max(run_numbers)

        self.save_dir = os.path.join(self.save_dir, f"run_{run_num}")
            
        self.canvas.update()
        if self.puzzle.buttons is not None:
            for button in self.puzzle.buttons:
                button.x_coord = button.winfo_rootx() + button.winfo_reqwidth() / 2
                button.y_coord = button.winfo_rooty() + button.winfo_reqheight() / 2
        self.puzzle.send_info_to_solver()
        self.expert.set_module(self.puzzle)
        self.solver.set_module(self.puzzle)

    # Utility function if user runs with GUI
    def display_solver_message(self, message):
        if self.gui:
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.insert(tk.END, "SOLVER: ", "solver_title")
            self.chat_display.insert(tk.END, message + "\n", "message")
            self.chat_display.config(state=tk.DISABLED)
            self.write_to_conversation({"from": "SOLVER", "value": message})

    # Utility function if user runs with GUI
    def display_expert_message(self, message):
        if self.gui:
            self.chat_display.config(state=tk.NORMAL)
            self.chat_display.insert(tk.END, "EXPERT: ", "expert_title")
            self.chat_display.insert(tk.END, message + "\n", "message")
            self.chat_display.config(state=tk.DISABLED)
            self.write_to_conversation({"from": "EXPERT", "value": message})

    # Save the conversation between agents to the outputs folder
    def write_to_conversation(self, message):
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir, exist_ok=True)
        with open(os.path.join(self.save_dir, "conversation.jsonl"), "a") as f:
            f.write(json.dumps(message) + "\n")

    # Main loop between solver and expert
    def send_message(self, event=None):

        self.solver.save_dir = self.save_dir
        self.expert.save_dir = self.save_dir
            
        message = "Hi! I'm the solver!"
        self.display_solver_message(message)
        self.write_to_conversation({"from": "SOLVER", "value": message})
        self.cur_step = 0

        while message and self.cur_step <= self.MAX_STEPS:
            self.cur_step += 1
            if self.cur_step > self.MAX_STEPS:
                print(
                    "ENVIRONMENT: Solver has reached the maximum number of steps. Moving on to the next puzzle.")
                with open(os.path.join(self.save_dir, "conversation.jsonl"), "a") as f:
                    f.write(json.dumps(
                        {"from": "ENVIRONMENT", "value": "Reached maximum number of steps without solution ..."}) + "\n")
                self.advance_puzzle(success=False)
                break

            history_puzzle = self.cur_puzzle
            old_mistakes = self.puzzle.mistakes
            
            if self.solver.step(self.puzzle, message):
                self.master.update()
                new_mistakes = self.puzzle.mistakes

                message = self.solver.get_feedback(
                    self.cur_puzzle, history_puzzle, old_mistakes < new_mistakes)

                if self.cur_puzzle != history_puzzle:
                    break

                self.display_solver_message(message)
                self.write_to_conversation(
                    {"from": "ENVIRONMENT", "value": message})
            

            ai_response = self.expert.respond(None, None, message)
            self.display_expert_message(ai_response)
            self.write_to_conversation(
                {"from": "EXPERT", "value": ai_response})
            
            message = self.solver.respond_with_image(
                self.width, self.height, self.puzzle.actions, ai_response)
            self.display_solver_message(message)
            self.write_to_conversation({"from": "SOLVER", "value": message})

            self.canvas.focus_set()

    def send_message_gui(self, event=None):
        message = self.entry.get()
        if message:
            self.display_solver_message(message)
            ai_response = self.expert.respond(None, None, message)
            self.display_expert_message(ai_response)
            self.entry.delete(0, tk.END)
            self.canvas.focus_set()


    # Draws the boundary of the module as well as info about it such as the
    # serial number and timer
    def draw(self):
        self.canvas.create_rectangle(
            12, 12, self.width - 12, self.height - 12, outline="black", width=2, fill="gray")
        self.canvas.create_text(
            250, 550, text=f"Serial Number: {self.serial_number}", fill="black", font=("Helvetica", 20))
        self.canvas.create_text(
            920, 22, text=f"The Solver and Expert chat here", fill="black", font=("Helvetica", 20))
        self.canvas.create_oval(self.width - 100, 100,
                                self.width - 50, 50, outline="black", width=2)
        self.timer_label = tk.Label(
            self.canvas, text="", font=("Arial", 20), bg="white")
        self.timer_label.place(x=250, y=512, anchor="center")


    # Advance to the next puzzle in the list. If success==True, it means the
    # agent solved the puzzle. Otherwise, puzzle is advancing due to the max
    # number of allowed turns being exceeded.
    def advance_puzzle(self, success=True):        
        if success:
            if not os.path.exists(os.path.join(self.save_dir, "conversation.jsonl")):
                os.makedirs(self.save_dir, exist_ok=True)
            with open(os.path.join(self.save_dir, "conversation.jsonl"), "a") as f:
                f.write(json.dumps(
                    {"from": "ENVIRONMENT", "value": "Puzzle successfully finished, moving on to the next puzzle ..."}) + "\n")
                
        
        score = self.puzzle.get_score()
        self.write_to_conversation(
            {"from": "ENVIRONMENT", "score": f"{score}"})

        self.cur_puzzle += 1

        self.cur_step = 0

        # Clear conversation history
        self.expert.clear()
        self.solver.clear()

        # Clear Chatbox
        self.chat_display.config(state=tk.NORMAL)
        self.chat_display.delete(1.0, tk.END)
        self.canvas.update()

        # Clear all buttons
        for widget in self.master.winfo_children():
            if isinstance(widget, tk.Button):
                widget.destroy()

        if hasattr(self.puzzle, 'stage_labels'):
            for label in self.puzzle.stage_labels:
                label.destroy()

        time.sleep(2)
        self.canvas.delete("all")

        self.draw()
        self.canvas.unbind("<Key>")
        self.canvas.unbind("<ButtonPress-1>")
        self.canvas.unbind("<ButtonRelease-1>")

        if self.cur_puzzle >= len(self.puzzles):
            exit()

        self.p_bar.n = self.cur_puzzle
        self.p_bar.refresh()

        self.setup_puzzle()

        if not self.gui:
            self.send_message()

    # Update the timer text every second
    def update_timer(self):
        if self.total_seconds > 0:
            minutes, seconds = divmod(self.total_seconds, 60)
            self.timer_label.config(
                text=f"Time Left: {minutes:02d}:{seconds:02d}")
            self.total_seconds -= 1
            self.canvas.after(1000, self.update_timer)
        else:
            self.timer_label.config(text="Time's up!")


    def on_escape(self, event):
        self.canvas.focus_set()


def main(args):
    root = tk.Tk()
    app = GameManager(root, args)
    root.overrideredirect(False)
    root.mainloop()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_sound", action="store_true", default=False)
    parser.add_argument("--config", type=str, default=None)
    parser.add_argument("--model_config", type=str, default=None)
    parser.add_argument("--save_folder", type=str, default="./outputs")
    parser.add_argument("--baseline", type=str, default="none")
    parser.add_argument("--gui", type=bool, default=False,
                        help="flag for if the mouse should move to perform actions")
    args = parser.parse_args()
    main(args)