"""
Hyper-RVC Main Processing Module
Enhanced with Advanced Progress Tracking System
"""
import spaces
import torch
import argparse
import gc
import hashlib
import json
import os
import shlex
import subprocess
import time
from contextlib import suppress
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Optional, Callable, List, Tuple, Any
from urllib.parse import urlparse, parse_qs
import shutil

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import sox
import yt_dlp
from pedalboard import Pedalboard, Reverb, Compressor, HighpassFilter
from pedalboard.io import AudioFile
from pydub import AudioSegment
import noisereduce as nr

from mdx import run_mdx
from rvc import Config, load_hubert, get_vc, rvc_infer
import logging

logging.getLogger("httpx").setLevel(logging.WARNING)

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
IS_ZERO_GPU = os.getenv("SPACES_ZERO_GPU")

mdxnet_models_dir = os.path.join(BASE_DIR, 'mdxnet_models')
rvc_models_dir = os.path.join(BASE_DIR, 'rvc_models')
output_dir = os.path.join(BASE_DIR, 'song_output')


# =============================================================================
# ENHANCED PROGRESS TRACKING SYSTEM
# =============================================================================

@dataclass
class ProgressStep:
    """Represents a single step in the processing pipeline."""
    name: str
    weight: float = 1.0  # Relative weight for progress calculation
    status: str = "pending"  # pending, in_progress, completed, error
    start_time: Optional[float] = None
    end_time: Optional[float] = None
    detail: str = ""


class EnhancedProgressTracker:
    """
    Advanced progress tracker with:
    - Step-by-step progress visualization
    - Time estimation (ETA)
    - Weighted progress calculation
    - Detailed status messages
    - Smooth progress updates
    """
    
    def __init__(self, progress: gr.Progress = None, is_webui: bool = True):
        self.progress = progress
        self.is_webui = is_webui
        self.steps: List[ProgressStep] = []
        self.current_step_index: int = 0
        self.total_weight: float = 0.0
        self.completed_weight: float = 0.0
        self.start_time: float = time.time()
        self.last_update_time: float = 0.0
        self.min_update_interval: float = 0.1  # Minimum seconds between UI updates
        
    def add_step(self, name: str, weight: float = 1.0) -> 'EnhancedProgressTracker':
        """Add a processing step."""
        step = ProgressStep(name=name, weight=weight)
        self.steps.append(step)
        self.total_weight += weight
        return self
    
    def set_steps(self, steps: List[Tuple[str, float]]) -> 'EnhancedProgressTracker':
        """Set multiple steps at once. Format: [(name, weight), ...]"""
        self.steps = []
        self.total_weight = 0.0
        for name, weight in steps:
            self.add_step(name, weight)
        return self
    
    def start(self, message: str = None):
        """Initialize and start tracking."""
        self.start_time = time.time()
        self.current_step_index = 0
        self.completed_weight = 0.0
        
        if message:
            self._update_progress(0, message)
        
        if self.steps:
            self._start_step(0)
    
    def _start_step(self, index: int):
        """Mark a step as started."""
        if index < len(self.steps):
            self.steps[index].status = "in_progress"
            self.steps[index].start_time = time.time()
            self.current_step_index = index
    
    def _complete_step(self, index: int):
        """Mark a step as completed."""
        if index < len(self.steps):
            self.steps[index].status = "completed"
            self.steps[index].end_time = time.time()
            self.completed_weight += self.steps[index].weight
    
    def _calculate_progress(self) -> float:
        """Calculate overall progress (0-100)."""
        if self.total_weight == 0:
            return 0.0
        
        # Completed weight contributes fully
        progress = (self.completed_weight / self.total_weight) * 100
        
        # Current step contributes partially based on sub-progress
        if self.current_step_index < len(self.steps):
            current_step = self.steps[self.current_step_index]
            if current_step.status == "in_progress":
                # Assume current step is 50% done for smooth progression
                step_contribution = (current_step.weight * 0.5 / self.total_weight) * 100
                progress += step_contribution
        
        return min(progress, 99.9)  # Cap at 99.9 until fully complete
    
    def _estimate_remaining_time(self) -> str:
        """Estimate time remaining based on progress."""
        elapsed = time.time() - self.start_time
        current_progress = self._calculate_progress()
        
        if current_progress > 0.1:
            total_estimated = (elapsed / current_progress) * 100
            remaining = max(0, total_estimated - elapsed)
            return str(timedelta(seconds=int(remaining)))
        return "Calculating..."
    
    def _get_status_message(self, custom_message: str = None) -> str:
        """Generate detailed status message."""
        if custom_message:
            return custom_message
        
        if self.current_step_index < len(self.steps):
            step = self.steps[self.current_step_index]
            eta = self._estimate_remaining_time()
            return f"[{self.current_step_index + 1}/{len(self.steps)}] {step.name} | ETA: {eta}"
        
        return "Processing..."
    
    def _update_progress(self, percent: float, message: str = None):
        """Update the Gradio progress bar."""
        current_time = time.time()
        
        # Throttle updates to prevent UI flooding
        if current_time - self.last_update_time < self.min_update_interval and percent < 100:
            return
        
        self.last_update_time = current_time
        
        if self.is_webui and self.progress:
            status_msg = self._get_status_message(message)
            self.progress(percent / 100.0, desc=status_msg)
        elif not self.is_webui:
            print(f"[{percent:.1f}%] {self._get_status_message(message)}")
    
    def update_step(self, step_name: str = None, detail: str = "", sub_progress: float = None):
        """
        Update progress within current or specified step.
        
        Args:
            step_name: Name of step to update (uses current if None)
            detail: Additional detail about current operation
            sub_progress: Sub-progress within current step (0-1)
        """
        # Find step index
        if step_name:
            for i, step in enumerate(self.steps):
                if step.name == step_name:
                    if step.status == "pending":
                        self._complete_step(self.current_step_index)
                        self._start_step(i)
                    break
        
        # Update detail
        if self.current_step_index < len(self.steps):
            self.steps[self.current_step_index].detail = detail
        
        # Calculate and display progress
        overall_progress = self._calculate_progress()
        
        # Adjust for sub-progress if provided
        if sub_progress is not None and self.current_step_index < len(self.steps):
            step = self.steps[self.current_step_index]
            step_progress = (step.weight * sub_progress / self.total_weight) * 100
            overall_progress = ((self.completed_weight + step.weight * sub_progress) / self.total_weight) * 100
        
        message = f"{detail}" if detail else None
        self._update_progress(min(overall_progress, 99.9), message)
    
    def next_step(self, detail: str = ""):
        """Complete current step and move to next."""
        if self.current_step_index < len(self.steps):
            self._complete_step(self.current_step_index)
            
            next_idx = self.current_step_index + 1
            if next_idx < len(self.steps):
                self._start_step(next_idx)
                
                # Update progress for new step
                overall_progress = self._calculate_progress()
                message = f"{detail}" if detail else None
                self._update_progress(overall_progress, message)
    
    def complete(self, final_message: str = "Complete!"):
        """Mark all processing as complete."""
        # Complete any remaining in-progress step
        if self.current_step_index < len(self.steps):
            self._complete_step(self.current_step_index)
        
        elapsed = time.time() - self.start_time
        completion_msg = f"{final_message} | Total time: {timedelta(seconds=int(elapsed))}"
        self._update_progress(100, completion_msg)
    
    def error(self, error_message: str):
        """Mark current step as errored."""
        if self.current_step_index < len(self.steps):
            self.steps[self.current_step_index].status = "error"
            self.steps[self.current_step_index].detail = error_message
        
        self._update_progress(self._calculate_progress(), f"ERROR: {error_message}")
    
    def get_progress_info(self) -> dict:
        """Get current progress information as dictionary."""
        return {
            "current_step": self.current_step_index,
            "total_steps": len(self.steps),
            "progress_percent": self._calculate_progress(),
            "eta": self._estimate_remaining_time(),
            "elapsed": str(timedelta(seconds=int(time.time() - self.start_time))),
            "steps": [
                {"name": s.name, "status": s.status, "detail": s.detail}
                for s in self.steps
            ]
        }


