import sys
import time
import os
import pathlib
import re
import traceback
from tkinter import filedialog

import rosu_pp_py as rosu # https://github.com/MaxOhn/rosu-pp-py
import ffmpeg # https://github.com/kkroening/ffmpeg-python

def get_line_by_prefix(lines: list[str], prefix: str) -> tuple[int, str]:
    i = next(i for i, line in enumerate(lines) if line.startswith(prefix))

    return i, lines[i].removeprefix(prefix).replace("\r", "")

def join(data: list, separators: list[str]) -> str:
    result = ""

    for item, separator in zip(data, separators):
        result += item + separator
    
    result += data[-1]

    return result

def split_timing_point_line(line: str) -> list[str]:
    data = line.split(",")

    return data

def read_timing_point_value(line: str, i: int) -> str:
    data = split_timing_point_line(line)

    return data[i]

def update_timing_point_line(line: str, i: int, offset: int) -> str:
    data = split_timing_point_line(line)
    data[i] = int(data[i])
    data[i] += offset

    return ",".join(str(n) for n in data)

def split_hit_object_line(line: str) -> tuple[list[str], list[str]]:
    separator_regex = "[,:|]"
    data = re.split(separator_regex, line)
    separators = re.findall(separator_regex, line)

    return data, separators

def read_hit_object_data(line: str, i: int) -> str:
    data, separators = split_hit_object_line(line)

    return data[i]

def update_hit_object_line(line: str, i: int, offset: int) -> str:
    data, separators = split_hit_object_line(line)
    data[i] = int(data[i])
    data[i] += offset

    return join([str(s) for s in data], separators)
    
def split_preview_point_line(line: str) -> tuple[str, str]:
    separator_regex = "[:]"
    data = re.split(separator_regex, line)
    separators = re.findall(separator_regex, line)

    return data, separators

def update_preview_point_line(line: str, offset: int) -> str:
    data, separators = split_preview_point_line(line)
    previewPoint = int(data[0])
    previewPoint += offset

    return f"PreviewTime: {previewPoint}"

def generate_osu_file(input_filename: str, offset: int = 0) -> str:
    with open(input_filename, "rb") as file:
        content: str = file.read().decode()

    lines = content.split("\n")
    offset_str = f"{offset}" if offset < 0 else f"+{offset}"

    _, artist =  get_line_by_prefix(lines, "Artist:")
    _, title =  get_line_by_prefix(lines, "Title:")
    _, creator =  get_line_by_prefix(lines, "Creator:")
    version_line_i, version =  get_line_by_prefix(lines, "Version:")

    filename = f"{artist} - {title} ({creator}) [{version}].osu"

    previewTime_line, previewTime =  get_line_by_prefix(lines, "PreviewTime:")
    lines[previewTime_line] = update_preview_point_line(previewTime, offset)

    try:
        video_offset_line, video_offset_and_name = get_line_by_prefix(lines, "Video,")
        video_offset_string = video_offset_and_name.split(",")[0]
        video_offset = int(video_offset_string) + offset
        new_video_offset_line = "Video," + str(video_offset) + "," + video_offset_and_name.split(",")[1]
        lines[video_offset_line] = new_video_offset_line
    except:
        pass

    timing_point_i, _ =  get_line_by_prefix(lines, "[TimingPoints]")
    i = timing_point_i + 1

    while True:
        if lines[i].strip() == "":
            break

        lines[i] = update_timing_point_line(lines[i], 0, offset)
        i += 1

    hit_objects_i, _ =  get_line_by_prefix(lines, "[HitObjects]")
    i = hit_objects_i + 1

    while True:
        if i >= len(lines) or lines[i].strip() == "":
            break

        lines[i] = update_hit_object_line(lines[i], 2, offset)
        i += 1
    
    updated_content = "\n".join(lines).replace("\r", "")
    
    return updated_content

