"""Complete post-processing pipeline for SalTTS."""

import numpy as np
import torch
from typing import Optional

from saltts.postprocessing.algo_proc_1 import AlgoProc1
from saltts.postprocessing.model_proc import ModelProc
from saltts.postprocessing.algo_proc_2 import AlgoProc2


class PostProcessingPipeline:
    """3-stage post-processing pipeline.

    Stage 1: Algorithm Processing (de-clicking, gating, compression)
    Stage 2: Model Processing (neural enhancement)
    Stage 3: Algorithm Processing (hard limiting, DC removal, fading)
    """

    def __init__(
        self,
        sample_rate: int = 24000,
        enable_stage1: bool = True,
        enable_stage2: bool = True,
        enable_stage3: bool = True,
        model_checkpoint: Optional[str] = None,
    ):
        """Initialize post-processing pipeline.

        Args:
            sample_rate: Audio sample rate
            enable_stage1: Enable algorithm processing stage 1
            enable_stage2: Enable model processing
            enable_stage3: Enable algorithm processing stage 2
            model_checkpoint: Path to model checkpoint
        """
        self.sample_rate = sample_rate
        self.enable_stage1 = enable_stage1
        self.enable_stage2 = enable_stage2
        self.enable_stage3 = enable_stage3

        # Initialize stages
        self.stage1 = AlgoProc1(sample_rate=sample_rate)
        self.stage2 = ModelProc()
        self.stage3 = AlgoProc2(sample_rate=sample_rate)

        # Load model checkpoint if provided
        if model_checkpoint and enable_stage2:
            self.stage2.load_state_dict(torch.load(model_checkpoint))
            self.stage2.eval()

    def process(self, audio: np.ndarray) -> np.ndarray:
        """Process audio through complete pipeline.

        Args:
            audio: Input audio waveform

        Returns:
            Processed audio
        """
        # Stage 1: Algorithm processing
        if self.enable_stage1:
            audio = self.stage1.process(audio)

        # Stage 2: Model processing
        if self.enable_stage2:
            audio = self.stage2.process(audio)

        # Stage 3: Algorithm processing
        if self.enable_stage3:
            audio = self.stage3.process(audio)

        return audio

    def process_batch(self, audio_list: list[np.ndarray]) -> list[np.ndarray]:
        """Process batch of audio.

        Args:
            audio_list: List of audio waveforms

        Returns:
            List of processed audio
        """
        return [self.process(audio) for audio in audio_list]
