import gradio as gr
import os
import allin1
import zipfile
import shutil
from subprocess import Popen 
from concurrent.futures import ThreadPoolExecutor

from pathlib import Path

from dissector import generate_dissector_data

HEADER = """
<header style="text-align: center;">
  <h1>
    All-In-One Music Structure Analyzer 🔮
  </h1>
  <p>
    <a href="https://github.com/mir-aidj/all-in-one">[Python Package]</a>
    <a href="https://arxiv.org/abs/2307.16425">[Paper]</a>
    <a href="https://taejun.kim/music-dissector/">[Visual Demo]</a>
  </p>
</header>
<main
  style="display: flex; justify-content: center;"
>
  <div
    style="display: inline-block;"
  >
    <p>
      This Space demonstrates the music structure analyzer predicts:
      <ul
        style="padding-left: 1rem;"
      >
        <li>BPM</li>
        <li>Beats</li>
        <li>Downbeats</li>
        <li>Functional segment boundaries</li>
        <li>Functional segment labels (e.g. intro, verse, chorus, bridge, outro)</li>
      </ul>
    </p>
    <p>
      It takes about 4:30 to analyze a 3-minute song.
      If you are in a hurry, you can <strong><u>try the examples at the bottom</u></strong>.
    </p>
    <p>
      For more information, please visit the links above ✨🧸
    </p>
  </div>
</main>
"""

CACHE_EXAMPLES = os.getenv('CACHE_EXAMPLES', '1') == '1'

BASIC_PITCH_DIR = "basic-pitch"


def compress_files(demucs_output_path, dissector_file, original_audio):
    """Compresses specific files in the specified folder into a .zip file"""
    # List of the files to be compressed
    wav_files = ['bass.wav', 'drums.wav', 'other.wav', 'vocals.wav']
    mp3_files = [file.replace('.wav', '.mp3') for file in wav_files]

    def convert_to_mp3(wav, mp3):
        """Utility function to run the ffmpeg command"""
        command = ["ffmpeg", "-i", os.path.join(demucs_output_path, wav), os.path.join(demucs_output_path, mp3)]
        process = Popen(command)
        process.communicate()

    # Use ThreadPoolExecutor to convert .wav files to .mp3 in parallel
    with ThreadPoolExecutor(max_workers=8) as executor:
        list(executor.map(convert_to_mp3, wav_files, mp3_files))

    try:
        shutil.rmtree(BASIC_PITCH_DIR)
    except:  # FileNotFoundError on first run
        pass
    if not os.path.exists(BASIC_PITCH_DIR):
        os.makedirs(BASIC_PITCH_DIR)
    command = [
        "basic-pitch",
        '--save-note-events',
        BASIC_PITCH_DIR,
        os.path.join(demucs_output_path, 'bass.wav'),
        os.path.join(demucs_output_path, 'vocals.wav'),
        os.path.join(demucs_output_path, 'other.wav')
    ]
    process = Popen(command)
    process.communicate()

    # Compress the files
    zip_path = demucs_output_path + ".zip"
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for mp3 in mp3_files:
            zipf.write(os.path.join(demucs_output_path, mp3), arcname=mp3)
        zipf.write(f'{BASIC_PITCH_DIR}/bass_basic_pitch.mid', arcname='bass_basic_pitch.mid')
        zipf.write(f'{BASIC_PITCH_DIR}/vocals_basic_pitch.mid', arcname='vocals_basic_pitch.mid')
        zipf.write(f'{BASIC_PITCH_DIR}/other_basic_pitch.mid', arcname='other_basic_pitch.mid')
        zipf.write(f'{BASIC_PITCH_DIR}/bass_basic_pitch.csv', arcname='bass_basic_pitch.csv')
        zipf.write(f'{BASIC_PITCH_DIR}/vocals_basic_pitch.csv', arcname='vocals_basic_pitch.csv')
        zipf.write(f'{BASIC_PITCH_DIR}/other_basic_pitch.csv', arcname='other_basic_pitch.csv')
        zipf.write(dissector_file, arcname='dissector.json')  # Fixed name
        zipf.write(original_audio, arcname='mixdown.mp3')  # Added original audio
    
    return zip_path


def analyze(path):
  path = Path(path)
  result = allin1.analyze(
    path,
    keep_byproducts=True,  # TODO(taejun): remove this
  )

  fig = allin1.visualize(result)
  fig.set_dpi(300)

  allin1.sonify(result, out_dir='./sonif')
  sonif_path = Path(f'./sonif/{path.stem}.sonif{path.suffix}').resolve().as_posix()
      
  dissector_file = generate_dissector_data(path.stem, result)  
  compressed_file = compress_files(f"demix/htdemucs/{path.stem}", dissector_file, path.as_posix())

  return result.bpm, fig, sonif_path, compressed_file


with gr.Blocks() as demo:
  gr.HTML(HEADER)

  input_audio_path = gr.Audio(
    label='Input',
    source='upload',
    type='filepath',
    format='mp3',
    show_download_button=False,
  )
  button = gr.Button('Analyze', variant='primary')
  output_viz = gr.Plot(label='Visualization')
  with gr.Row():
    output_bpm = gr.Textbox(label='BPM', scale=1)
    output_sonif = gr.Audio(
      label='Sonification',
      type='filepath',
      format='mp3',
      show_download_button=False,
      scale=9,
    )
    output_compressed = gr.File(
        label="Compressed Files",
        type="file",
    )
  gr.Examples(
    examples=[
        './assets/kult.mp3',
        # './assets/krematorii.mp3',
      # './assets/NewJeans - Super Shy.mp3',
      # './assets/Bruno Mars - 24k Magic.mp3'
    ],
    inputs=input_audio_path,
    outputs=[output_bpm, output_viz, output_sonif, output_compressed],
    fn=analyze,
    cache_examples=CACHE_EXAMPLES,
  )
  button.click(
    fn=analyze,
    inputs=input_audio_path,
    outputs=[output_bpm, output_viz, output_sonif, output_compressed],
    api_name='analyze',
  )

if __name__ == '__main__':
  demo.launch()