"""Main SalTTS inference engine."""

import torch
import numpy as np
from typing import Optional, Dict
import time

from saltts.core.stream_transformer import StreamTransformer
from saltts.core.side_processor import SideProcessor
from saltts.core.aux_config_gen import AuxConfigGenerator
from saltts.core.mixer import Mixer
from saltts.models.generator_a.flow import FlowGenerator
from saltts.models.generator_a.diffusion import DiffusionGenerator
from saltts.models.generator_a.gan import GANGenerator
from saltts.models.generator_b.parametric import ParametricGenerator
from saltts.postprocessing.pipeline import PostProcessingPipeline
from saltts.rl.agent import SACAgent
from saltts.nmdb_integration.channel_manager import ChannelManager
from saltts.utils.config import Config
from saltts.utils.logging import setup_logger
from saltts.core.data_structures import AudioOutput


class SalTTSEngine:
    """Main SalTTS inference engine.

    Integrates all components for end-to-end speech synthesis.
    """

    def __init__(
        self,
        config_dir: str = "config",
        device: str = "cuda",
        generator_a_type: str = "flow",
    ):
        """Initialize SalTTS engine.

        Args:
            config_dir: Configuration directory
            device: Device (cuda or cpu)
            generator_a_type: Type of Generator A (flow/diffusion/gan)
        """
        self.device = device
        self.generator_a_type = generator_a_type

        # Load configuration
        self.config = Config(config_dir)
        self.model_config = self.config.load("model")
        self.runtime_config = self.config.load("runtime")

        # Setup logger
        self.logger = setup_logger("saltts_engine")

        # Initialize components
        self._init_components()

        self.logger.info("SalTTS engine initialized successfully")

    def _init_components(self) -> None:
        """Initialize all components."""
        # Prosody planning pipeline
        self.stream_transformer = StreamTransformer().to(self.device)
        self.side_processor = SideProcessor().to(self.device)
        self.aux_config_gen = AuxConfigGenerator().to(self.device)
        self.mixer = Mixer().to(self.device)

        # Generator A
        if self.generator_a_type == "flow":
            self.generator_a = FlowGenerator().to(self.device)
        elif self.generator_a_type == "diffusion":
            self.generator_a = DiffusionGenerator().to(self.device)
        elif self.generator_a_type == "gan":
            self.generator_a = GANGenerator().to(self.device)
        else:
            raise ValueError(f"Unknown generator type: {self.generator_a_type}")

        # Generator B
        self.generator_b = ParametricGenerator().to(self.device)

        # Post-processing
        self.postproc = PostProcessingPipeline()

        # RL agent
        self.rl_agent = SACAgent(device=self.device)

        # Set to eval mode
        self.stream_transformer.eval()
        self.side_processor.eval()
        self.aux_config_gen.eval()
        self.mixer.eval()
        self.generator_a.eval()
        self.generator_b.eval()

    def synthesize(
        self,
        text: str,
        emotion: str = "neutral",
        speaker_id: int = 0,
        use_rl: bool = True,
    ) -> AudioOutput:
        """Synthesize speech from text.

        Args:
            text: Input text
            emotion: Emotion label
            speaker_id: Speaker ID
            use_rl: Whether to use RL-based mixing

        Returns:
            AudioOutput with waveform and metadata
        """
        start_time = time.time()

        with torch.no_grad():
            # Placeholder: Convert text to token IDs
            text_ids = torch.randint(0, 100, (1, len(text))).to(self.device)

            # Placeholder: Get embeddings (would come from NMDB)
            emb1 = torch.randn(1, 384).to(self.device)
            emb2 = torch.randn(1, 512).to(self.device)
            emb3 = torch.randn(1, 512).to(self.device)
            data3 = torch.randn(1, 128).to(self.device)

            # Prosody planning
            cfg = self.stream_transformer(text_ids, emb2)
            side_feat = self.side_processor(emb1, emb2)
            aux_params = self.aux_config_gen(data3, emb3)
            mixed_params = self.mixer(cfg, side_feat, aux_params)

            # Generate audio with both generators
            conditioning = torch.randn(1, 256).to(self.device)
            mel_a = self.generator_a.infer(mixed_params, conditioning)
            audio_b_list = self.generator_b.synthesize(mixed_params)

            # RL-based mixing (if enabled)
            if use_rl:
                # Placeholder: Use RL agent to determine mixing weights
                pass

            # For now, use Generator A output
            # Convert mel to audio (placeholder)
            audio = torch.randn(24000 * 2).numpy()

            # Post-processing
            audio = self.postproc.process(audio)

        latency_ms = (time.time() - start_time) * 1000

        return AudioOutput(
            waveform=audio,
            sample_rate=24000,
            latency_ms=latency_ms,
            generator_type="A",
        )

    def save_audio(self, audio_output: AudioOutput, path: str) -> None:
        """Save audio to file.

        Args:
            audio_output: Audio output
            path: Output path
        """
        import soundfile as sf
        sf.write(path, audio_output.waveform, audio_output.sample_rate)
        self.logger.info(f"Audio saved to {path}")