def create_progress_tracker(progress: gr.Progress = None, is_webui: bool = True) -> EnhancedProgressTracker:
    """
    Factory function to create a pre-configured progress tracker 
    for the song cover pipeline.
    """
    tracker = EnhancedProgressTracker(progress=progress, is_webui=is_webui)
    
    # Define pipeline steps with weights reflecting relative processing time
    tracker.set_steps([
        ("Initializing Pipeline", 1),
        ("Downloading/Loading Audio", 3),
        ("Converting to Stereo", 1),
        ("Separating Vocals from Instrumental", 15),
        ("Separating Main Vocals from Backup", 12),
        ("Applying DeReverb to Vocals", 10),
        ("Extracting Voiceless Track", 8),
        ("Converting Voice with RVC", 25),
        ("Applying Audio Effects", 8),
        ("Applying Pitch Shift", 5),
        ("Combining Audio Tracks", 7),
        ("Finalizing Output", 5),
    ])
    
    return tracker


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def clean_old_folders(base_path: str, max_age_seconds: int = 10800):
    """Clean up old output folders to save disk space."""
    if not os.path.isdir(base_path):
        logging.warning(f"Error: {base_path} is not a valid directory.")
        return

    now = time.time()
    cleaned_count = 0

    for folder_name in os.listdir(base_path):
        folder_path = os.path.join(base_path, folder_name)
        if os.path.isdir(folder_path):
            last_modified = os.path.getmtime(folder_path)
            if now - last_modified > max_age_seconds:
                try:
                    shutil.rmtree(folder_path)
                    cleaned_count += 1
                except Exception as e:
                    logging.warning(f"Failed to delete {folder_path}: {e}")
    
    if cleaned_count > 0:
        logging.info(f"Cleaned up {cleaned_count} old folders from {base_path}")


def get_youtube_video_id(url, ignore_playlist=True):
    """
    Extract video ID from various YouTube URL formats.
    
    Examples:
    - http://youtu.be/SA2iWivDJiE
    - http://www.youtube.com/watch?v=_oPAwA_Udwc&feature=feedu
    - http://www.youtube.com/embed/SA2iWivDJiE
    """
    if "m.youtube.com" in url:
        url = url.replace("m.youtube.com", "www.youtube.com")
    
    query = urlparse(url)
    
    if query.hostname == 'youtu.be':
        if query.path[1:] == 'watch':
            return query.query[2:]
        return query.path[1:]

    if query.hostname in {'www.youtube.com', 'youtube.com', 'music.youtube.com'}:
        if not ignore_playlist:
            with suppress(KeyError):
                return parse_qs(query.query)['list'][0]
        if query.path == '/watch':
            return parse_qs(query.query)['v'][0]
        if query.path[:7] == '/watch/':
            return query.path.split('/')[1]
        if query.path[:7] == '/embed/':
            return query.path.split('/')[2]
        if query.path[:3] == '/v/':
            return query.path.split('/')[1]

    return None


