import time
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Union

import mlx.core as mx
from mlx_lm.generate import stream_generate
from mlx_lm.models.qwen3 import Model as Qwen3Model
from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig
from mlx_lm.sample_utils import make_logits_processors, make_sampler
from tqdm import tqdm

from mlx_audio.codec.models.snac import SNAC
from mlx_audio.utils import load_audio

from ..base import GenerationResult

# VyvoTTS special token IDs (Qwen3-based tokenizer)
TOKENIZER_LENGTH = 151669
START_OF_TEXT = 151643
END_OF_TEXT = 151645
START_OF_SPEECH = TOKENIZER_LENGTH + 1  # 151670
END_OF_SPEECH = TOKENIZER_LENGTH + 2  # 151671
START_OF_HUMAN = TOKENIZER_LENGTH + 3  # 151672
END_OF_HUMAN = TOKENIZER_LENGTH + 4  # 151673
START_OF_AI = TOKENIZER_LENGTH + 5  # 151674
END_OF_AI = TOKENIZER_LENGTH + 6  # 151675
PAD_TOKEN = TOKENIZER_LENGTH + 7  # 151676
AUDIO_TOKENS_START = TOKENIZER_LENGTH + 10  # 151679


@dataclass
class ModelConfig(Qwen3ModelConfig):
    tokenizer_name: str = None
    sample_rate: int = 24000


snac_model = SNAC.from_pretrained("mlx-community/snac_24khz").eval()


