# 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()

    # Ensure SF2 exists, try runtime download if not
    if not sf2_path.exists():
        print(f"Soundfont not found at {sf2_path!s}")
        # try to call a downloader if the package exposes one
        try:
            # assume download_models lives in the package top-level module or a util
            downloader = importlib.import_module(__package__ + ".download_models")
            if hasattr(downloader, "download_models"):
                print("Attempting to download required models...")
                downloader.download_models()
        except Exception:
            # ignore import/download errors here; check path again and raise useful message below
            pass

    if not sf2_path.exists():
        # fallback: search common locations under package
        candidates = list(models_dir.glob("*")) if models_dir.exists() else []
        if candidates:
            sf2_path = candidates[0]
            print(f"Using candidate soundfont {sf2_path.name}")
        else:
            # Helpful error that tells the user what to do
            raise FileNotFoundError(
                f"Required soundfont {sf2_name} not found under {models_dir!s}. "
                "Run the package downloader or install package with model files."
            )

    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