def yt_download(link, progress_callback=None):
    """
    Download audio from YouTube URL with progress callback support.
    
    Args:
        link: YouTube video URL
        progress_callback: Optional function(status, percent) for progress updates
    
    Returns:
        Path to downloaded audio file
    """
    if not link.strip():
        raise ValueError("You need to provide a download link.")
    
    def progress_hook(d):
        if progress_callback and d['status'] == 'downloading':
            if 'downloaded_bytes' in d and 'total_bytes' in d and d['total_bytes'] > 0:
                percent = (d['downloaded_bytes'] / d['total_bytes']) * 100
                progress_callback(f"Downloading: {percent:.1f}%", min(percent, 90))
    
    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': '%(title)s',
        'nocheckcertificate': True,
        'ignoreerrors': True,
        'no_warnings': True,
        'quiet': True,
        'extractaudio': True,
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'mp3',
            'preferredquality': '192',
        }],
        'progress_hooks': [progress_hook] if progress_callback else [],
    }
    
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        result = ydl.extract_info(link, download=True)
        if result is None:
            raise Exception("Failed to extract video information")
        download_path = ydl.prepare_filename(result, outtmpl='%(title)s.mp3')

    return download_path


def raise_exception(error_msg, is_webui):
    """Raise appropriate exception based on context."""
    if is_webui:
        raise gr.Error(error_msg)
    else:
        raise Exception(error_msg)


def get_rvc_model(voice_model, is_webui):
    """
    Locate RVC model files (.pth and .index) in model directory.
    
    Returns:
        Tuple of (model_path, index_path)
    """
    rvc_model_filename, rvc_index_filename = None, None
    model_dir = os.path.join(rvc_models_dir, voice_model)
    
    if not os.path.isdir(model_dir):
        error_msg = f'Model directory does not exist: {model_dir}'
        raise_exception(error_msg, is_webui)
    
    for file in os.listdir(model_dir):
        file_path = os.path.join(model_dir, file)
        
        # Handle nested directories
        if os.path.isdir(file_path):
            for ff in os.listdir(file_path):
                ext = os.path.splitext(ff)[1]
                if ext == '.pth':
                    rvc_model_filename = ff
                elif ext == '.index':
                    rvc_index_filename = ff
        else:
            ext = os.path.splitext(file)[1]
            if ext == '.pth':
                rvc_model_filename = file
            elif ext == '.index':
                rvc_index_filename = file

    if rvc_model_filename is None:
        error_msg = f'No model file (.pth) found in {model_dir}.'
        raise_exception(error_msg, is_webui)

    model_path = os.path.join(model_dir, rvc_model_filename)
    index_path = os.path.join(model_dir, rvc_index_filename) if rvc_index_filename else ''
    
    return model_path, index_path


def get_audio_paths(song_dir):
    """Extract various audio paths from processed song directory."""
    orig_song_path = None
    instrumentals_path = None
    main_vocals_dereverb_path = None
    backup_vocals_path = None

    for file in os.listdir(song_dir):
        if file.endswith('_Instrumental.wav'):
            instrumentals_path = os.path.join(song_dir, file)
            orig_song_path = instrumentals_path.replace('_Instrumental', '')

        elif file.endswith('_Vocals_Main_DeReVerb.wav'):
            main_vocals_dereverb_path = os.path.join(song_dir, file)

        elif file.endswith('_Vocals_Backup.wav'):
            backup_vocals_path = os.path.join(song_dir, file)

    return orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path


def get_audio_with_suffix(song_dir, suffix="_mysuffix.wav"):
    """Find audio file with specific suffix in directory."""
    for file in os.listdir(song_dir):
        if file.endswith(suffix):
            return os.path.join(song_dir, file)
    return None


def convert_to_stereo(audio_path):
    """
    Convert mono audio to stereo if needed.
    FIXED: Added file existence validation and better error handling
    """
    # Validate file exists first
    if not audio_path or not os.path.exists(audio_path):
        raise FileNotFoundError(f"Audio file not found for stereo conversion: {audio_path}")
    
    try:
        # Try loading with soundfile first (more reliable)
        try:
            data, sr = sf.read(audio_path)
            if len(data.shape) == 1 or data.shape[1] == 1:
                # Mono file - need conversion
                stereo_path = f'{os.path.splitext(audio_path)[0]}_stereo.wav'
                command = shlex.split(f'ffmpeg -y -loglevel error -i "{audio_path}" -ac 2 -f wav "{stereo_path}"')
                result = subprocess.run(command, capture_output=True, text=True, timeout=60)
                if result.returncode != 0:
                    logging.error(f"FFmpeg stereo conversion failed: {result.stderr}")
                    return audio_path  # Return original if conversion fails
                return stereo_path
            else:
                return audio_path  # Already stereo
        except Exception as sf_error:
            logging.debug(f"soundfile failed, trying librosa: {sf_error}")
            
            # Fallback to librosa
            wave, sr_librosa = librosa.load(audio_path, mono=False, sr=44100)

            # Check if mono
            if type(wave[0]) != np.ndarray:
                stereo_path = f'{os.path.splitext(audio_path)[0]}_stereo.wav'
                command = shlex.split(f'ffmpeg -y -loglevel error -i "{audio_path}" -ac 2 -f wav "{stereo_path}"')
                result = subprocess.run(command, capture_output=True, text=True, timeout=60)
                if result.returncode != 0:
                    logging.error(f"FFmpeg stereo conversion failed: {result.stderr}")
                    return audio_path
                return stereo_path
            else:
                return audio_path
                
    except FileNotFoundError:
        raise  # Re-raise FileNotFoundError with original message
    except Exception as e:
        logging.warning(f"Stereo conversion warning for {audio_path}: {e}")
        # Return original path if conversion fails completely
        return audio_path