def export_best_nm_map(input_filename: str, output_folder : str, best_offset_nm: int, map_name: str) -> None:
    best_nomod_map = generate_osu_file(input_filename=input_filename, offset = best_offset_nm)

    best_nomod_map_lines = best_nomod_map.split("\n")
    audio_line_i, audio_name =  get_line_by_prefix(best_nomod_map_lines, "AudioFilename: ")

    best_nomod_map_lines[audio_line_i] = f"AudioFilename: {audio_name}"
    version_line_i, version =  get_line_by_prefix(best_nomod_map_lines, "Version:")
    best_nomod_map_lines[version_line_i] = f"Version:{version}"

    best_nomod_map_path = os.path.join(output_folder, map_name)

    best_nomod_file = "\n".join(best_nomod_map_lines).replace("\r", "")
    os.makedirs(output_folder, exist_ok=True)
    with open(best_nomod_map_path, "w", encoding="utf-8") as file:
        file.write(best_nomod_file)  

    return

def export_best_nm_mapset(input_filename: str, output_folder : str, best_offset_nm: int, map_name: str) -> None:
    files = [""]
    mapset_folder = os.path.dirname(input_filename)
    for file in os.listdir(mapset_folder):
        if(file.endswith(".osu")):
            files.append(file)
    files.pop(0)

    os.makedirs(output_folder, exist_ok=True)
    for file in files:
        input_filename = os.path.join(mapset_folder, file)
        export_best_nm_map(input_filename, output_folder, best_offset_nm, file)
    return

def export_best_dt_map(input_filename: str, output_folder : str, best_offset_dt: int, map_name: str) -> None:
    best_dt_map = generate_osu_file(input_filename=input_filename, offset = best_offset_dt)

    best_dt_map_lines = best_dt_map.split("\n")

    audio_line_i, audio_name =  get_line_by_prefix(best_dt_map_lines, "AudioFilename: ")
    best_dt_map_lines[audio_line_i] = f"AudioFilename: {audio_name}"

    version_line_i, version =  get_line_by_prefix(best_dt_map_lines, "Version:")
    best_dt_map_lines[version_line_i] = f"Version:{version}"

    best_dt_map_path = os.path.join(output_folder, map_name)

    best_dt_file = "\n".join(best_dt_map_lines).replace("\r", "")
    os.makedirs(output_folder, exist_ok=True)
    with open(best_dt_map_path, "w", encoding="utf-8") as file:
        file.write(best_dt_file)  

    return

def export_best_dt_mapset(input_filename: str, output_folder : str, best_offset_dt: int, map_name: str) -> None:
    files = [""]
    mapset_folder = os.path.dirname(input_filename)
    for file in os.listdir(mapset_folder):
        if(file.endswith(".osu")):
            files.append(file)
    files.pop(0)

    os.makedirs(output_folder, exist_ok=True)
    for file in files:
        input_filename = os.path.join(mapset_folder, file)
        export_best_dt_map(input_filename, output_folder, best_offset_dt, file)
    return

def export_best_nm_audio(output_folder: str, best_offset_nm : int, lines: list[str]) -> None:
    audio_line_number, audio_filename =  get_line_by_prefix(lines, "AudioFilename: ")
    audio_filepath = output_folder.removesuffix("\\Offsetmaxxed")
    audio_filepath = os.path.join(audio_filepath, audio_filename)
    audio = ffmpeg.input(audio_filepath, itsoffset=best_offset_nm)

    audio_export_filepath = os.path.join(output_folder, audio_filename)

    (
    ffmpeg
        .output(audio, audio_export_filepath, af=f"adelay={best_offset_nm}|{best_offset_nm}", audio_bitrate="192k", loglevel="quiet")
        .overwrite_output()
        .run(capture_stdout=True, capture_stderr=True)
    )
    return

