import tkinter as tk
from tkinter import filedialog, messagebox, simpledialog
from mutagen.mp3 import MP3, EasyMP3
import base64
import tempfile
import json
import os
import pygame

from mutagen.id3 import ID3, APIC


class LevelDesigner:
    def __init__(self, root):
        self.root = root
        self.root.title("Level Designer")
        initial_width = 600
        initial_height = 400

        self.root.geometry(
            f"{initial_width}x{initial_height}"
        )  # Set initial window size
        self.root.minsize(width=initial_width, height=initial_height)
        self.song_path = ""
        self.notes = []

        self.select_button = tk.Button(
            root, text="Select Song", command=self.select_song
        )
        self.select_button.pack()

        self.canvas_frame = tk.Frame(root)
        self.canvas_frame.pack(fill="both", expand=True)

        self.canvas = tk.Canvas(self.canvas_frame, bg="white")
        self.canvas.pack(side="top", fill="both", expand=True)

        self.h_scrollbar = tk.Scrollbar(
            self.canvas_frame, orient="horizontal", command=self.canvas.xview
        )
        self.h_scrollbar.pack(side="bottom", fill="x")
        self.canvas.configure(xscrollcommand=self.h_scrollbar.set)

        self.canvas.bind("<Button-1>", self.add_note)
        self.canvas.bind("<Button-3>", self.delete_note)

        self.root.bind("<space>", self.auto_add_note)

        self.save_button_visible = False
        self.save_button_height = 0

        self.arrows_visible = False

        self.save_button = tk.Button(root, text="Save Level", command=self.save_level)
        self.save_button.pack()
        self.save_button_height = self.save_button.winfo_height()
        self.save_button.pack_forget()

        self.arrow_buttons_frame = tk.Frame(root)
        self.arrow_buttons_frame.pack()

        self.song = None
        self.meta = None

        self.left_button = tk.Button(
            self.arrow_buttons_frame,
            text="←",
            command=lambda: self.select_arrow("left"),
            bg="red",
        )
        self.left_button.pack(side="left")

        self.up_button = tk.Button(
            self.arrow_buttons_frame,
            text="↑",
            command=lambda: self.select_arrow("up"),
            bg="green",
        )
        self.up_button.pack(side="left")

        self.down_button = tk.Button(
            self.arrow_buttons_frame,
            text="↓",
            command=lambda: self.select_arrow("down"),
            bg="blue",
        )
        self.down_button.pack(side="left")

        self.right_button = tk.Button(
            self.arrow_buttons_frame,
            text="→",
            command=lambda: self.select_arrow("right"),
            bg="yellow",
        )
        self.right_button.pack(side="left")

        self.canvas_frame.bind("<Configure>", self.update_canvas_height)
        self.canvas.bind("<Configure>", self.redraw_notes)

        self.add_controls()
        self.current_pos: float = 0.0

        pygame.mixer.init()
        self.song_playing = False

    def init_song(self):
        if self.song_path:
            pygame.mixer.music.load(self.song_path)
            self.song_playing = False
            self.cursor = None

    def update_cursor(self):
        self.canvas.delete("cursor")
        if self.song_playing:
            current_time = self.current_pos + pygame.mixer.music.get_pos() / 1000.0
        else:
            current_time = self.current_pos

        x = current_time / self.song_length * self.canvas_width

        canvas_width = self.canvas.winfo_width()  # Get current canvas width
        visible_x1 = self.canvas.canvasx(0)  # Get the leftmost visible x coordinate
        visible_x2 = visible_x1 + canvas_width  # Get the rightmost visible x coordinate

        # # Calculate the new scroll region based on the cursor's position
        if x < visible_x1 + 50:  # If cursor is about to leave from the left
            new_x1 = max(0, x - canvas_width / 2)
            self.canvas.xview_moveto(new_x1 / self.canvas_width)
        elif x > visible_x2 - 50:  # If cursor is about to leave from the right
            new_x1 = max(0, x - canvas_width / 2)
            self.canvas.xview_moveto(new_x1 / self.canvas_width)

        # Draw the cursor
        self.cursor = self.canvas.create_line(
            x, 0, x, self.canvas.winfo_height(), fill="red", tags="cursor"
        )

        if self.song_playing:
            self.root.after(100, self.update_cursor)

    def play_pause_song(self):
        if self.song_playing:
            print("Paused")
            pygame.mixer.music.pause()
            self.current_pos += pygame.mixer.music.get_pos() / 1000.0
            self.song_playing = False
        else:
            print("Unpaused")
            pygame.mixer.music.play(start=self.current_pos)
            self.song_playing = True
        self.update_cursor()

    def move_cursor(self, amount):
        if self.song_playing:
            self.current_pos += pygame.mixer.music.get_pos() / 1000.0
            self.current_pos = max(0, min(self.current_pos + amount, self.song_length))
            pygame.mixer.music.play(start=self.current_pos)
        else:
            self.current_pos = max(0, min(self.current_pos + amount, self.song_length))
        self.update_cursor()

    def backward_1_second(self):
        self.move_cursor(-1)

    def backward_half_second(self):
        self.move_cursor(-0.5)

    def forward_1_second(self):
        self.move_cursor(1)

    def forward_half_second(self):
        self.move_cursor(0.5)

    def backward_5_seconds(self):
        self.move_cursor(-5)

    def forward_5_seconds(self):
        self.move_cursor(5)

    def add_controls(self):
        control_frame = tk.Frame(self.root)
        control_frame.pack()

        backward_half_second_button = tk.Button(
            control_frame, text="<< 0.5s", command=self.backward_half_second
        )
        backward_half_second_button.pack(side="left")

        backward_1_second_button = tk.Button(
            control_frame, text="<< 1s", command=self.backward_1_second
        )
        backward_1_second_button.pack(side="left")

        backward_button = tk.Button(
            control_frame, text="<< 5s", command=self.backward_5_seconds
        )
        backward_button.pack(side="left")

        play_pause_button = tk.Button(
            control_frame, text="PP", command=self.play_pause_song
        )
        play_pause_button.pack(side="left")

        forward_button = tk.Button(
            control_frame, text=">> 5s", command=self.forward_5_seconds
        )
        forward_button.pack(side="left")

        forward_1_second_button = tk.Button(
            control_frame, text=">> 1s", command=self.forward_1_second
        )
        forward_1_second_button.pack(side="left")

        forward_half_second_button = tk.Button(
            control_frame, text=">> 0.5s", command=self.forward_half_second
        )
        forward_half_second_button.pack(side="left")

    def get_canvas_width(self):
        canvas_width = self.canvas.winfo_width()
        return canvas_width

    def redraw_notes(self, event):
        self.canvas.delete("notes")
        canvas_height = (
            self.canvas_frame.winfo_height() - self.arrow_buttons_frame.winfo_height()
        )
        self.canvas_height = canvas_height  # Store canvas height for future use
        for note_time, arrow in self.notes:
            x = (
                note_time
                / float(MP3(self.song_path).info.length)
                * self.get_canvas_width()
            )
            color = self.get_arrow_color(arrow)
            self.canvas.create_rectangle(
                x - 2, 0, x + 2, self.canvas_height, fill=color, tags="notes"
            )

    def update_canvas_height(self, event):
        canvas_height = (
            self.canvas_frame.winfo_height()
            - self.arrow_buttons_frame.winfo_height()
            - self.save_button_height
        )
        self.canvas.config(height=canvas_height)
        self.generate_visualizations()

    def select_song(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("MP3 files", "*.mp3"), ("RHYS files", "*.rhys")]
        )
        if file_path:
            self.song_path = file_path  # Assign the song path here
            if file_path.endswith(".rhys"):
                self.load_rhys_file(file_path)
            self.init_song()
            self.generate_visualizations()
            self.save_button.pack()
            self.arrow_buttons_frame.pack()

    def generate_visualizations(self):
        audio = MP3(self.song_path)
        meta = EasyMP3(self.song_path)
        song_length = audio.info.length
        song_name = meta["title"][0]
        song_artist = meta["artist"][0]
        self.song_length = song_length

        self.canvas.delete("all")

        mins = int(song_length / 60)
        secs = song_length % 60
        self.canvas.create_text(
            20, 10, text=f"{song_name}", anchor="nw", font=("Arial", 12)
        )
        self.canvas.create_text(
            20, 25, text=f"By: {song_artist}", anchor="nw", font=("Arial", 12)
        )
        self.canvas.create_text(
            20,
            40,
            text=f"Song Length: {mins}m:{secs}:s",
            anchor="nw",
            font=("Arial", 12),
        )

        canvas_height = (
            self.canvas_frame.winfo_height() - self.arrow_buttons_frame.winfo_height()
        )
        self.canvas.config(height=canvas_height)

        self.canvas_width = song_length * 100
        self.canvas.config(scrollregion=(0, 0, self.canvas_width, canvas_height))

        for sec in range(int(song_length) + 1):
            x = sec * 100
            self.canvas.create_line(x, 0, x, canvas_height, fill="gray")
            if sec % 5 == 0:
                self.canvas.create_text(
                    x + 10,
                    canvas_height - 20,
                    text=f"{sec}s",
                    font=("Arial", 10),
                    anchor="nw",
                )
            elif sec % 1 == 0:
                self.canvas.create_line(x, 0, x, canvas_height, fill="lightgray")

        self.update_cursor()

    def load_rhys_file(self, file_path):
        with open(file_path, "r") as f:
            data = json.load(f)
            encoded_audio = data["audio"]
            self.notes = data["notes"]

            # Decode the base64 audio and save it temporarily
            decoded_audio = base64.b64decode(encoded_audio)
            temp_audio_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
            temp_audio_file.write(decoded_audio)
            temp_audio_file.close()

            self.song_path = (
                temp_audio_file.name
            )  # Update the song_path to the temporary file

            self.generate_visualizations()
            for note_time, arrow in self.notes:
                x = note_time * 100
                color = self.get_arrow_color(arrow)
                self.canvas.create_rectangle(x - 2, 0, x + 2, 200, fill=color)

    def auto_add_note(self, event):
        print("Auto add")
        
        if self.song_playing:
            current_time = self.current_pos + pygame.mixer.music.get_pos() / 1000.0
        else:
            current_time = self.current_pos

        x = current_time / self.song_length * self.canvas_width
        
        time = x / self.get_canvas_width() * float(MP3(self.song_path).info.length)
        self.notes.append((time, self.selected_arrow))
        color = self.get_arrow_color(self.selected_arrow)
        self.canvas.create_rectangle(
            x - 2,
            0,
            x + 2,
            self.canvas_frame.winfo_height(),
            fill=color,
            tags=f"{time}",
        )

    def add_note(self, event):
        x = self.canvas.canvasx(
            event.x
        )  # Adjust x coordinate based on canvas scroll position
        time = x / self.get_canvas_width() * float(MP3(self.song_path).info.length)
        self.notes.append((time, self.selected_arrow))
        color = self.get_arrow_color(self.selected_arrow)
        self.canvas.create_rectangle(
            x - 2,
            0,
            x + 2,
            self.canvas_frame.winfo_height(),
            fill=color,
            tags=f"{time}",
        )

    def delete_note(self, event):
        x = self.canvas.canvasx(
            event.x
        )  # Adjust x coordinate based on canvas scroll position
        canvas_width = self.get_canvas_width()
        closest_note = None
        closest_distance = float("inf")

        delta = 4

        # Find the closest note to the mouse click
        for note_time, _ in self.notes:
            note_x = note_time / float(MP3(self.song_path).info.length) * canvas_width
            distance = abs(note_x - x)
            if distance < closest_distance and distance <= delta:
                closest_note = (note_time, _)
                self.canvas.delete(f"{note_time}")
                closest_distance = distance

        if closest_note:
            self.notes.remove(closest_note)
            self.redraw_notes(event)

    def select_arrow(self, direction):
        self.selected_arrow = direction

    def get_arrow_color(self, direction):
        if direction == "up":
            return "green"
        elif direction == "down":
            return "blue"
        elif direction == "left":
            return "red"
        elif direction == "right":
            return "yellow"

    def save_from_mp3(self):
        encoded_audio = ""
        with open(self.song_path, "rb") as f:
            encoded_audio = base64.b64encode(f.read()).decode("utf-8")

        # Get metadata information
        audio = EasyMP3(self.song_path)
        artist = audio.get("artist", [""])[0]
        album = audio.get("album", [""])[0]
        title = audio.get("title", [""])[0]
        track_time = int(audio.info.length)

        # Fallback to MP3 and ID3 to get cover art
        mp3_audio = MP3(self.song_path, ID3=ID3)
        cover_art_data = ""
        cover_art_filename = ""

        # Extract cover art
        for tag in mp3_audio.tags.values():
            if isinstance(tag, APIC):
                cover_art_data = base64.b64encode(tag.data).decode("utf-8")
                cover_art_filename = (
                    f"{os.path.splitext(os.path.basename(self.song_path))[0]}_cover.jpg"
                )


        formatted_notes = [{"t": note[0], "a": note[1]} for note in self.notes]

        # Fail-safe for missing metadata
        if not artist:
            artist = "Unknown Artist"
        if not album:
            album = "Unknown Album"
        if not title:
            title = "Unknown Track"

        song_data = {
            "notes": formatted_notes,
            "artist": artist,
            "album": album,
            "title": title,
            "track_time": track_time,
            "audio": encoded_audio,
            "cover_art": cover_art_data,
        }

        file_name = filedialog.asksaveasfilename(
            parent=root, title="Save as rhys:", filetypes=[("RHYS files", "*.rhys")]
        )

        if file_name:
            with open(f"{file_name}", "w", encoding="utf-8") as f:
                json.dump(song_data, f)
            print(f"Level saved to {file_name}")

            messagebox.showinfo("Success", f"Level saved to {file_name}")
        else:
            messagebox.showerror("Error", "You must provide a filename")

    def update_rhys(self):
        with open(self.song_path, "r+", encoding="utf-8") as f:
            data = json.load(f)
            data["notes"] = self.notes
            f.seek(0)
            json.dump(data, f)
            f.truncate()
        print("Notes added to rhys file")

    def save_level(self):
        if self.song_path:
            file_extension = os.path.splitext(self.song_path)[1]
            if file_extension == ".mp3":
                self.save_from_mp3()
            elif file_extension == ".rhys":
                self.update_rhys()


root = tk.Tk()
app = LevelDesigner(root)

root.mainloop()