def pitch_shift(audio_path, pitch_change):
    """Apply pitch shift to audio file."""
    output_path = f'{os.path.splitext(audio_path)[0]}_p{pitch_change}.wav'
    if not os.path.exists(output_path):
        y, sr = sf.read(audio_path)
        tfm = sox.Transformer()
        tfm.pitch(pitch_change)
        y_shifted = tfm.build_array(input_array=y, sample_rate_in=sr)
        sf.write(output_path, y_shifted, sr)

    return output_path


def get_hash(filepath):
    """Generate short hash of file for identification."""
    with open(filepath, 'rb') as f:
        file_hash = hashlib.blake2b()
        while chunk := f.read(8192):
            file_hash.update(chunk)
    return file_hash.hexdigest()[:11]


# =============================================================================
# AUDIO PROCESSING FUNCTIONS
# =============================================================================

def add_audio_effects(audio_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping):
    """Apply professional audio effects (HPF, Compressor, Reverb)."""
    output_path = f'{os.path.splitext(audio_path)[0]}_mixed.wav'

    board = Pedalboard(
        [
            HighpassFilter(cutoff_frequency=80),
            Compressor(ratio=4, threshold_db=-15),
            Reverb(room_size=reverb_rm_size, dry_level=reverb_dry, wet_level=reverb_wet, damping=reverb_damping)
         ]
    )

    with AudioFile(audio_path) as f:
        with AudioFile(output_path, 'w', f.samplerate, f.num_channels) as o:
            while f.tell() < f.frames:
                chunk = f.read(int(f.samplerate))
                effected = board(chunk, f.samplerate, reset=False)
                o.write(effected)

    return output_path


def combine_audio(audio_paths, output_path, main_gain, backup_gain, inst_gain, output_format):
    """Combine multiple audio tracks with gain adjustments."""
    try:
        main_vocal_audio = AudioSegment.from_wav(audio_paths[0]) - 4 + main_gain
        backup_vocal_audio = AudioSegment.from_wav(audio_paths[1]) - 6 + backup_gain
        instrumental_audio = AudioSegment.from_wav(audio_paths[2]) - 7 + inst_gain
        
        combined = main_vocal_audio.overlay(backup_vocal_audio).overlay(instrumental_audio)
        combined.export(output_path, format=output_format)
    except Exception as e:
        logging.error(f"Audio combining error: {e}")
        raise


def apply_noisereduce(audio_list, type_output="wav"):
    """Apply noise reduction to audio files."""
    result = []
    for audio_path in audio_list:
        out_path = f"{os.path.splitext(audio_path)[0]}_nr.{type_output}"

        try:
            audio = AudioSegment.from_file(audio_path)
            samples = np.array(audio.get_array_of_samples())
            
            reduced_noise = nr.reduce_noise(
                samples, 
                sr=audio.frame_rate, 
                prop_decrease=0.6,
                n_std_thresh_stationary=1.5
            )

            reduced_audio = AudioSegment(
                reduced_noise.tobytes(), 
                frame_rate=audio.frame_rate, 
                sample_width=audio.sample_width,
                channels=audio.channels
            )

            reduced_audio.export(out_path, format=type_output)
            result.append(out_path)

        except Exception as e:
            logging.warning(f"Noise reduction error for {audio_path}: {e}")
            result.append(audio_path)

    return result


# =============================================================================
# MAIN PROCESSING PIPELINE
# =============================================================================

def preprocess_song(mdx_model_params, song_output_dir, orig_song_path, keep_orig=False, 
                   tracker: EnhancedProgressTracker = None):
    """
    Preprocess song: separate vocals/instrumentals, apply dereverb.
    
    Returns:
        Tuple of all processed audio paths
    """
    # Step: Vocal Separation
    if tracker:
        tracker.update_step(detail="Running MDX vocal separation...")
    
    vocals_path, instrumentals_path = run_mdx(
        mdx_model_params, 
        song_output_dir, 
        os.path.join(mdxnet_models_dir, 'UVR-MDX-NET-Voc_FT.onnx'), 
        orig_song_path, 
        denoise=True, 
        keep_orig=keep_orig
    )
    
    if tracker:
        tracker.next_step(detail="Separating main from backup vocals...")
    
    # Step: Main/Backup Vocal Separation
    backup_vocals_path, main_vocals_path = run_mdx(
        mdx_model_params, 
        song_output_dir, 
        os.path.join(mdxnet_models_dir, 'UVR_MDXNET_KARA_2.onnx'), 
        vocals_path, 
        suffix='Backup', 
        invert_suffix='Main', 
        denoise=True
    )
    
    if tracker:
        tracker.next_step(detail="Applying dereverb processing...")
    
    # Step: Dereverb
    _, main_vocals_dereverb_path = run_mdx(
        mdx_model_params, 
        song_output_dir, 
        os.path.join(mdxnet_models_dir, 'Reverb_HQ_By_FoxJoy.onnx'), 
        main_vocals_path, 
        invert_suffix='DeReverb', 
        exclude_main=True, 
        denoise=True
    )
    
    return orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path


