"""
Audio Encoder + Speech Response Module for Neon VLA.

Enables the robot to:
1. HEAR — Understand spoken commands via audio waveform encoding
2. SPEAK — Generate spoken responses via PersonaPlex TTS

Two audio encoding strategies:
- Native: Use Qwen2.5-Omni which processes audio tokens natively
- Whisper: Use Whisper encoder as a separate audio feature extractor

PersonaPlex voices give the robot personality — it speaks back to you
while executing actions. The robot can narrate what it's doing.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


@dataclass
class AudioConfig:
    """Configuration for audio input/output."""

    # Input
    enabled: bool = True
    encoder_type: str = "whisper"  # "whisper", "omni", or "mel"
    whisper_model: str = "openai/whisper-base"  # For whisper encoder
    sample_rate: int = 16000  # Whisper expects 16kHz
    max_audio_seconds: float = 15.0  # Max audio clip length
    freeze_encoder: bool = True
    # Feature projection
    audio_hidden_size: int = 512  # Whisper base hidden size
    projection_dim: int = 0  # 0 = match backbone hidden_size
    # Output (PersonaPlex TTS)
    enable_speech_output: bool = True
    personaplex_voice: str = "NATM1"  # Default robot voice
    personaplex_host: str = "http://192.168.1.151:8400"  # Thor
    # Training
    use_audio_text_pairs: bool = True  # Train on paired audio + transcript


class AudioEncoder(nn.Module):
    """Encodes audio waveforms into feature vectors for the VLA.

    Supports:
    1. Whisper encoder — robust ASR features, works with any backbone
    2. Mel spectrogram — lightweight, no pretrained model needed
    3. Omni — native audio processing (when using Qwen2.5-Omni backbone)
    """

    def __init__(self, config: AudioConfig, backbone_hidden_size: int = 2048):
        super().__init__()
        self.config = config
        self._encoder = None
        self._processor = None

        # Projection from audio features to backbone dimension
        proj_dim = config.projection_dim if config.projection_dim > 0 else backbone_hidden_size
        self.projection = nn.Sequential(
            nn.Linear(config.audio_hidden_size, proj_dim),
            nn.GELU(),
            nn.LayerNorm(proj_dim),
        )
        self._proj_dim = proj_dim

    @property
    def output_dim(self) -> int:
        return self._proj_dim

    def load(self):
        """Load the audio encoder model."""
        if self._encoder is not None:
            return

        if self.config.encoder_type == "whisper":
            self._load_whisper()
        elif self.config.encoder_type == "mel":
            self._build_mel_encoder()
        elif self.config.encoder_type == "omni":
            # Omni mode — audio is handled natively by the backbone
            # No separate encoder needed
            logger.info("Audio encoder: omni mode (backbone handles audio)")
            return
        else:
            raise ValueError(f"Unknown encoder_type: {self.config.encoder_type}")

    def _load_whisper(self):
        """Load Whisper encoder for audio feature extraction."""
        from transformers import WhisperModel, WhisperFeatureExtractor

        logger.info(f"Loading Whisper encoder: {self.config.whisper_model}")
        model = WhisperModel.from_pretrained(self.config.whisper_model)
        self._encoder = model.encoder  # Only the encoder, not the decoder
        self._processor = WhisperFeatureExtractor.from_pretrained(self.config.whisper_model)

        # Update hidden size from actual model
        actual_hidden = self._encoder.config.d_model  # 512 for base, 768 for small, etc.
        if actual_hidden != self.config.audio_hidden_size:
            logger.info(f"Updating audio_hidden_size: {self.config.audio_hidden_size} → {actual_hidden}")
            # Rebuild projection with correct input dim
            self.projection = nn.Sequential(
                nn.Linear(actual_hidden, self._proj_dim),
                nn.GELU(),
                nn.LayerNorm(self._proj_dim),
            )

        if self.config.freeze_encoder:
            for param in self._encoder.parameters():
                param.requires_grad = False
            logger.info("Whisper encoder frozen")

        logger.info(f"Whisper encoder loaded: hidden_size={actual_hidden}")

    def _build_mel_encoder(self):
        """Build a lightweight mel spectrogram encoder (no pretrained model)."""
        import torchaudio

        # Mel spectrogram → small CNN → features
        self._mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=self.config.sample_rate,
            n_fft=400,
            hop_length=160,
            n_mels=80,
        )
        self._encoder = nn.Sequential(
            nn.Conv1d(80, 256, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Conv1d(256, self.config.audio_hidden_size, kernel_size=3, padding=1),
            nn.GELU(),
            nn.AdaptiveAvgPool1d(1),  # Pool to single vector
        )
        logger.info("Mel spectrogram encoder built")

    def encode(self, audio: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Encode audio waveform into features.

        Args:
            audio: Raw waveform (batch, samples) at self.config.sample_rate
            attention_mask: Optional mask for padded audio

        Returns:
            (batch, output_dim) audio features
        """
        if self._encoder is None:
            self.load()

        if self.config.encoder_type == "whisper":
            return self._encode_whisper(audio)
        elif self.config.encoder_type == "mel":
            return self._encode_mel(audio)
        else:
            raise ValueError(f"Cannot encode with encoder_type={self.config.encoder_type}")

    def _encode_whisper(self, audio: torch.Tensor) -> torch.Tensor:
        """Encode using Whisper's encoder."""
        device = next(self._encoder.parameters()).device

        # Process audio → mel features
        # Whisper expects 30s of audio at 16kHz, padded/truncated
        if audio.ndim == 1:
            audio = audio.unsqueeze(0)  # Add batch dim

        # Convert to numpy for feature extractor
        audio_np = audio.cpu().numpy()

        # Process each sample in batch
        features_list = []
        for i in range(audio_np.shape[0]):
            inputs = self._processor(
                audio_np[i],
                sampling_rate=self.config.sample_rate,
                return_tensors="pt",
            )
            input_features = inputs.input_features.to(device)  # (1, 80, 3000)

            # Forward through Whisper encoder
            with torch.no_grad() if self.config.freeze_encoder else torch.enable_grad():
                encoder_output = self._encoder(input_features)
                hidden = encoder_output.last_hidden_state  # (1, seq_len, hidden_size)

            # Pool: mean across time
            pooled = hidden.mean(dim=1)  # (1, hidden_size)
            features_list.append(pooled)

        features = torch.cat(features_list, dim=0)  # (batch, hidden_size)

        # Project to backbone dimension
        return self.projection(features)

    def _encode_mel(self, audio: torch.Tensor) -> torch.Tensor:
        """Encode using lightweight mel spectrogram CNN."""
        if audio.ndim == 1:
            audio = audio.unsqueeze(0)

        # Mel spectrogram
        mel = self._mel_transform(audio)  # (batch, n_mels, time)

        # CNN encoder
        features = self._encoder(mel)  # (batch, hidden_size, 1)
        features = features.squeeze(-1)  # (batch, hidden_size)

        # Project
        return self.projection(features)

    def forward(self, audio: torch.Tensor) -> torch.Tensor:
        """Forward pass — encode audio to features."""
        return self.encode(audio)


