#=====================================================
# https://huggingface.co/spaces/asigalov61/MIDI-Doctor
#=====================================================

"""
Check and heal your MIDI :)
"""

# -----------------------------
# CORE MODULES
# -----------------------------

import os
import time as reqtime
import datetime
from pytz import timezone

import json

import gradio as gr

# -----------------------------
# CONFIGURATION & GLOBALS
# -----------------------------
SEP = '=' * 70
PDT = timezone('US/Pacific')

# -----------------------------
# PRINT START-UP INFO
# -----------------------------
def print_sep():
    print(SEP)

print_sep()
print("MIDI Doctor Gradio App")
print_sep()
print("Loading MIDI Doctor module...")

import mididoctor

print_sep()
print("Done loading modules!")
print_sep()

# -----------------------------
# HELPER FUNCTIONS
# -----------------------------

def pretty_print_dict(d, indent=0, as_string=True):
    """Recursively prints and optionally returns a pretty-formatted string of a nested dictionary."""
    lines = []

    def _format_dict(d, indent):
        for key, value in d.items():
            line = '    ' * indent + str(key) + ':'
            if isinstance(value, dict):
                lines.append(line)
                _format_dict(value, indent + 1)
            else:
                lines.append(line + ' ' + str(value))

    _format_dict(d, indent)

    for line in lines:
        print(line)

    if as_string:
        return '\n'.join(lines)

# -----------------------------

import json
import os

def save_nested_dict_to_json(data, filepath, indent=4, ensure_ascii=False, fsync_enabled=False):
    """
    Saves a nested dictionary to a JSON file, handling non-serializable objects
    and ensuring all data is flushed (and optionally fsynced) to disk.
    """
    def _default(obj):
        try:
            return obj.__dict__
        except Exception:
            return str(obj)

    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(
                data,
                f,
                indent=indent,
                ensure_ascii=ensure_ascii,
                default=_default
            )
            f.flush()
            if fsync_enabled:
                os.fsync(f.fileno())
    except Exception as e:
        print(f"Failed to save JSON to '{filepath}': {e}")

# -----------------------------
# MAIN FUNCTION
# -----------------------------

def Heal_MIDI(input_midi):
    """
    Heal MIDI
    """
    print_sep()
    print("Request start time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))

    print_sep()
    if input_midi is not None:
        fn = os.path.basename(input_midi.name)
        fn1 = fn.split('.')[0]
        print('Input file name:', fn)
        print_sep()

        print('Healing MIDI...')
        healed_MIDI_stats = mididoctor.heal_midi(input_midi.name, 
                                                 return_midi_hashes=True,
                                                 return_original_midi_stats=True,
                                                 return_healed_midi_stats=True,
                                                 return_repair_stats=True
                                                )

        out_fname = './healed_midis/'+fn

        print('Done!')
        print_sep()
        print('Saving stats dictionary...')
        save_nested_dict_to_json(healed_MIDI_stats, fn1+'_healed_MIDI_stats.json')
        print('Done!')
        print_sep()
        print('Healed MIDI stats:')
        print('-' * 70)
        healed_MIDI_stats_str = pretty_print_dict(healed_MIDI_stats)
        print_sep()
        print("Request end time:", datetime.datetime.now(PDT).strftime("%Y-%m-%d %H:%M:%S"))
        print_sep()
        
        return out_fname, fn1+'_healed_MIDI_stats.json', healed_MIDI_stats_str

# -----------------------------
# GRADIO INTERFACE SETUP
# -----------------------------
with gr.Blocks() as demo:

    gr.Markdown("<h1 style='text-align: left; margin-bottom: 1rem'>MIDI Doctor</h1>")
    gr.HTML("""
        Check out <a href="https://github.com/asigalov61/mididoctor">MIDI Doctor</a> on GitHub
        <p>
            <a href="https://huggingface.co/spaces/asigalov61/MIDI-Doctor?duplicate=true">
                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-md.svg" alt="Duplicate in Hugging Face">
            </a>
        </p>
    """)

    gr.Markdown(
        """
        ## MIDI Doctor will check and heal
        * MIDI signature
        * MIDI channels range
        * MIDI pitches range
        * MIDI velocities range
        * MIDI patch changes range
        * Duplicate pitches
        * Bad chords
        * Excessively short durations
        * Excessively long durations
        * Overlapping durations
        * Text events inconsistencies
        * Excessively quiet dynamics
        * Excessively flat dynamics
        """
    )

    gr.Markdown("## Upload your MIDI")
    input_midi = gr.File(label="Input MIDI", file_types=[".midi", ".mid", ".kar"])
    input_midi.upload()
    
    generate_btn = gr.Button("Heal", variant="primary")

    healed_MIDI_file = gr.File(label="Healed MIDI file")
    healed_MIDI_stats = gr.Textbox(label="Healed MIDI stats")
    healed_MIDI_stats_json = gr.File(label="Healed MIDI stats json")
            
    generate_btn.click(
        Heal_MIDI,
        [input_midi],
        [healed_MIDI_file,
         healed_MIDI_stats_json,
         healed_MIDI_stats
        ]
    )

demo.launch()