@spaces.GPU(duration=65)
def voice_change(voice_model, vocals_path, output_path, pitch_change, f0_method, 
                 index_rate, filter_radius, rms_mix_rate, protect, crepe_hop_length, 
                 is_webui, steps, tracker: EnhancedProgressTracker = None):
    """
    Perform RVC voice conversion on vocals.
    """
    if tracker:
        tracker.update_step(detail="Loading RVC model...")
    
    rvc_model_path, rvc_index_path = get_rvc_model(voice_model, is_webui)

    cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)
    
    if tracker:
        tracker.update_step(detail="Running voice conversion (this may take a while)...")

    global hubert_model
    rvc_infer(
        rvc_index_path, index_rate, vocals_path, output_path, pitch_change, 
        f0_method, cpt, version, net_g, filter_radius, tgt_sr, rms_mix_rate, 
        protect, crepe_hop_length, vc, hubert_model, steps
    )
    
    del cpt
    gc.collect()
    torch.cuda.empty_cache() if torch.cuda.is_available() else None


@spaces.GPU(duration=65)
def process_song(
    song_dir, mdx_model_params, song_id, is_webui, input_type,
    keep_files, pitch_change, pitch_change_all, voice_model, index_rate, 
    filter_radius, rms_mix_rate, protect, f0_method, crepe_hop_length, 
    output_format, keep_orig, orig_song_path, steps,
    tracker: EnhancedProgressTracker = None
):
    """
    Process song through the full pipeline.
    """
    if not os.path.exists(song_dir):
        os.makedirs(song_dir)
        orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path = preprocess_song(
            mdx_model_params, song_dir, orig_song_path, keep_orig, tracker
        )
    else:
        vocals_path, main_vocals_path = None, None
        paths = get_audio_paths(song_dir)

        if any(path is None for path in paths):
            orig_song_path, vocals_path, instrumentals_path, main_vocals_path, backup_vocals_path, main_vocals_dereverb_path = preprocess_song(
                mdx_model_params, song_dir, orig_song_path, keep_orig, tracker
            )
        else:
            orig_song_path, instrumentals_path, main_vocals_dereverb_path, backup_vocals_path = paths

    # Voiceless track extraction
    ins_path = get_audio_with_suffix(song_dir, "_Voiceless.wav")
    if not ins_path:
        if tracker:
            tracker.next_step(detail="Extracting voiceless instrumental track...")
        
        instrumentals_path, _ = run_mdx(
            mdx_model_params,
            song_dir,
            os.path.join(mdxnet_models_dir, "UVR-MDX-NET-Inst_HQ_4.onnx"),
            instrumentals_path,
            exclude_inversion=True,
            suffix="Voiceless",
            denoise=False,
            keep_orig=True,
            base_device=("cuda" if IS_ZERO_GPU else "")
        )
        ins_path = get_audio_with_suffix(song_dir, "_Voiceless.wav")
    else:
        instrumentals_path = ins_path

    # Prepare output paths
    pitch_change_total = pitch_change * 12 + pitch_change_all
    ai_vocals_path = os.path.join(
        song_dir, 
        f'{os.path.splitext(os.path.basename(orig_song_path))[0]}_{voice_model}_p{pitch_change_total}_i{index_rate}_fr{filter_radius}_rms{rms_mix_rate}_pro{protect}_{f0_method}{"" if f0_method != "mangio-crepe" else f"_{crepe_hop_length}"}_s{steps}.wav'
    )
    ai_cover_path = os.path.join(
        song_dir, 
        f'{os.path.splitext(os.path.basename(orig_song_path))[0]} ({voice_model} Ver).{output_format}'
    )

    # RVC Voice Conversion
    if not os.path.exists(ai_vocals_path):
        if tracker:
            tracker.next_step()
        voice_change(
            voice_model, main_vocals_dereverb_path, ai_vocals_path, pitch_change_total, 
            f0_method, index_rate, filter_radius, rms_mix_rate, protect, 
            crepe_hop_length, is_webui, steps, tracker
        )

    return ai_vocals_path, ai_cover_path, instrumentals_path, backup_vocals_path, vocals_path, main_vocals_path, ins_path


