import os
import shutil
import subprocess
from pathlib import Path
import numpy as np
import torch
import folder_paths
from scipy.io.wavfile import write

class LikeSpiderUINode:
    """
    Base class for declarative UI-driven ComfyUI nodes.
    Extend this class and define `UI_CONFIG` and `run(...)`.
    """

    @classmethod
    def INPUT_TYPES(cls):
        config = cls.UI_CONFIG
        required = {}

        for field in config["inputs"]:
            name = field["name"]
            type_ = field["type"]
            default = field.get("default")

            if type_ == "string":
                required[name] = ("STRING", {"default": default or ""})
            elif type_ == "select":
                required[name] = (field["options"], {"default": default})
            elif type_ == "audio":
                required[name] = ("AUDIO",)
            else:
                raise ValueError(f"Unsupported field type: {type_}")

        return {"required": required}

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("audio",)
    FUNCTION = "interface_entry"
    OUTPUT_NODE = True
    CATEGORY = "LikeSpiderAI"

    def interface_entry(self, **kwargs):
        return self.run(**kwargs)

    def save_audio_with_ffmpeg(self, audio_data, filename_base, format="mp3", bitrate=192):
        samples = audio_data["waveform"]
        if isinstance(samples, torch.Tensor):
            samples = samples.detach().cpu().numpy()
        if samples.ndim == 3 and samples.shape[0] == 1:
            samples = samples[0]

        max_val = np.max(np.abs(samples))
        if max_val > 1.0:
            samples = samples / max_val
        samples = np.clip(samples, -1.0, 1.0)
        audio_int16 = (samples * 32767).astype(np.int16)

        if audio_int16.ndim == 2 and audio_int16.shape[0] in [1, 2]:
            wav_data = audio_int16.T
        elif audio_int16.ndim == 1:
            wav_data = audio_int16
        else:
            raise Exception(f"❌ Unsupported audio shape: {audio_int16.shape}")

        output_dir = Path(folder_paths.get_output_directory()) / "audio"
        output_dir.mkdir(parents=True, exist_ok=True)

        existing = sorted(output_dir.glob(f"{filename_base}_*.{format}"))
        index = len(existing)
        file_path = output_dir / f"{filename_base}_{index:05d}.{format}"

        temp_wav = output_dir / "temp_input.wav"
        write(temp_wav, 44100, wav_data)

        if format == "mp3":
            subprocess.run([
                "ffmpeg", "-y", "-i", str(temp_wav),
                "-codec:a", "libmp3lame", "-b:a", f"{bitrate}k",
                str(file_path)
            ], check=True)
        elif format == "wav":
            shutil.copy(temp_wav, file_path)
        elif format == "flac":
            subprocess.run([
                "ffmpeg", "-y", "-i", str(temp_wav),
                "-codec:a", "flac",
                str(file_path)
            ], check=True)
        else:
            raise Exception(f"Unsupported format: {format}")

        os.remove(temp_wav)
        print(f"[LikeSpiderUINode] ✅ Saved: {file_path}")
        return file_path