from typing import List

import soundfile as sf
import torch
from snac import SNAC
from transformers import AutoModelForCausalLM, AutoTokenizer


class OrpheusTTS:
    """
    A neural speech synthesis model that converts text to speech using SNAC audio codec.

    This class integrates a causal language model with SNAC (Speech Neural Audio Codec)
    to generate natural-sounding speech with paralinguistic features like laughter,
    sighing, and other vocal expressions.

    Example:
        >>> generator = OrpheusTTS()
        >>> prompts = ["Hello, I'm excited to meet you!"]
        >>> generator(prompts, voice="Sarah", output_prefix="greeting")
        # Creates greeting_0.wav file

    Attributes:
        device (str): Computation device ('cuda' or 'cpu')
        model: The causal language model for text-to-audio-token generation
        tokenizer: Tokenizer for text preprocessing
        snac_model: SNAC model for audio token decoding
    """

    def __init__(self, model_path: str = "canopylabs/orpheus-3b-0.1-ft", device: str = "cuda"):
        """
        Initialize the speech generator with specified model and device.

        Args:
            model_path: HuggingFace model path for the language model
            device: Computing device ('cuda' for GPU, 'cpu' for CPU)
        """
        self.device = device
        self._load_models(model_path)

    def _load_models(self, model_path: str):
        """Load SNAC and language models."""
        # Load SNAC audio codec model for decoding audio tokens
        self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to("cpu")

        # Load the main language model for text-to-audio-token generation
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path, torch_dtype=torch.bfloat16).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)

    def _prepare_inputs(self, prompts: List[str], voice: str):
        """
        Tokenize and format inputs for generation.

        Args:
            prompts: List of text strings to convert to speech
            voice: Voice identifier to prepend to each prompt

        Returns:
            Tuple of (input_ids, attention_mask) tensors ready for model generation
        """
        # Add voice prefix to each prompt
        prompts = [f"{voice}: {p}" for p in prompts]

        # Tokenize all prompts
        input_ids_list = [self.tokenizer(prompt, return_tensors="pt").input_ids for prompt in prompts]

        # Add special tokens: Start of Human (SOH) and End of Text/Human (EOT/EOH)
        start_token = torch.tensor([[128259]], dtype=torch.int64)  # SOH
        end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # EOT, EOH

        formatted_inputs = [torch.cat([start_token, ids, end_tokens], dim=1) for ids in input_ids_list]

        # Pad to same length for batch processing
        max_len = max(inp.shape[1] for inp in formatted_inputs)
        pad_token = 128263  # Padding token ID

        padded_inputs = []
        attention_masks = []

        for inp in formatted_inputs:
            padding_len = max_len - inp.shape[1]
            padding = torch.full((1, padding_len), pad_token, dtype=torch.int64)

            padded = torch.cat([padding, inp], dim=1)
            mask = torch.cat(
                [
                    torch.zeros((1, padding_len), dtype=torch.int64),
                    torch.ones((1, inp.shape[1]), dtype=torch.int64),
                ],
                dim=1,
            )

            padded_inputs.append(padded)
            attention_masks.append(mask)

        return (
            torch.cat(padded_inputs, dim=0).to(self.device),
            torch.cat(attention_masks, dim=0).to(self.device),
        )

    def _redistribute_codes(self, codes: List[int]):
        """
        Convert token codes to SNAC format and decode to audio.

        SNAC uses a hierarchical codec with 3 layers. This method redistributes
        the flat token sequence into the proper layer structure.

        Args:
            codes: List of integer codes from the language model

        Returns:
            Decoded audio tensor
        """
        # SNAC processes audio in groups of 7 tokens
        n_groups = (len(codes) + 1) // 7

        layer_1, layer_2, layer_3 = [], [], []

        for i in range(n_groups):
            base_idx = 7 * i
            # Layer 1: Coarse audio features (every 7th token)
            layer_1.append(codes[base_idx])
            # Layer 2: Medium resolution features
            layer_2.extend([codes[base_idx + 1] - 4096, codes[base_idx + 4] - 4 * 4096])
            # Layer 3: Fine-grained audio details
            layer_3.extend([
                codes[base_idx + 2] - 2 * 4096,
                codes[base_idx + 3] - 3 * 4096,
                codes[base_idx + 5] - 5 * 4096,
                codes[base_idx + 6] - 6 * 4096,
            ])

        snac_codes = [
            torch.tensor(layer_1).unsqueeze(0),
            torch.tensor(layer_2).unsqueeze(0),
            torch.tensor(layer_3).unsqueeze(0),
        ]

        return self.snac_model.decode(snac_codes)

    def _postprocess_tokens(self, generated_ids: torch.Tensor) -> List[List[int]]:
        """
        Extract and clean generated audio tokens.

        Args:
            generated_ids: Raw token output from language model

        Returns:
            List of cleaned token sequences, one per input prompt
        """
        # Find last occurrence of start token (128257) to locate audio tokens
        start_token, end_token = 128257, 128258
        start_positions = (generated_ids == start_token).nonzero(as_tuple=True)

        if len(start_positions[1]) > 0:
            last_start = start_positions[1][-1].item()
            tokens = generated_ids[:, last_start + 1:]
        else:
            tokens = generated_ids

        # Process each sequence
        code_lists = []
        for row in tokens:
            # Remove end tokens to get clean audio token sequence
            clean_tokens = row[row != end_token]

            # Trim to multiple of 7 (SNAC requirement) and adjust token values
            trim_len = (clean_tokens.size(0) // 7) * 7
            trimmed = clean_tokens[:trim_len]
            # Adjust token IDs to SNAC's expected range
            adjusted = [t.item() - 128266 for t in trimmed]

            code_lists.append(adjusted)

        return code_lists

    def __call__(self, prompts: List[str], voice: str, output_prefix: str = "sample"):
        """
        Generate speech from text prompts.

        Args:
            prompts: List of text strings to convert to speech
            voice: Voice identifier (e.g., "Sarah", "John")
            output_prefix: Prefix for output WAV files (default: "sample")

        Example:
            >>> generator = OrpheusTTS()
            >>> prompts = ["Hello world!", "How are you today?"]
            >>> generator.generate(prompts, "Emma", "greeting")
            # Creates greeting_0.wav and greeting_1.wav
        """
        input_ids, attention_mask = self._prepare_inputs(prompts, voice)

        # Generate audio tokens using the language model
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=1200,
                do_sample=True,
                temperature=0.6,
                top_p=0.95,
                repetition_penalty=1.1,
                num_return_sequences=1,
                eos_token_id=128258,
            )

        code_lists = self._postprocess_tokens(generated_ids)

        # Generate and save audio files
        for i, codes in enumerate(code_lists):
            audio = self._redistribute_codes(codes)
            # Save as 24kHz WAV file
            sf.write(
                f"{output_prefix}_{i}.wav",
                audio.detach().squeeze().cpu().numpy(),
                24000,
            )

        return code_lists