def export_best_dt_audio(output_folder: str, best_offset_dt : int, lines: list[str]) -> None:
    audio_line_number, audio_filename =  get_line_by_prefix(lines, "AudioFilename: ")
    audio_filepath = output_folder.removesuffix("\\offset")
    audio_filepath = os.path.join(audio_filepath, audio_filename)
    audio = ffmpeg.input(audio_filepath, itsoffset=best_offset_dt)

    audio_export_filepath = os.path.join(output_folder, audio_filename)
    (
    ffmpeg
        .output(audio, audio_export_filepath, af=f"adelay={best_offset_dt}|{best_offset_dt}", audio_bitrate="192k", loglevel="quiet")
        .overwrite_output()
        .run(capture_stdout=True, capture_stderr=True)
    )
    return

def restart_main():
    print("Restarting.")
    print()
    print()
    time.sleep(1)
    sys.stdout.flush()
    Offsetmaxxer()
    # # Restart the application...
    # executable = sys.executable
    # executable_filename = os.path.split(executable)[1]
    # if executable_filename.lower().startswith('python'):
    #     # application is running within a python interpreter
    #     python = executable
    #     os.execv(python, [python, ] + sys.argv)
    #     pass
    # else:
    #     # application is running as a standalone executable
    #     os.execv(executable, sys.argv)
    #     pass
    # pass

def show_startup_info():
    print("Created by The Fart Lord (users/7912447) & Fred727 (users/5055826). Message The Fart Lord on osu! if you encounter any errors or if you have any suggestions!")
    print("Thanks to Badewanne3 for the fast python pp calculator!")
    print("More info can be found in the README.md")
    print()
    time.sleep(0.5)

def main():
    show_startup_info()
    Offsetmaxxer()