def decode_audio_from_codes(code_list):
    layer_1 = []
    layer_2 = []
    layer_3 = []
    for i in range((len(code_list) + 1) // 7):
        layer_1.append(code_list[7 * i])
        layer_2.append(code_list[7 * i + 1] - 4096)
        layer_3.append(code_list[7 * i + 2] - (2 * 4096))
        layer_3.append(code_list[7 * i + 3] - (3 * 4096))
        layer_2.append(code_list[7 * i + 4] - (4 * 4096))
        layer_3.append(code_list[7 * i + 5] - (5 * 4096))
        layer_3.append(code_list[7 * i + 6] - (6 * 4096))
    codes = [
        mx.expand_dims(mx.array(layer_1), 0),
        mx.expand_dims(mx.array(layer_2), 0),
        mx.expand_dims(mx.array(layer_3), 0),
    ]
    audio_hat = snac_model.decode(codes).squeeze(-1)
    return audio_hat


def encode_audio_to_codes(audio):
    audio = audio[None, None, :]

    codes = snac_model.encode(audio)

    layer_1 = codes[0].squeeze(0).tolist()
    layer_2 = codes[1].squeeze(0).tolist()
    layer_3 = codes[2].squeeze(0).tolist()

    code_list = []
    num_groups = len(layer_1)
    for i in range(num_groups):
        code_list.append(layer_1[i])
        code_list.append(layer_2[2 * i] + 4096)
        code_list.append(layer_3[4 * i] + 2 * 4096)
        code_list.append(layer_3[4 * i + 1] + 3 * 4096)
        code_list.append(layer_2[2 * i + 1] + 4 * 4096)
        code_list.append(layer_3[4 * i + 2] + 5 * 4096)
        code_list.append(layer_3[4 * i + 3] + 6 * 4096)

    return mx.array(code_list)[None, :]


class Model(Qwen3Model):
    def __init__(self, config: ModelConfig, **kwargs):
        super().__init__(config)
        self.config = config
        self.model_type = config.model_type
        self.tokenizer = None

    @property
    def layers(self):
        return self.model.layers

    @classmethod
    def post_load_hook(cls, model: "Model", model_path: Path):
        """
        Hook called after model weights are loaded.
        Used to initialize the tokenizer which is required for text input.
        """
        from transformers import AutoTokenizer

        if model.tokenizer is None:
            model.tokenizer = AutoTokenizer.from_pretrained(str(model_path))
        return model

    @property
    def sample_rate(self):
        return self.config.sample_rate

    def parse_output(self, input_ids):
        token_to_find = START_OF_SPEECH  # 151670
        token_to_remove = END_OF_SPEECH  # 151671

        # MLX doesn't have nonzero, so we need to create indices manually
        mask = input_ids == token_to_find
        indices = []
        for i in range(mask.shape[0]):
            for j in range(mask.shape[1]):
                if mask[i, j]:
                    indices.append((i, j))
        token_indices = [[], []]
        for i, j in indices:
            token_indices[0].append(i)
            token_indices[1].append(j)

        # Check if we found any tokens BEFORE converting to MLX array
        # to avoid "Cannot do a non-empty take from an array with zero elements" error
        if len(token_indices[1]) > 0:
            token_indices = mx.array(token_indices)
            last_occurrence_idx = int(token_indices[1][-1])
            cropped_tensor = input_ids[:, last_occurrence_idx + 1 :]
        else:
            cropped_tensor = input_ids

        mask = cropped_tensor != token_to_remove

        processed_rows = []

        for row in cropped_tensor:
            # Create a mask and filter manually since boolean indexing isn't supported
            row_list = row.tolist()
            masked_row = mx.array([val for val in row_list if val != token_to_remove])
            processed_rows.append(masked_row)

        code_lists = []

        for row in processed_rows:
            row_length = row.shape[0]
            new_length = (row_length // 7) * 7
            trimmed_row = row[:new_length]
            trimmed_row = [t - AUDIO_TOKENS_START for t in trimmed_row]  # 151679
            code_lists.append(trimmed_row)

        return code_lists

    def prepare_zeroprompt(
        self,
        ref_audio: mx.array,
        ref_text: str,
    ):
        """Prepare the reference audio context (zeroprompt) for voice cloning."""
        print(
            "\033[93mWARNING: Audio cloning doesn't work reliably on this model.\033[0m"
        )
        audio_input_ids = encode_audio_to_codes(ref_audio) + AUDIO_TOKENS_START
        audio_transcript_ids = self.tokenizer(ref_text, return_tensors="mlx").input_ids

        start_token = mx.array([[START_OF_HUMAN]], dtype=mx.int64)  # 151672
        end_tokens = mx.array(
            [[END_OF_TEXT, END_OF_HUMAN]], dtype=mx.int64
        )  # 151645, 151673
        audio_start_tokens = mx.array(
            [[START_OF_AI, START_OF_SPEECH]], dtype=mx.int64
        )  # 151674, 151670
        audio_end_tokens = mx.array(
            [[END_OF_SPEECH, END_OF_AI]], dtype=mx.int64
        )  # 151671, 151675

        # [SOH] [transcript] [EOT EOH] [SOA SOS] [audio_tokens] [EOS EOA]
        zeroprompt = mx.concatenate(
            [
                start_token,
                audio_transcript_ids,
                end_tokens,
                audio_start_tokens,
                audio_input_ids,
                audio_end_tokens,
            ],
            axis=1,
        )
        return zeroprompt

    def prepare_input_ids(
        self,
        prompt: Union[str, List[str]],
        voice: Optional[str] = None,
        zeroprompt: Optional[mx.array] = None,
    ):
        """Prepare input ids for a single prompt or batch of prompts, optionally with zeroprompt prefix."""
        # Handle single prompt vs batch
        if isinstance(prompt, str):
            prompts = [prompt]
        else:
            prompts = prompt

        all_input_ids = []
        for p in prompts:
            if voice is not None and zeroprompt is None:
                p = f"{voice}: {p}"

            start_token = mx.array([[START_OF_HUMAN]], dtype=mx.int64)  # 151672
            end_tokens = mx.array(
                [[END_OF_TEXT, END_OF_HUMAN]], dtype=mx.int64
            )  # 151645, 151673

            input_ids = self.tokenizer(p, return_tensors="mlx").input_ids
            prompt_input_ids = mx.concatenate(
                [start_token, input_ids, end_tokens], axis=1
            )  # [SOH] [text] [EOT EOH]

            if zeroprompt is not None:
                prompt_input_ids = mx.concatenate(
                    [zeroprompt, prompt_input_ids], axis=1
                )

            all_input_ids.append(prompt_input_ids)

        # If only one prompt, return as-is (keeps original behavior)
        if len(all_input_ids) == 1:
            return all_input_ids[0]

        # Pad to same length and stack for batch
        max_len = max(ids.shape[1] for ids in all_input_ids)
        padded = []
        for ids in all_input_ids:
            if ids.shape[1] < max_len:
                padding = mx.zeros((1, max_len - ids.shape[1]), dtype=mx.int64)
                ids = mx.concatenate([ids, padding], axis=1)
            padded.append(ids)

        return mx.concatenate(padded, axis=0)

    def generate_result(
        self, audio, start_time: float, token_count: int, segment_idx: int, **kwargs
    ) -> GenerationResult:
        """Helper to create a GenerationResult from audio."""
        samples = audio.shape[0] if audio is not None else 0
        assert samples > 0, "No audio generated"

        sample_rate = self.config.sample_rate
        audio_duration_seconds = samples / sample_rate

        elapsed_time = time.perf_counter() - start_time
        rtf = audio_duration_seconds / elapsed_time if elapsed_time > 0 else 0

        duration_hours = int(audio_duration_seconds // 3600)
        duration_mins = int((audio_duration_seconds % 3600) // 60)
        duration_secs = int(audio_duration_seconds % 60)
        duration_ms = int((audio_duration_seconds % 1) * 1000)
        duration_str = (
            f"{duration_hours:02d}:{duration_mins:02d}:"
            f"{duration_secs:02d}.{duration_ms:03d}"
        )

        return GenerationResult(
            audio=audio,
            samples=samples,
            sample_rate=sample_rate,
            segment_idx=segment_idx,
            token_count=token_count,
            audio_duration=duration_str,
            real_time_factor=rtf,
            prompt={
                "tokens": token_count,
                "tokens-per-sec": (
                    round(token_count / elapsed_time, 2) if elapsed_time > 0 else 0
                ),
            },
            audio_samples={
                "samples": samples,
                "samples-per-sec": (
                    round(samples / elapsed_time, 2) if elapsed_time > 0 else 0
                ),
            },
            processing_time_seconds=elapsed_time,
            peak_memory_usage=mx.get_peak_memory() / 1e9,
        )

    def generate(
        self,
        text,
        voice: str,
        temperature: float = 0.6,
        top_p: float = 0.8,
        split_pattern: str = "\n",
        max_tokens: int = 1200,
        verbose: bool = False,
        ref_audio: Optional[Union[str, mx.array]] = None,
        ref_text: Optional[str] = None,
        stream: bool = False,
        streaming_interval: float = 2.0,
        **kwargs,
    ):
        # Load reference audio if provided (handles file paths and mx.array)
        if ref_audio is not None:
            ref_audio = load_audio(ref_audio, sample_rate=self.sample_rate)

        prompt_text = text.replace("\\n", "\n").replace("\\t", "\t")
        prompts = [p for p in prompt_text.split(split_pattern) if p.strip()]

        # Prepare zeroprompt once if voice cloning is requested
        zeroprompt = None
        if ref_audio is not None and ref_text is not None:
            zeroprompt = self.prepare_zeroprompt(ref_audio, ref_text)

        sampler = make_sampler(temperature, top_p, top_k=kwargs.get("top_k", -1))
        logits_processors = make_logits_processors(
            kwargs.get("logit_bias", None),
            kwargs.get("repetition_penalty", 1.3),
            kwargs.get("repetition_context_size", 20),
        )

        streaming_token_interval = int(streaming_interval * 137.5)

        # Process each prompt segment individually for memory efficiency
        for segment_idx, segment_prompt in enumerate(prompts):
            time_start = time.perf_counter()

            # Prepare input_ids for this segment (with zeroprompt if voice cloning)
            input_ids = self.prepare_input_ids(
                segment_prompt,
                voice,
                zeroprompt,
            )

            generated_token_count = 0
            yielded_token_count = 0
            yielded_frame_count = 0

            # Generate tokens for this segment
            for i, response in enumerate(
                tqdm(
                    stream_generate(
                        self,
                        tokenizer=self.tokenizer,
                        prompt=input_ids.squeeze(0),
                        max_tokens=max_tokens,
                        sampler=sampler,
                        logits_processors=logits_processors,
                    ),
                    total=max_tokens,
                    disable=not verbose,
                    desc=f"Segment {segment_idx + 1}/{len(prompts)}",
                )
            ):
                next_token = mx.array([response.token])
                input_ids = mx.concatenate([input_ids, next_token[None, :]], axis=1)
                generated_token_count += 1

                if i % 50 == 0:
                    mx.clear_cache()

                # Stream partial audio at intervals
                if stream and generated_token_count % streaming_token_interval == 0:
                    code_lists = self.parse_output(input_ids)
                    if code_lists and len(code_lists[0]) > 0:
                        audio = decode_audio_from_codes(code_lists[0])[0]
                        if audio.shape[0] > yielded_frame_count:
                            yield self.generate_result(
                                audio=audio[yielded_frame_count:],
                                start_time=time_start,
                                token_count=generated_token_count - yielded_token_count,
                                segment_idx=segment_idx,
                            )
                            yielded_token_count = generated_token_count
                            yielded_frame_count = audio.shape[0]
                            time_start = time.perf_counter()

                if next_token == END_OF_SPEECH:  # 151671
                    break

            # Decode and yield remaining audio for this segment
            code_lists = self.parse_output(input_ids)

            for code_list in code_lists:
                if len(code_list) == 0:
                    continue
                audio = decode_audio_from_codes(code_list)[0]

                if audio.shape[0] > yielded_frame_count:
                    yield self.generate_result(
                        audio=audio[yielded_frame_count:],
                        start_time=time_start,
                        token_count=generated_token_count - yielded_token_count,
                        segment_idx=segment_idx,
                    )

            # Clear cache after each segment to avoid memory buildup
            mx.clear_cache()