@spaces.GPU(duration=140)
def song_cover_pipeline(
    song_input, voice_model, pitch_change, keep_files,
    is_webui=0, main_gain=0, backup_gain=0, inst_gain=0, index_rate=0.5, 
    filter_radius=3, rms_mix_rate=0.25, f0_method='rmvpe', crepe_hop_length=128, 
    protect=0.33, pitch_change_all=0, reverb_rm_size=0.15, reverb_wet=0.2, 
    reverb_dry=0.8, reverb_damping=0.7, output_format='mp3', extra_denoise=False, 
    steps=1, progress=gr.Progress()
):
    """
    Main pipeline for generating AI song covers with enhanced progress tracking.
    """
    # Initialize enhanced progress tracker
    tracker = create_progress_tracker(progress=progress, is_webui=bool(is_webui))
    
    # Cleanup old folders
    if not keep_files or IS_ZERO_GPU:
        clean_old_folders("./song_output", 14400)

    if IS_ZERO_GPU:
        clean_old_folders("./rvc_models", 10800)

    try:
        # Input validation - with detailed error messages
        if not voice_model:
            raise_exception('Please select a voice model first!', is_webui)
        
        # Validate voice model exists
        model_dir = os.path.join(rvc_models_dir, voice_model)
        if not os.path.isdir(model_dir):
            raise_exception(f'Voice model not found: "{voice_model}". Please download or select a valid model.', is_webui)
        
        if song_input is None or (isinstance(song_input, str) and len(song_input.strip()) == 0):
            raise_exception('Please provide an audio file (select from folder, upload, or paste URL).', is_webui)
        
        # Clean up input path
        if isinstance(song_input, str):
            song_input = song_input.strip().strip('\"').strip("'")
        
        # CRITICAL: Validate file exists BEFORE processing
        if not urlparse(song_input).scheme:  # Not a URL
            if not os.path.exists(song_input):
                raise_exception(
                    f'Audio file not found: "{song_input}"\n\n'
                    'The file may have been deleted or moved. Please:\n'
                    '1. Re-select the audio from the folder dropdown\n'
                    '2. Re-upload the file\n'
                    '3. Check that the file path is correct', 
                    is_webui
                )
            
            # Verify file is readable and has content
            try:
                file_size = os.path.getsize(song_input)
                if file_size == 0:
                    raise_exception(f'Audio file is empty (0 bytes): {song_input}', is_webui)
                logging.info(f"Audio file validated: {song_input} ({file_size/1024/1024:.2f} MB)")
            except OSError as e:
                raise_exception(f'Cannot read audio file: {song_input}\nError: {e}', is_webui)
        
        logging.info(f"Starting pipeline - Input: {song_input}, Model: {voice_model}")
        
        # Start progress tracking
        tracker.start("Initializing Hyper-RVC pipeline...")
        tracker.next_step(detail="Validating inputs...")

        # Load MDX model parameters
        with open(os.path.join(mdxnet_models_dir, 'model_data.json')) as infile:
            mdx_model_params = json.load(infile)

        # Determine input type
        if urlparse(song_input).scheme == 'https':
            input_type = 'yt'
            song_id = get_youtube_video_id(song_input)
            if song_id is None:
                raise_exception('Invalid YouTube URL.', is_webui)
        else:
            input_type = 'local'
            song_input = song_input.strip('"')
            if os.path.exists(song_input):
                song_id = get_hash(song_input)
            else:
                error_msg = f'File not found: {song_input}'
                raise_exception(error_msg, is_webui)
                song_id = None

        song_dir = os.path.join(output_dir, song_id)

        # Download/Load audio
        if input_type == 'yt':
            def yt_progress(status, percent):
                tracker.update_step(detail=status)
            
            tracker.next_step(detail="Downloading from YouTube...")
            orig_song_path = yt_download(song_input, progress_callback=yt_progress)
            keep_orig = False
        else:
            orig_song_path = song_input
            keep_orig = True
            tracker.next_step(detail="Loading local audio file...")

        # Convert to stereo
        tracker.next_step(detail="Converting to stereo if needed...")
        orig_song_path = convert_to_stereo(orig_song_path)

        start_time = time.time()

        # Run main processing pipeline
        (
            ai_vocals_path,
            ai_cover_path,
            instrumentals_path,
            backup_vocals_path,
            vocals_path,
            main_vocals_path,
            ins_path
        ) = process_song(
            song_dir,
            mdx_model_params,
            song_id,
            is_webui,
            input_type,
            keep_files,
            pitch_change,
            pitch_change_all,
            voice_model,
            index_rate,
            filter_radius,
            rms_mix_rate,
            protect,
            f0_method,
            crepe_hop_length,
            output_format,
            keep_orig,
            orig_song_path,
            steps,
            tracker
        )

        end_time = time.time()
        logging.info(f"Pipeline execution time: {end_time - start_time:.2f}s")

        # Apply audio effects
        tracker.next_step(detail="Applying reverb and effects...")

        if extra_denoise:
            tracker.update_step(detail="Applying noise reduction...")
            ai_vocals_path = apply_noisereduce([ai_vocals_path])[0]

        ai_vocals_mixed_path = add_audio_effects(
            ai_vocals_path, reverb_rm_size, reverb_wet, reverb_dry, reverb_damping
        )
        
        # Pitch shift for instrumentals if needed
        if pitch_change_all != 0:
            tracker.next_step(detail="Applying overall pitch change...")
            instrumentals_path = pitch_shift(instrumentals_path, pitch_change_all)
            backup_vocals_path = pitch_shift(backup_vocals_path, pitch_change_all)

        # Combine tracks
        tracker.next_step(detail="Mixing final audio tracks...")
        combine_audio(
            [ai_vocals_mixed_path, backup_vocals_path, instrumentals_path], 
            ai_cover_path, main_gain, backup_gain, inst_gain, output_format
        )

        # Cleanup intermediate files
        if not keep_files:
            tracker.next_step(detail="Cleaning up temporary files...")
            intermediate_files = [vocals_path, main_vocals_path, ai_vocals_mixed_path]
            
            if extra_denoise and ai_vocals_path.endswith('_nr.wav'):
                intermediate_files.append(ai_vocals_path)
            if pitch_change_all != 0:
                intermediate_files += [instrumentals_path, backup_vocals_path]
                
            for file in intermediate_files:
                if file and os.path.exists(file) and file != ins_path:
                    try:
                        os.remove(file)
                    except Exception as e:
                        logging.debug(f"Cleanup: Could not remove {file}")

        # Complete!
        tracker.complete("AI Cover generated successfully!")
        
        return ai_cover_path

    except Exception as e:
        tracker.error(str(e))
        raise_exception(str(e), is_webui)


