import torch


class LoadAudioandSplit:
    """
    Splits audio into chunks where each chunk's duration is calculated from
    connected frame counts and FPS. Duration = frames / fps
    Auto-detects how many frames_N inputs are connected.
    """
    RETURN_TYPES = ("INT", "FLOAT") + tuple(["AUDIO"] * 8)
    RETURN_NAMES = ("scene_count", "final_timestamp") + tuple([f"audio_{i}" for i in range(1, 9)])
    FUNCTION = "split_audio"
    CATEGORY = "wankeyframebuilder"

    @classmethod
    def INPUT_TYPES(cls):
        optional_inputs = {}
        for i in range(1, 9):
            optional_inputs[f"frames_{i}"] = ("INT", {"default": None, "min": 1, "forceInput": True})
        
        return {
            "required": {
                "audio": ("AUDIO",),
                "fps": ("FLOAT", {"default": 25.0, "min": 0.001, "max": 120.0}),
                "overlap_frames": ("INT", {"default": 0, "min": 0, "max": 500}),
            },
            "optional": optional_inputs
        }

    def split_audio(self, audio, fps, overlap_frames=0, **kwargs):
        make_stereo = True

        waveform = audio["waveform"]
        src_sample_rate = int(audio.get("sample_rate", 44100))

        if waveform.ndim == 2:
            waveform = waveform.unsqueeze(0)

        total_samples = waveform.shape[-1]
        audio_total_duration = float(total_samples) / float(src_sample_rate)

        # auto-detect connected frame inputs (contiguous from frames_1)
        frame_counts = []
        for i in range(1, 9):
            frame_key = f"frames_{i}"
            val = kwargs.get(frame_key, None)
            if val is None:
                break
            frame_counts.append(int(val))

        scene_count = len(frame_counts)

        # if nothing connected, default to 1 scene with 97 frames
        if scene_count == 0:
            frame_counts = [97]
            scene_count = 1

        # calculate durations from frame counts and fps
        durations = [float(fc) / float(fps) for fc in frame_counts]
        overlap_seconds = float(overlap_frames) / float(fps)

        # compute start times with overlap
        starts = []
        current_time = 0.0
        for idx, d in enumerate(durations):
            starts.append(current_time)
            # next chunk starts at (duration - overlap) after this one
            if idx < len(durations) - 1:
                current_time += d - overlap_seconds

        segments = []

        for idx, start_time in enumerate(starts):
            requested_duration = durations[idx]
            start_samp = max(0, int(start_time * src_sample_rate))
            end_samp = min(total_samples, int(start_samp + requested_duration * src_sample_rate))

            seg = waveform[..., start_samp:end_samp]

            # if segment is empty, create silence
            if seg.numel() == 0:
                silence_samples = int(requested_duration * src_sample_rate)
                seg = torch.zeros((1, 2, silence_samples), dtype=waveform.dtype)

            if make_stereo and seg.shape[1] == 1:
                seg = seg.repeat(1, 2, 1)

            segments.append({"waveform": seg, "sample_rate": src_sample_rate})

        # pad with None for unused outputs
        while len(segments) < 8:
            segments.append(None)

        # calculate end timestamp of last chunk
        last_chunk_end = starts[-1] + durations[-1]

        return (scene_count, float(last_chunk_end), *tuple(segments))


NODE_CLASS_MAPPINGS = {
    "LoadAudioandSplit": LoadAudioandSplit
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadAudioandSplit": "Load Audio and Split"
}