backend/services/audio_service.py

import os
import shutil
import zipfile
import subprocess
import time
from pathlib import Path
from typing import Callable
from faster_whisper import WhisperModel

from audio_separator.separator import Separator

from pydub import AudioSegment
import torch
import logging
import threading

class ProgressSimulator:
    def __init__(self, callback, start_pct, end_pct, duration_sec, message):
        self.callback = callback
        self.start_pct = start_pct
        self.end_pct = end_pct
        self.duration_sec = duration_sec if duration_sec > 0 else 60
        self.message = message
        self.stop_event = threading.Event()
        self.thread = threading.Thread(target=self._run, daemon=True)
        
    def start(self):
        self.thread.start()
        
    def stop(self):
        self.stop_event.set()
        if self.thread.is_alive():
            self.thread.join(timeout=1.0)
        
    def _run(self):
        start_time = time.time()
        while not self.stop_event.is_set():
            elapsed = time.time() - start_time
            if elapsed >= self.duration_sec:
                current_pct = self.end_pct - 1
            else:
                progress = elapsed / self.duration_sec
                # Linear progress
                current_pct = int(self.start_pct + (self.end_pct - self.start_pct) * progress)
                current_pct = min(current_pct, self.end_pct - 1)
                
            self.callback(current_pct, self.message)
            self.stop_event.wait(3.0)

# Base directory for all temporary files during processing
WORK_DIR = Path("temp_workdir")

# Factory for Separator instances — create fresh to avoid stale model state
def create_separator(output_dir=None, output_format="WAV"):
    sep = Separator(log_level=logging.WARNING)
    if output_dir:
        sep.output_dir = str(output_dir)
    sep.output_format = output_format
    return sep