# =============================================================================
# BATCH PROCESSING
# =============================================================================

@spaces.GPU(duration=120)
def batch_process_files(
    input_files, voice_model, pitch_change, f0_method, index_rate, protect, 
    output_format, progress=gr.Progress()
):
    """
    Process multiple audio files with the same RVC model settings.
    Features enhanced progress tracking with per-file progress.
    """
    if not input_files:
        raise gr.Error("No files provided for batch processing!")
    
    if not voice_model:
        raise gr.Error("Please select a voice model!")
    
    # Initialize batch progress tracker
    tracker = EnhancedProgressTracker(progress=progress, is_webui=True)
    tracker.start("Starting batch processing...")
    
    results = []
    total_files = len(input_files)
    batch_output_dir = os.path.join(output_dir, 'batch_output')
    os.makedirs(batch_output_dir, exist_ok=True)
    
    # Load RVC model once for efficiency
    tracker.update_step(detail="Loading RVC model (shared across all files)...")
    rvc_model_path, rvc_index_path = get_rvc_model(voice_model, True)
    cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)
    
    global hubert_model
    
    for i, file_info in enumerate(input_files):
        try:
            # Get file path
            if hasattr(file_info, 'name'):
                input_path = file_info.name
            elif isinstance(file_info, str):
                input_path = file_info
            else:
                continue
            
            filename = os.path.basename(input_path)
            base_name = os.path.splitext(filename)[0]
            
            # Update progress with file info
            file_progress = ((i) / total_files) * 100
            tracker.update_step(
                detail=f"Processing file {i+1}/{total_files}: {filename}",
                sub_progress=0
            )
            
            # Load and convert audio
            audio, sr = librosa.load(input_path, sr=16000, mono=True)
            
            temp_input = os.path.join(batch_output_dir, f'temp_{base_name}.wav')
            sf.write(temp_input, audio, 16000)
            
            output_path = os.path.join(batch_output_dir, f'{base_name}_{voice_model}_converted.{output_format}')
            
            # Update progress for voice conversion
            tracker.update_step(
                detail=f"[{i+1}/{total_files}] Converting voice: {filename}",
                sub_progress=0.5
            )
            
            # Run voice conversion
            rvc_infer(
                rvc_index_path, index_rate, temp_input, 
                output_path.replace(f'.{output_format}', '.wav'),
                pitch_change * 12, f0_method, cpt, version, net_g,
                3, tgt_sr, 0.25, protect, 128, vc, hubert_model, 1
            )
            
            # Convert to desired format if needed
            if output_format != 'wav':
                wav_path = output_path.replace(f'.{output_format}', '.wav')
                y, sr_out = sf.read(wav_path)
                sf.write(output_path, y, sr_out)
                os.remove(wav_path)
            else:
                output_path = output_path.replace(f'.{output_format}', '.wav')
            
            results.append(output_path)
            logging.info(f"[+] Batch processed: {filename}")
            
            # Mark file complete
            tracker.update_step(
                detail=f"[✓] Completed: {filename}",
                sub_progress=1.0
            )
            
        except Exception as e:
            logging.error(f"[-] Error processing {getattr(file_info, 'name', 'unknown')}: {str(e)}")
            continue
    
    # Cleanup GPU resources
    del cpt
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    tracker.complete(f"Batch complete! {len(results)}/{total_files} files processed.")
    
    return results


# =============================================================================
# PREVIEW GENERATION
# =============================================================================

