# Neural Piano main Python module

import os
import io
from pathlib import Path
import importlib

import librosa
import soundfile as sf

import midirenderer

from .music2latent.inference import EncoderDecoder

from .master import master_mono_piano

def neuralpiano(input_midi_file,
                output_audio_file,
                return_audio=False
                ):
    
    filepath = Path(__file__).resolve()
    lib_root = filepath.parent
    models_dir = lib_root / "models"
    sf2_name = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"
    sf2_path = models_dir / sf2_name

    print('=' * 70)
    print('Neural Piano')
    print('=' * 70)
    
    print('Prepping model...')
    encdec = EncoderDecoder()

    pkg_root = Path(__file__).resolve().parent
    models_dir = pkg_root / "models"
    sf2_name = "SGM-v2.01-YamahaGrand-Guit-Bass-v2.7.sf2"
    sf2_path = models_dir / sf2_name

    if not sf2_path.exists():
        # call package downloader (same module or adjust import)
        try:
            # if download_models is in the same package module file, import accordingly
            from .your_downloader_module import download_models
            download_models()
        except Exception:
            pass

    if not sf2_path.exists():
        raise FileNotFoundError(f"Required soundfont {sf2_name} not found under {models_dir!s}.")

    print('Reading and rendering MIDI file...')
    wav_data = midirenderer.render_wave_from(
        sf2_path.read_bytes(),
        Path(input_midi_file).read_bytes()
    )
    
    print('Loading rendered MIDI...')
    with io.BytesIO(wav_data) as byte_stream:
        wv, sr = librosa.load(byte_stream, sr=48000)  # Music2Latent supports 48kHz audio as well
    
    print('Encoding...')
    latent = encdec.encode(wv)
    
    print('Rendering...')
    wv_rec = encdec.decode(latent)
    
    print('Mastering...')
    stereo, diag = master_mono_piano(wv_rec)
    
    print('Saving final audio...')
    sf.write(output_audio_file, stereo.squeeze(), sr)
    
    print('=' * 70)
    print('Done!')
    print('=' * 70)
    
    if return_audio:
        return stereo