class SpeechResponseHead(nn.Module):
    """Generates text responses for PersonaPlex TTS synthesis.

    The robot can narrate what it's doing while performing actions:
    - "I see the red cup. Reaching for it now."
    - "Moving to the table on my left."
    - "Task complete!"

    This head predicts a text response from the same features used for actions.
    The text is then synthesized to speech via PersonaPlex TTS.
    """

    def __init__(self, hidden_size: int, vocab_size: int = 32000, max_response_tokens: int = 50):
        super().__init__()
        self.max_tokens = max_response_tokens
        # Simple token predictor (lightweight — real generation happens via backbone)
        self.response_classifier = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.GELU(),
            nn.Linear(hidden_size // 2, vocab_size),
        )
        # Response type classifier (what kind of speech to generate)
        self.response_type = nn.Sequential(
            nn.Linear(hidden_size, 64),
            nn.GELU(),
            nn.Linear(64, 5),  # 5 types: narrate, confirm, warn, ask, silent
        )

    def forward(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Predict speech response from features.

        Args:
            features: (batch, hidden_size) from fusion layer

        Returns:
            Dict with "response_type" logits and "token_logits"
        """
        return {
            "response_type": self.response_type(features),
            "token_logits": self.response_classifier(features),
        }


class PersonaPlexSpeaker:
    """Interface to PersonaPlex TTS for robot speech output.

    Connects to the PersonaPlex server (running on Thor) to synthesize
    speech from text. The robot speaks in a consistent voice.

    Usage:
        speaker = PersonaPlexSpeaker(voice="NATM1")
        audio_path = speaker.speak("I see the red cup. Reaching for it now.")
    """

    RESPONSE_TEMPLATES = {
        "narrate": [
            "I see {object}. {action} now.",
            "Moving to {location}.",
            "{action} the {object}.",
        ],
        "confirm": [
            "Got it. {action}.",
            "Understood. Working on it.",
            "On it.",
        ],
        "warn": [
            "I can't reach that safely.",
            "Obstacle detected. Adjusting path.",
            "That seems dangerous. Should I proceed?",
        ],
        "ask": [
            "Which {object} do you mean?",
            "Could you point to it?",
            "I see multiple options. Can you clarify?",
        ],
    }

    def __init__(self, config: AudioConfig):
        self.config = config
        self.voice = config.personaplex_voice
        self.host = config.personaplex_host

    def speak(self, text: str, output_path: str = None) -> Optional[str]:
        """Synthesize speech from text via PersonaPlex.

        Args:
            text: Text to speak
            output_path: Where to save the WAV file

        Returns:
            Path to generated audio file, or None if failed
        """
        if not text or not self.config.enable_speech_output:
            return None

        try:
            import requests
            import tempfile

            if output_path is None:
                output_path = tempfile.mktemp(suffix=".wav", prefix="neon_speech_")

            # Try PersonaPlex server
            response = requests.post(
                f"{self.host}/generate",
                json={
                    "text": text,
                    "voice": self.voice,
                },
                timeout=10,
            )

            if response.status_code == 200:
                with open(output_path, "wb") as f:
                    f.write(response.content)
                logger.info(f"Speech generated: {output_path}")
                return output_path

        except Exception as e:
            logger.warning(f"PersonaPlex TTS failed: {e}")

        # Fallback: try local PersonaPlex
        try:
            import subprocess

            result = subprocess.run(
                [
                    "python3", "-m", "personaplex",
                    "--text", text,
                    "--voice", self.voice,
                    "--output", output_path or "/tmp/neon_speech.wav",
                ],
                capture_output=True,
                timeout=60,
            )
            if result.returncode == 0:
                return output_path

        except Exception as e:
            logger.warning(f"Local PersonaPlex fallback failed: {e}")

        return None

    def speak_action_narration(self, response_type: str, context: Dict[str, str] = None) -> Optional[str]:
        """Generate and speak an action narration.

        Args:
            response_type: "narrate", "confirm", "warn", "ask", "silent"
            context: Dict with {object, location, action} for template filling

        Returns:
            Path to generated audio, or None
        """
        if response_type == "silent" or not self.config.enable_speech_output:
            return None

        templates = self.RESPONSE_TEMPLATES.get(response_type, self.RESPONSE_TEMPLATES["confirm"])

        # Pick a template and fill in context
        import random
        template = random.choice(templates)
        context = context or {}
        try:
            text = template.format(**context)
        except KeyError:
            text = template.split("{")[0].strip() or "Working on it."

        return self.speak(text)