def Offsetmaxxer():
    max_offset = 601 #redundancy, two extra calcs shouldn't slow things down much anyway
    offset_iterator = 0

    print("Select an .osu file.")
    print()
    time.sleep(0.5)
    input_filename = filedialog.askopenfilename(filetypes=[("osu file", ".osu")])
    output_folder = str(pathlib.Path(input_filename).parent.resolve()) + "\\Offsetmaxxer"

    if input_filename.strip() == "":
        print("No file was selected.")
        sys.exit()

    selected_beatmap = rosu.Beatmap(path = input_filename)

    nomod_pp_stable = rosu.Performance(lazer=False)
    nomod_pp = rosu.Performance()
    attrs = nomod_pp.calculate(selected_beatmap)
    best_star_rating_nm = attrs.difficulty.stars
    best_pp_nm = attrs.pp
    best_offset_nm = 0
    worst_star_rating_nm = attrs.difficulty.stars
    worst_pp_nm = attrs.pp
    worst_offset_nm = 0

    dt_pp_stable = rosu.Performance(lazer=False)
    dt_pp = rosu.Performance()
    dt_pp_stable.set_mods(64) #DT is mod 64
    dt_pp.set_mods(64)
    dt_attrs = dt_pp.calculate(selected_beatmap)
    best_star_rating_dt = dt_attrs.difficulty.stars
    best_pp_dt = dt_attrs.pp
    best_offset_dt = 0

    worst_star_rating_dt = dt_attrs.difficulty.stars
    worst_pp_dt = dt_attrs.pp
    worst_offset_dt = 0

    map_name = os.path.basename(input_filename)
    with open(input_filename, "rb") as file:
        content: str = file.read().decode()
    lines = content.split("\n")
    line_number, _ =  get_line_by_prefix(lines, "[TimingPoints]")
    original_offset = lines[line_number+1].split(",")[0]

    print(f"Selected beatmap: {map_name}")
    print()
    print(f"Original NM stats:")
    print("---------")
    print(f"Star rating: {attrs.difficulty.stars:.3f}*")
    print(f"Lazer pp: {attrs.pp:.3f}")
    print(f"Stable pp: {nomod_pp_stable.calculate(selected_beatmap).pp:.3f}")
    print("---------")
    print()
    print(f"Original DT stats:")
    print("---------")
    print(f"Star rating: {dt_attrs.difficulty.stars:.3f}*")
    print(f"Lazer pp: {dt_attrs.pp:.3f}")
    print(f"Stable pp: {dt_pp_stable.calculate(selected_beatmap).pp:.3f}")
    print("---------")
    print(f"Original offset: {original_offset}")
    print("---------")
    print()

    print("Select action:")
    print("1. Find perfect Nomod and DT offsets for selected difficulty.")
    print("2. Find perfect Nomod offset for selected difficulty.")
    print("3. Find perfect DT offset for selected difficulty.")
    print()
    selected_action = input()
    print(f"Selected: {selected_action}")

    match selected_action:
        case "1":
            print("Finding perfect offset...")
            print()
            while(offset_iterator < max_offset):
                offsetted_beatmap_file =  generate_osu_file(input_filename = input_filename, offset = offset_iterator)
                offsetted_beatmap = rosu.Beatmap(content = offsetted_beatmap_file)

                nm_attrs = nomod_pp.calculate(offsetted_beatmap)
                if(nm_attrs.difficulty.stars >= best_star_rating_nm):
                    best_star_rating_nm = nm_attrs.difficulty.stars
                    best_pp_nm = nm_attrs.pp
                    best_pp_nm_stable = nomod_pp_stable.calculate(offsetted_beatmap).pp
                    best_offset_nm = offset_iterator

                if(nm_attrs.difficulty.stars <= worst_star_rating_nm):
                    worst_star_rating_nm = nm_attrs.difficulty.stars
                    worst_pp_nm = nm_attrs.pp
                    worst_pp_nm_stable = nomod_pp_stable.calculate(offsetted_beatmap).pp
                    worst_offset_nm = offset_iterator

                dt_attrs = dt_pp.calculate(offsetted_beatmap)
                if(dt_attrs.difficulty.stars >= best_star_rating_dt):
                    best_star_rating_dt = dt_attrs.difficulty.stars
                    best_pp_dt = dt_attrs.pp
                    best_pp_dt_stable = dt_pp_stable.calculate(offsetted_beatmap).pp
                    best_offset_dt = offset_iterator

                if(dt_attrs.difficulty.stars <= worst_star_rating_dt):
                    worst_star_rating_dt = dt_attrs.difficulty.stars
                    worst_pp_dt = dt_attrs.pp
                    worst_pp_dt_stable = dt_pp_stable.calculate(offsetted_beatmap).pp
                    worst_offset_dt = offset_iterator

                offset_iterator = offset_iterator + 1
            print(f"Best NM stats:")
            print("---------")
            print(f"Star rating: {best_star_rating_nm:.3f}*")
            print(f"Lazer pp: {best_pp_nm:.3f}")
            print(f"Stable pp: {best_pp_nm_stable:.3f}")
            print(f"Offset: {best_offset_nm + int(original_offset)}")
            print("---------")

            print(f"Worst NM stats:")
            print("---------")
            print(f"Star rating: {worst_star_rating_nm:.3f}*")
            print(f"Lazer pp: {worst_pp_nm:.3f}")
            print(f"Stable pp: {worst_pp_nm_stable:.3f}")
            print(f"Offset: {worst_offset_nm + int(original_offset)}")
            print("---------")

            print(f"Best DT stats:")
            print("---------")
            print(f"Star rating: {best_star_rating_dt:.3f}*")
            print(f"Lazer pp: {best_pp_dt:.3f}")
            print(f"Stable pp: {best_pp_dt_stable:.3f}")
            print(f"Offset: {best_offset_dt + int(original_offset)}")
            print("---------")
            
            print(f"Worst DT stats:")
            print("---------")
            print(f"Star rating: {worst_star_rating_dt:.3f}*")
            print(f"Lazer pp: {worst_pp_dt:.3f}")
            print(f"Stable pp: {worst_pp_dt_stable:.3f}")
            print(f"Offset: {worst_offset_dt + int(original_offset)}")
            print("---------")

        case "2":
            print("Finding perfect offset...")
            print()
            while(offset_iterator < max_offset):
                offsetted_beatmap_file =  generate_osu_file(input_filename = input_filename, offset = offset_iterator)
                offsetted_beatmap = rosu.Beatmap(content = offsetted_beatmap_file)

                nm_attrs = nomod_pp.calculate(offsetted_beatmap)
                if(nm_attrs.difficulty.stars >= best_star_rating_nm):
                    best_star_rating_nm = nm_attrs.difficulty.stars
                    best_pp_nm = nm_attrs.pp
                    best_pp_nm_stable = nomod_pp_stable.calculate(offsetted_beatmap).pp
                    best_offset_nm = offset_iterator
                    
                if(nm_attrs.difficulty.stars <= worst_star_rating_nm):
                    worst_star_rating_nm = nm_attrs.difficulty.stars
                    worst_pp_nm = nm_attrs.pp
                    worst_pp_nm_stable = nomod_pp_stable.calculate(offsetted_beatmap).pp
                    worst_offset_nm = offset_iterator

                offset_iterator = offset_iterator + 1
            print(f"Best NM stats:")
            print("---------")
            print(f"Star rating: {best_star_rating_nm:.3f}*")
            print(f"Lazer pp: {best_pp_nm:.3f}")
            print(f"Stable pp: {best_pp_nm_stable:.3f}")
            print(f"Offset: {best_offset_nm + int(original_offset)}")
            print("---------")

            print(f"Worst NM stats:")
            print("---------")
            print(f"Star rating: {worst_star_rating_nm:.3f}*")
            print(f"Lazer pp: {worst_pp_nm:.3f}")
            print(f"Stable pp: {worst_pp_nm_stable:.3f}")
            print(f"Offset: {worst_offset_nm + int(original_offset)}")
            print("---------")
        case "3":
            print("Finding perfect offset...")
            print()
            while(offset_iterator < max_offset):
                offsetted_beatmap_file =  generate_osu_file(input_filename = input_filename, offset = offset_iterator)
                offsetted_beatmap = rosu.Beatmap(content = offsetted_beatmap_file)

                dt_attrs = dt_pp.calculate(offsetted_beatmap)
                if(dt_attrs.difficulty.stars >= best_star_rating_dt):
                    best_star_rating_dt = dt_attrs.difficulty.stars
                    best_pp_dt = dt_attrs.pp
                    best_pp_dt_stable = dt_pp_stable.calculate(offsetted_beatmap).pp
                    best_offset_dt = offset_iterator
                    
                if(dt_attrs.difficulty.stars <= worst_star_rating_dt):
                    worst_star_rating_dt = dt_attrs.difficulty.stars
                    worst_pp_dt = dt_attrs.pp
                    worst_pp_dt_stable = dt_pp_stable.calculate(offsetted_beatmap).pp
                    worst_offset_dt = offset_iterator

                offset_iterator = offset_iterator + 1
            print(f"Best DT stats:")
            print("---------")
            print(f"Star rating: {best_star_rating_dt:.3f}*")
            print(f"Lazer pp: {best_pp_dt:.3f}")
            print(f"Stable pp: {best_pp_dt_stable:.3f}")
            print(f"Offset: {best_offset_dt + int(original_offset)}")
            print("---------")
            
            print(f"Worst DT stats:")
            print("---------")
            print(f"Star rating: {worst_star_rating_dt:.3f}*")
            print(f"Lazer pp: {worst_pp_dt:.3f}")
            print(f"Stable pp: {worst_pp_dt_stable:.3f}")
            print(f"Offset: {worst_offset_dt + int(original_offset)}")
            print("---------")
        case _:
            print("Invalid input.")
            restart_main()

    print("Select action:")
    print("1. Apply best Nomod offset to all difficulties.")
    print("2. Apply best DT offset to all difficulties.")
    print("3. Apply best Nomod offset to selected difficulty only")
    print("4. Apply best DT offset to selected difficulty only")
    print("5. Select another map file.")
    print("6. Exit.")
    print()
    selected_action = input()
    print(f"Selected: {selected_action}")
    print()

    match selected_action:
        case "1":
            print("Applying perfect offset...")
            print()
            export_best_nm_mapset(input_filename, output_folder, best_offset_nm, map_name)
        case "2":
            print("Applying perfect offset...")
            print()
            export_best_dt_mapset(input_filename, output_folder, best_offset_dt, map_name)
        case "3":
            print("Applying perfect offset...")
            print()
            export_best_nm_map(input_filename, output_folder, best_offset_nm, map_name)
        case "4":
            print("Applying perfect offset...")
            print()
            export_best_dt_map(input_filename, output_folder, best_offset_dt, map_name)
        case "5":
            restart_main()
        case "6":
            sys.exit()
        case _:
            print("Invalid input.")
            restart_main()

    path = os.path.realpath(output_folder)
    os.startfile(path)
    print("Beatmap file/s have been generated!")
    print()

    print("Select action:")
    print("1. Export best DT offset audio. (Requires FFmpeg!)")
    print("2. Export best NM offset audio. (Requires FFmpeg!)")
    print("3. Select another map file.")
    print("4. What is FFmpeg?.")
    print("5. Exit.")
    print()
    selected_action = input()
    print(f"Selected: {selected_action}")
    
    try:
        match selected_action:
            case "1":
                export_best_dt_audio(output_folder, best_offset_dt, lines)
            case "2":
                export_best_nm_audio(output_folder, best_offset_nm, lines)
            case "3":
                restart_main()
            case "4":
                print("FFmpeg is required to generate audio files.")
                print("https://github.com/BtbN/FFmpeg-Builds/releases")
                print("If you have FFmpeg installed, try running this executable as administrator or placing \"ffmpeg.exe\" into the same folder.")
                print()
                print("Simple instructions for Windows users:")
                print("Direct zip download link for Windows: https://github.com/BtbN/FFmpeg-Builds/releases/download/autobuild-2024-12-30-12-59/ffmpeg-n7.1-62-gb168ed9b14-win64-gpl-7.1.zip")
                print("Extract and place \"ffmpeg.exe\" into the same directory as this executable then try again.")
                print()
                input("Press enter to exit. ")
                sys.exit()
            case "5":
                sys.exit()
            case _:
                print("Invalid input.")
                restart_main()
    except ffmpeg.Error as ffmpegerror:
        print('stdout:', ffmpegerror.stdout)
        print('stderr:', ffmpegerror.stderr)
        raise ffmpegerror
    except Exception as e:
        print("FFmpeg is required to generate audio files.")
        print("https://github.com/BtbN/FFmpeg-Builds/releases")
        print("If you have FFmpeg installed, try running Offsetmaxxer as administrator or placing \"ffmpeg.exe\" into the same folder.")
        print()
        print("Simple instructions for Windows users:")
        print("Direct zip download link for Windows: https://github.com/BtbN/FFmpeg-Builds/releases/download/autobuild-2024-12-30-12-59/ffmpeg-n7.1-62-gb168ed9b14-win64-gpl-7.1.zip")
        print("Extract and place \"ffmpeg.exe\" into the same directory as Offsetmaxxer then try again.")

    print()
    print("Audio file has been generated! Check your beatmap folder!")
    print()
    time.sleep(1)
    print("Select action:")
    print("1. Select another map file.")
    print("2. Exit.")
    print()
    selected_action = input()
    match selected_action:
        case "1":
            restart_main
        case "2":
            sys.exit()
        case _:
            print("Invalid input.")
            restart_main()
    time.sleep(1)
    restart_main()


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        traceback.print_exc()
        print()
        print("Some unexpected error happened")
        input("Press enter to exit. ")
