import os
import folder_paths
import torch
import torchaudio
import tempfile
import shutil
from audio_separator.separator import Separator

# Определяем путь к моделям: ComfyUI/models/UVR
COMFY_MODELS_DIR = folder_paths.models_dir
UVR_MODELS_DIR = os.path.join(COMFY_MODELS_DIR, "UVR")

# Если папки нет - создаем
if not os.path.exists(UVR_MODELS_DIR):
    try:
        os.makedirs(UVR_MODELS_DIR)
    except Exception:
        pass

class AudioSeparatorNode:
    def __init__(self):
        pass
    
    @classmethod
    def INPUT_TYPES(s):
        # Сканируем папку models/UVR на наличие моделей
        model_files = []
        if os.path.exists(UVR_MODELS_DIR):
            model_files = [f for f in os.listdir(UVR_MODELS_DIR) if f.endswith(('.pth', '.onnx', '.bin'))]
        
        # Если папка пустая, добавляем дефолтные имена, чтобы список не был пустым
        if not model_files:
            model_files = ["1_HP-UVR.pth", "UVR-DeNoise.pth", "UVR-MDX-NET-Inst_HQ_3.onnx"]
            
        return {
            "required": {
                "audio": ("AUDIO",),
                "model_filename": (model_files, {"default": model_files[0] if model_files else "1_HP-UVR.pth"}),
                "window_size": ("INT", {"default": 512, "min": 320, "max": 1024, "step": 64, "tooltip": "VR Model Window Size. 512 is High Quality."}),
                "aggression": ("INT", {"default": 10, "min": 1, "max": 50, "step": 1, "tooltip": "VR Model Aggression. Higher = stronger removal."}),
                "batch_size": ("INT", {"default": 4, "min": 1, "max": 16, "step": 1, "tooltip": "Batch size for GPU processing."}),
            },
            "optional": {
                "normalization": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0, "step": 0.1}),
            }
        }

    RETURN_TYPES = ("AUDIO", "AUDIO")
    RETURN_NAMES = ("Vocals (Clean)", "Instrumental (Noise)")
    FUNCTION = "separate_audio"
    CATEGORY = "Audio/Processing"

    def separate_audio(self, audio, model_filename, window_size, aggression, batch_size, normalization):
        # 1. Подготовка аудио из Comfy
        waveform = audio['waveform']
        # Берем sample_rate входного файла только для сохранения времянки
        input_sample_rate = audio['sample_rate']
        
        # Если батч > 1, берем только первый элемент
        if waveform.dim() == 3:
            waveform = waveform[0]
            
        # Создаем временную папку
        temp_dir = tempfile.mkdtemp()
        input_path = os.path.join(temp_dir, "input_temp.wav")
        
        # Сохраняем временный файл
        torchaudio.save(input_path, waveform, input_sample_rate)

        # 2. Инициализация Сепаратора
        print(f"🎧 Audio Separator: Initializing with model '{model_filename}'...")
        
        separator = Separator(
            model_file_dir=UVR_MODELS_DIR,  # Ищем модели в ComfyUI/models/UVR
            output_dir=temp_dir,            # Результат кладем во временную папку
            output_format="WAV",
            normalization_threshold=normalization,
            vr_params={
                "batch_size": batch_size,
                "window_size": window_size,
                "aggression": aggression,
            },
            mdx_params={
                "batch_size": batch_size,
            }
        )

        # Загружаем модель
        try:
            separator.load_model(model_filename=model_filename)
        except Exception as e:
            shutil.rmtree(temp_dir)
            raise ValueError(f"Ошибка загрузки модели '{model_filename}'. Убедитесь, что она лежит в {UVR_MODELS_DIR}. \nОшибка: {e}")

        # 3. Обработка
        print(f"🎧 Audio Separator: Processing on GPU...")
        output_files = separator.separate(input_path)
        
        # Определяем, какой файл есть какой
        vocals_path = os.path.join(temp_dir, output_files[0])
        instr_path = os.path.join(temp_dir, output_files[1])
        
        for f in output_files:
            full_p = os.path.join(temp_dir, f)
            if "(Vocals)" in f: vocals_path = full_p
            if "(Instrumental)" in f or "(Noise)" in f: instr_path = full_p

        # 4. Возврат в Comfy
        try:
            # ВАЖНО: Читаем Sample Rate из обработанного файла!
            # Модели часто меняют частоту на 44100, и если вернуть старую - будет слоу-мо.
            vocals_wave, vocals_sr = torchaudio.load(vocals_path)
            instr_wave, instr_sr = torchaudio.load(instr_path)
        except Exception as e:
            shutil.rmtree(temp_dir)
            raise RuntimeError(f"Ошибка чтения результата: {e}")
        
        # Формируем формат (Batch, Channels, Samples)
        vocals_out = {"waveform": vocals_wave.unsqueeze(0), "sample_rate": vocals_sr}
        instr_out = {"waveform": instr_wave.unsqueeze(0), "sample_rate": instr_sr}

        # Удаляем временные файлы
        try:
            shutil.rmtree(temp_dir)
        except:
            pass
            
        print("🎧 Audio Separator: Done!")

        return (vocals_out, instr_out)