@spaces.GPU(duration=30)
def generate_preview(
    audio_path, duration=10, voice_model=None, pitch_change=0, f0_method='rmvpe+', 
    index_rate=0.5, protect=0.33, progress=gr.Progress()
):
    """
    Generate a short preview clip of the voice conversion.
    Useful for testing settings before full conversion.
    """
    tracker = EnhancedProgressTracker(progress=progress, is_webui=True)
    tracker.set_steps([
        ("Loading Audio", 2),
        ("Extracting Preview Segment", 1),
        ("Loading RVC Model", 3),
        ("Converting Preview", 4),
        ("Finalizing", 1)
    ])
    tracker.start()
    
    if not audio_path:
        raise gr.Error("Please provide an audio file for preview!")
    
    tracker.next_step(detail="Reading audio file...")
    
    # Load audio
    audio, sr = librosa.load(audio_path, sr=16000, mono=True)
    
    # Extract preview segment (from middle of audio for best representation)
    max_samples = int(duration * sr)
    start_sample = min(len(audio) // 2 - max_samples // 2, len(audio) - max_samples)
    start_sample = max(0, start_sample)
    preview_audio = audio[start_sample:start_sample + max_samples]
    
    # Save preview input
    preview_dir = os.path.join(output_dir, 'preview')
    os.makedirs(preview_dir, exist_ok=True)
    preview_input_path = os.path.join(preview_dir, 'preview_input.wav')
    sf.write(preview_input_path, preview_audio, 16000)
    
    tracker.next_step(detail="Preview extracted, loading model...")
    
    # If no model specified, return original audio
    if not voice_model:
        tracker.complete("Returning original audio preview.")
        return preview_input_path
    
    # Load model
    rvc_model_path, rvc_index_path = get_rvc_model(voice_model, True)
    cpt, version, net_g, tgt_sr, vc = get_vc(device, config.is_half, config, rvc_model_path)
    
    global hubert_model
    
    tracker.next_step(detail="Generating voice conversion preview...")
    
    # Convert preview
    preview_output_path = os.path.join(preview_dir, 'preview_output.wav')
    rvc_infer(
        rvc_index_path, index_rate, preview_input_path,
        preview_output_path, pitch_change * 12, f0_method, cpt, version, net_g,
        3, tgt_sr, 0.25, protect, 128, vc, hubert_model, 1
    )
    
    del cpt
    gc.collect()
    
    tracker.complete("Preview generated successfully!")
    
    return preview_output_path


# =============================================================================
# GLOBAL INITIALIZATION & CLI
# =============================================================================

device = "cuda:0" if torch.cuda.is_available() else "cpu"
compute_half = True if torch.cuda.is_available() else False
config = Config(device, compute_half)
hubert_model = load_hubert("cuda", config.is_half, None)

print(f"Device: {device}, Half precision: {config.is_half}")
print("RVC models and Hubert loaded successfully.")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Generate an AI cover song in the song_output/id directory.', 
        add_help=True
    )
    parser.add_argument('-i', '--song-input', type=str, required=True, help='Link to a YouTube video or filepath to local mp3/wav')
    parser.add_argument('-dir', '--rvc-dirname', type=str, required=True, help='Name of folder in rvc_models containing RVC model')
    parser.add_argument('-p', '--pitch-change', type=int, required=True, help='Pitch change in octaves (+1 M→F, -1 F→M)')
    parser.add_argument('-k', '--keep-files', action=argparse.BooleanOptionalAction, help='Keep intermediate audio files')
    parser.add_argument('-ir', '--index-rate', type=float, default=0.5, help='Index rate for timbre control (0-1)')
    parser.add_argument('-fr', '--filter-radius', type=int, default=3, help='Filter radius for median filtering (0-7)')
    parser.add_argument('-rms', '--rms-mix-rate', type=float, default=0.25, help='RMS mix rate (0=original loudness, 1=fixed)')
    parser.add_argument('-palgo', '--pitch-detection-algo', type=str, default='rmvpe', help='Pitch detection algorithm')
    parser.add_argument('-hop', '--crepe-hop-length', type=int, default=128, help='Crepe hop length for mangio-crepe')
    parser.add_argument('-pro', '--protect', type=float, default=0.33, help='Protect voiceless consonants (0-0.5)')
    parser.add_argument('-mv', '--main-vol', type=int, default=0, help='Main vocals volume (dB)')
    parser.add_argument('-bv', '--backup-vol', type=int, default=0, help='Backup vocals volume (dB)')
    parser.add_argument('-iv', '--inst-vol', type=int, default=0, help='Instrumentals volume (dB)')
    parser.add_argument('-pall', '--pitch-change-all', type=int, default=0, help='Overall pitch change (semitones)')
    parser.add_argument('-rsize', '--reverb-size', type=float, default=0.15, help='Reverb room size (0-1)')
    parser.add_argument('-rwet', '--reverb-wetness', type=float, default=0.2, help='Reverb wet level (0-1)')
    parser.add_argument('-rdry', '--reverb-dryness', type=float, default=0.8, help='Reverb dry level (0-1)')
    parser.add_argument('-rdamp', '--reverb-damping', type=float, default=0.7, help='Reverb damping (0-1)')
    parser.add_argument('-oformat', '--output-format', type=str, default='mp3', help='Output format (mp3/wav)')
    args = parser.parse_args()

    rvc_dirname = args.rvc_dirname
    if not os.path.exists(os.path.join(rvc_models_dir, rvc_dirname)):
        raise Exception(f'The folder {os.path.join(rvc_models_dir, rvc_dirname)} does not exist.')

    cover_path = song_cover_pipeline(
        args.song_input, rvc_dirname, args.pitch_change, args.keep_files,
        main_gain=args.main_vol, backup_gain=args.backup_vol, inst_gain=args.inst_vol,
        index_rate=args.index_rate, filter_radius=args.filter_radius,
        rms_mix_rate=args.rms_mix_rate, f0_method=args.pitch_detection_algo,
        crepe_hop_length=args.crepe_hop_length, protect=args.protect,
        pitch_change_all=args.pitch_change_all,
        reverb_rm_size=args.reverb_size, reverb_wet=args.reverb_wetness,
        reverb_dry=args.reverb_dryness, reverb_damping=args.reverb_damping,
        output_format=args.output_format
    )
    print(f'[+] Cover generated at {cover_path}')