import gradio as gr
import os
import shutil
import zipfile
import librosa
import numpy as np
from pydub import AudioSegment
from moviepy.editor import AudioFileClip, ImageClip
import subprocess
from pathlib import Path
import sys

# --- Configuration ---
OUTPUT_DIR = Path("nightpulse_output")
TEMP_DIR = Path("temp_processing")

def process_track(audio_file, cover_art_image):
    # Initialize return variables
    zip_path = None
    video_path = None

    try:
        # --- 0. Input Validation ---
        if not audio_file:
            raise ValueError("No audio file provided.")
            
        # --- 1. Setup Directories ---
        if OUTPUT_DIR.exists():
            shutil.rmtree(OUTPUT_DIR)
        if TEMP_DIR.exists():
            shutil.rmtree(TEMP_DIR)
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        TEMP_DIR.mkdir(parents=True, exist_ok=True)
        
        filename = Path(audio_file).stem
        
        # --- 2. Analyze BPM & Key (Librosa) ---
        print(f"Analyzing {filename}...")
        try:
            # Mono load for robust BPM detection
            y, sr = librosa.load(audio_file, duration=60, mono=True)
            tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
            
            # Handle different librosa return types (float vs array)
            if np.ndim(tempo) > 0:
                detected_bpm = int(round(tempo[0]))
            else:
                detected_bpm = int(round(tempo))
            print(f"Detected BPM: {detected_bpm}")
        except Exception as e:
            print(f"BPM Warning: {e}")
            detected_bpm = 120 # Safe Fallback
        
        # --- 3. AI Stem Separation (Demucs) ---
        print("Separating stems...")
        try:
            # Call Demucs via subprocess to ensure clean execution
            subprocess.run([
                sys.executable, "-m", "demucs", 
                "-n", "htdemucs", 
                "--out", str(TEMP_DIR), 
                audio_file
            ], check=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Demucs failed: {e.stderr.decode()}")
        
        # Locate Stems (Robust Search)
        demucs_out = TEMP_DIR / "htdemucs"
        track_folder = next(demucs_out.iterdir(), None)
        
        if not track_folder:
             raise FileNotFoundError("Demucs output folder missing.")

        drums_path = track_folder / "drums.wav"
        melody_path = track_folder / "other.wav"
        bass_path = track_folder / "bass.wav"

        # --- 4. Loop Logic ---
        if detected_bpm <= 0: detected_bpm = 120
        ms_per_beat = (60 / detected_bpm) * 1000
        eight_bars_ms = ms_per_beat * 4 * 8
        
        def create_loop(source_path, output_name):
            if not source_path.exists(): return None, None
            
            audio = AudioSegment.from_wav(str(source_path))
            # Grab middle 8 bars
            start = len(audio) // 3
            end = start + eight_bars_ms
            
            if len(audio) < end:
                start = 0
                end = min(len(audio), eight_bars_ms)
                
            loop = audio[start:int(end)].fade_in(15).fade_out(15).normalize()
            out_name = OUTPUT_DIR / f"{detected_bpm}BPM_{output_name}.wav"
            loop.export(out_name, format="wav")
            return out_name, loop

        loop_drums, _ = create_loop(drums_path, "DrumLoop")
        loop_melody, _ = create_loop(melody_path, "MelodyLoop")
        create_loop(bass_path, "BassLoop")
        
        # --- 5. Video Generation ---
        if cover_art_image and loop_melody:
            print("Rendering Video...")
            try:
                vid_out = OUTPUT_DIR / "Promo_Video.mp4"
                audio_clip = AudioFileClip(str(loop_melody))
                img_clip = ImageClip(cover_art_image)
                
                # Resize to 1080w (maintain aspect ratio)
                img_clip = img_clip.resize(width=1080)
                img_clip = img_clip.set_duration(audio_clip.duration)
                img_clip = img_clip.set_audio(audio_clip)
                img_clip.fps = 24
                
                img_clip.write_videofile(str(vid_out), codec="libx264", audio_codec="aac", logger=None)
                video_path = str(vid_out)
            except Exception as e:
                print(f"Video skipped: {e}")

        # --- 6. Zip Export ---
        zip_file = "NightPulse_Pack.zip"
        with zipfile.ZipFile(zip_file, 'w') as zf:
            for f in OUTPUT_DIR.iterdir():
                zf.write(f, f.name)
        zip_path = zip_file
        
        return zip_path, video_path

    except Exception as e:
        raise gr.Error(f"System Error: {str(e)}")

# --- UI Definition (Corrected) ---
iface = gr.Interface(
    fn=process_track,
    inputs=[
        gr.Audio(type="filepath", label="Upload Suno Track"),
        gr.Image(type="filepath", label="Upload Cover Art")
    ],
    outputs=[
        gr.File(label="Download ZIP"),
        gr.Video(label="Preview Video")
    ],
    title="Night Pulse Audio | Automator",
    description="Night Pulse Pipeline v1.1 (Stable)"
)

if __name__ == "__main__":
    iface.launch()