def process_audio_file(
    file_path: str, 
    task_id: str, 
    progress_callback: Callable[..., None],
    isolate_vocals: bool = False,

    enhance_speech: bool = False,
    lyric_sync: bool = False,
    fast_mode: bool = True
) -> str:
    """
    1. Converts uploaded file to WAV format.
    2. Isolates stems natively using Demucs (no API limits).
    3. Merges files and returns the zip path.
    """
    task_dir = WORK_DIR / task_id
    task_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        # Step 1: Convert audio to standard WAV format
        progress_callback(10, "Normalizing audio formats...")
        downloaded_audio_path = task_dir / "converted.wav"
        
        try:
            subprocess.run([
                "ffmpeg", "-y", "-i", str(file_path),
                "-ar", "44100", "-ac", "2", str(downloaded_audio_path)
            ], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"FFmpeg failed to convert the file. Is it a valid audio/video file? {e}")

        final_dir = task_dir / "final_stems"
        final_dir.mkdir(exist_ok=True)
        shutil.copy(str(downloaded_audio_path), str(final_dir / "original.wav"))
        
        # Determine audio duration for progress simulators
        audio_duration = 0
        try:
            audio_info = AudioSegment.from_wav(str(downloaded_audio_path))
            audio_duration = len(audio_info) / 1000.0  # in seconds
        except:
            pass
        
        # Step 2: Stem Separation
        if isolate_vocals:
            progress_callback(30, "Analyzing audio and isolating core stems (takes a few minutes)...")
            from core.state import state
                
            sep = create_separator(output_dir=final_dir, output_format="WAV")
            
            # HIGH QUALITY MODE (MDX23C)
            progress_callback(40, "Running high-quality separation...")
            sep.load_model(model_filename='UVR-MDX-NET-Inst_HQ_3.onnx')
            
            if sep.model_instance is None:
                raise RuntimeError(f"Failed to load separation model. The model may not have downloaded correctly.")
            
            sim = ProgressSimulator(progress_callback, 40, 69, audio_duration * 1.5, "Running high-quality separation (this will take a while)...")
            sim.start()
            try:
                out_files = sep.separate(str(downloaded_audio_path))
            finally:
                sim.stop()
            
            # audio-separator may return relative paths or just filenames.
            # Resolve them against the output directory to get the actual file path.
            for f in out_files:
                f_path = Path(f)
                if not f_path.is_absolute() or not f_path.exists():
                    f_path = final_dir / f_path.name
                if not f_path.exists():
                    print(f"WARNING: Separator output file not found: {f} (resolved: {f_path})")
                    continue
                f_name = f_path.name.lower()
                
                if "vocals" in f_name or "vocal" in f_name:
                    shutil.move(str(f_path), str(final_dir / "vocals.wav"))
                elif "instrumental" in f_name or "inst" in f_name:
                    shutil.move(str(f_path), str(final_dir / "instrumental.wav"))
                        
            # Free memory explicitly to prevent Hugging Face OOM crashes
            del sep
            import gc
            gc.collect()
        
        # Step 3: Enhance Speech (Optional)
        if enhance_speech:
            progress_callback(70, "Enhancing vocal clarity using Studio Denoise AI...")
            from df.enhance import enhance, init_df, load_audio, save_audio
            
            target_for_enhance = final_dir / "vocals.wav"
            if not target_for_enhance.exists():
                target_for_enhance = downloaded_audio_path
                
            model, df_state, _ = init_df()  # Load default DeepFilterNet3 model
            audio, _ = load_audio(str(target_for_enhance), sr=df_state.sr())
            
            sim = ProgressSimulator(progress_callback, 70, 79, audio_duration * 0.2, "Enhancing vocal clarity using Studio Denoise AI...")
            sim.start()
            try:
                enhanced = enhance(model, df_state, audio)
            finally:
                sim.stop()
            
            out_path = final_dir / "enhanced.wav"
            save_audio(str(out_path), enhanced, df_state.sr())
            
            # Replace the vocals file with the enhanced version if it exists
            if target_for_enhance.name == "vocals.wav":
                shutil.copy(str(out_path), str(final_dir / "vocals.wav"))
                
            # Free memory explicitly to prevent OOM
            del model
            del df_state
            del audio
            del enhanced
            import gc
            gc.collect()
            
        # Step 5: Whisper Lyric Sync (Optional)
        target_audio_for_whisper = None
        if (final_dir / "vocals.wav").exists():
            target_audio_for_whisper = final_dir / "vocals.wav"
        else:
            target_audio_for_whisper = final_dir / "original.wav"
            
        if lyric_sync and target_audio_for_whisper.exists():
            progress_callback(80, "Transcribing vocals and generating synced lyrics...")
            import json
            import torch
            
            # Auto-detect GPU for massive speed boosts if upgraded on Hugging Face
            device = "cuda" if torch.cuda.is_available() else "cpu"
            compute_type = "float16" if device == "cuda" else "int8"
            
            # Use large-v3-turbo model for significantly better Telugu accuracy
            model = WhisperModel("deepdml/faster-whisper-large-v3-turbo-ct2", device=device, compute_type=compute_type)
            segments, info = model.transcribe(
                str(target_audio_for_whisper), 
                word_timestamps=True,
                language="te",  # Hardcoded exactly to Telugu
                condition_on_previous_text=False
            )
            
            srt_path = final_dir / "lyrics.srt"
            json_path = final_dir / "transcript.json"
            word_list = []
            
            # Dataset Setup
            dataset_dir = final_dir / "dataset"
            dataset_audio_dir = dataset_dir / "audio"
            dataset_audio_dir.mkdir(parents=True, exist_ok=True)
            dataset_metadata_path = dataset_dir / "metadata.csv"
            
            # Load full audio for slicing
            full_audio = AudioSegment.from_wav(str(target_audio_for_whisper))
            
            timeline_markers = []
            prev_end = 0.0
            
            with open(srt_path, "w", encoding="utf-8") as f, open(dataset_metadata_path, "w", encoding="utf-8") as metadata_f:
                for i, segment in enumerate(segments, start=1):
                    # Native progress reporting based on timestamp
                    if audio_duration > 0:
                        pct = 80 + int(15 * (segment.end / audio_duration))
                        progress_callback(min(pct, 95), "Transcribing vocals and generating synced lyrics...")
                        
                    def format_time(seconds):
                        hours = int(seconds // 3600)
                        minutes = int((seconds % 3600) // 60)
                        secs = int(seconds % 60)
                        millis = int((seconds - int(seconds)) * 1000)
                        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
                    
                    # SRT Writing
                    f.write(f"{i}\n")
                    f.write(f"{format_time(segment.start)} --> {format_time(segment.end)}\n")
                    text = segment.text.strip()
                    f.write(f"{text}\n\n")
                    
                    # Timeline markers
                    if segment.start - prev_end > 2.0 and prev_end > 0:
                        timeline_markers.append({
                            "type": "silence",
                            "start": prev_end,
                            "end": segment.start,
                            "duration": round(segment.start - prev_end, 2)
                        })
                    prev_end = segment.end
                    timeline_markers.append({
                        "type": "speech",
                        "start": segment.start,
                        "end": segment.end,
                        "text": text
                    })
                    
                    # Dataset Generation
                    start_ms = int(segment.start * 1000)
                    end_ms = int(segment.end * 1000)
                    chunk = full_audio[start_ms:end_ms]
                    wav_filename = f"{i:04d}.wav"
                    chunk.export(str(dataset_audio_dir / wav_filename), format="wav")
                    metadata_f.write(f"audio/{wav_filename}|{text}\n")
                    
                    if hasattr(segment, 'words') and segment.words:
                        for word_info in segment.words:
                            word_list.append({
                                "word": word_info.word.strip(),
                                "start": word_info.start,
                                "end": word_info.end
                            })
                            
            with open(json_path, "w", encoding="utf-8") as jf:
                json.dump({"words": word_list}, jf, indent=2)
                
            csv_path = final_dir / "transcript.csv"
            import csv
            with open(csv_path, "w", newline="", encoding="utf-8") as cf:
                writer = csv.writer(cf)
                writer.writerow(["word", "start", "end"])
                for w in word_list:
                    writer.writerow([w["word"], w["start"], w["end"]])
            
            # Write Timeline Intelligence (Silence Detection)
            timeline_path = final_dir / "timeline_markers.json"
            with open(timeline_path, "w", encoding="utf-8") as tf:
                json.dump({"markers": timeline_markers}, tf, indent=2)
                
            # Free memory explicitly
            del model
            del full_audio
            import gc
            gc.collect()

        # Step 9: Zip the results
        progress_callback(90, "Finalizing and packaging your stems...")
        zip_path = WORK_DIR / f"vocals_{task_id}.zip"
        
        with zipfile.ZipFile(str(zip_path), 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(str(final_dir)):
                for file in files:
                    file_path = os.path.join(root, file)
                    rel_path = os.path.relpath(file_path, str(final_dir))
                    
                    # Determine target folder for architecture
                    arc_folder = ""
                    if rel_path.startswith("dataset"):
                        # Preserve dataset folder structure exactly
                        arcname = f"Stemify_Project/Transcripts/{rel_path.replace(os.sep, '/')}"
                        zipf.write(file_path, arcname=arcname)
                        continue
                        
                    if file in ["vocals.wav", "instrumental.wav", "bass.wav", "drums.wav", "original.wav"]:
                        arc_folder = "Audio_Stems"
                    elif file in ["vocals_enhanced.wav", "vocals_dry.wav"]:
                        arc_folder = "Enhanced_Audio"
                    elif file.startswith("speaker_") and file.endswith(".wav"):
                        arc_folder = "Speaker_Separation"
                    elif file.endswith(".mid"):
                        arc_folder = "MIDI_Notes"
                    elif file in ["lyrics.srt", "transcript.json", "transcript.csv", "aligned_timestamps.json", "aligned_timestamps.csv"]:
                        arc_folder = "Transcripts"
                    else:
                        arc_folder = "Misc"
                        
                    arcname = f"Stemify_Project/{arc_folder}/{file}"
                    zipf.write(file_path, arcname=arcname)
                    
        progress_callback(100, "Done!")
        return str(zip_path)
        
    finally:
        pass



def package_custom_download(
    task_id: str, 
    stems: list[str], 
    output_format: str = "wav", 
    chunked: bool = False,
    folder_name: str = "custom_stems"
) -> str:
    """
    Packages a custom download based on the UI selections.
    Supports on-the-fly format conversion (mp3/wav) and chunking into 50s segments.
    """
    task_dir = WORK_DIR / task_id
    final_stems_dir = task_dir / "final_stems"
    
    if not final_stems_dir.exists():
        raise FileNotFoundError("Processing not complete or files missing.")
        
    # Create a unique staging directory for this specific download configuration
    timestamp = int(time.time())
    staging_dir = task_dir / f"custom_pkg_{timestamp}"
    
    # Nest inside the custom folder name so the ZIP extracts cleanly into a single folder
    base_dir = staging_dir / folder_name
    base_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        for stem in stems:
            stem_file = final_stems_dir / f"{stem.lower()}.wav"
            if not stem_file.exists():
                continue
                
            # Create subfolder for this stem e.g., "Vocals/"
            stem_folder = base_dir / stem.capitalize()
            stem_folder.mkdir(exist_ok=True)
            
            if chunked:
                # Use FFmpeg to split into 50 second chunks ad0001, ad0002...
                ext = output_format.lower()
                out_pattern = str(stem_folder / f"ad%04d.{ext}")
                
                cmd = [
                    "ffmpeg", "-y", "-i", str(stem_file),
                    "-f", "segment", "-segment_time", "50",
                    "-segment_start_number", "1"
                ]
                
                if ext == "mp3":
                    cmd.extend(["-c:a", "libmp3lame", "-b:a", "320k"])
                
                cmd.append(out_pattern)
                subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            else:
                # No chunking, just convert/copy
                ext = output_format.lower()
                out_file = stem_folder / f"{stem.lower()}.{ext}"
                
                if ext == "wav":
                    shutil.copy(str(stem_file), str(out_file))
                else:
                    cmd = [
                        "ffmpeg", "-y", "-i", str(stem_file),
                        "-c:a", "libmp3lame", "-b:a", "320k",
                        str(out_file)
                    ]
                    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                    
        # Handle subtitles and transcripts
        if chunked:
            import json
            import csv
            transcript_folder = base_dir / "Transcripts"
            transcript_folder.mkdir(exist_ok=True)
            
            def chunk_transcript(src_json, prefix):
                if not src_json.exists():
                    return
                with open(src_json, "r", encoding="utf-8") as f:
                    data = json.load(f)
                
                words = data.get("words", data.get("alignment", []))
                if not words:
                    return
                
                chunks = {}
                for w in words:
                    chunk_idx = int(w["start"] // 50) + 1
                    if chunk_idx not in chunks:
                        chunks[chunk_idx] = []
                    
                    offset = (chunk_idx - 1) * 50
                    chunks[chunk_idx].append({
                        "word": w["word"],
                        "start": round(max(0, w["start"] - offset), 3),
                        "end": round(w["end"] - offset, 3)
                    })
                    
                for idx, chunk_words in chunks.items():
                    out_json = transcript_folder / f"{prefix}_ad{idx:04d}.json"
                    with open(out_json, "w", encoding="utf-8") as f:
                        json.dump({"words": chunk_words}, f, indent=2)
                        
                    out_csv = transcript_folder / f"{prefix}_ad{idx:04d}.csv"
                    with open(out_csv, "w", newline="", encoding="utf-8") as f:
                        writer = csv.writer(f)
                        writer.writerow(["word", "start", "end"])
                        for cw in chunk_words:
                            writer.writerow([cw["word"], cw["start"], cw["end"]])

            chunk_transcript(final_stems_dir / "transcript.json", "whisper")
            chunk_transcript(final_stems_dir / "aligned_timestamps.json", "aligned")
            
            # Also copy the global .srt files to Transcripts folder just in case
            for text_file in ["lyrics.srt"]:
                src_file = final_stems_dir / text_file
                if src_file.exists():
                    shutil.copy(str(src_file), str(transcript_folder / text_file))
            
        else:
            # Copy subtitles and transcripts if they exist normally
            transcript_folder = base_dir / "Transcripts"
            transcript_folder.mkdir(exist_ok=True)
            for text_file in ["lyrics.srt", "transcript.json", "transcript.csv", "aligned_timestamps.json", "aligned_timestamps.csv"]:
                src_file = final_stems_dir / text_file
                if src_file.exists():
                    shutil.copy(str(src_file), str(transcript_folder / text_file))
                    
        # Also copy the dataset folder for ASR training if it exists (always included)
        transcript_folder = base_dir / "Transcripts"
        src_dataset = final_stems_dir / "dataset"
        if src_dataset.exists():
            shutil.copytree(str(src_dataset), str(transcript_folder / "dataset"), dirs_exist_ok=True)
                    
        # Zip everything up
        zip_path = task_dir / f"custom_download_{timestamp}.zip"
        with zipfile.ZipFile(str(zip_path), 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(str(staging_dir)):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, str(staging_dir))
                    zipf.write(file_path, arcname=arcname)
                    
        return str(zip_path)
    finally:
        # Clean up the staging directory to save space
        if staging_dir.exists():
            shutil.rmtree(str(staging_dir), ignore_errors=True)