"""
Modified From https://github.com/XXXXRT666/GPT-SoVITS
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Callable, List, MutableSequence, Protocol, Type, cast

import mlx.core as mx
import torch

from ..PyTorch.structs import T2SRequest, T2SResult
from .sample_funcs_mlx import SampleProtocolMLX, sample_naive

Tensor = torch.Tensor
Array = mx.array


@dataclass(slots=True)
class T2SRequestMLX:
    x: List[Array]
    x_lens: Array
    prompts: Array
    bert_feature: List[Array]
    valid_length: int
    top_k: int = 5
    top_p: float = 1
    early_stop_num: int = -1
    temperature: float = 1.0
    repetition_penalty: float = 1.35

    @classmethod
    def from_torch(cls, request: T2SRequest) -> T2SRequestMLX:
        x = list(map(lambda tensor: mx.array(tensor.cpu()), request.x))
        x_lens = mx.array(request.x_lens.cpu())
        prompts = mx.array(request.prompts.cpu())
        bert_feature = list(map(lambda tensor: mx.array(tensor.cpu()), request.bert_feature))

        return cls(
            x,
            x_lens,
            prompts,
            bert_feature,
            request.valid_length,
            request.top_k,
            request.top_p,
            request.early_stop_num,
            request.temperature,
            request.repetition_penalty,
        )


class KVCacheProtocol(Protocol):
    k_cache: Array
    v_cache: Array

    def empty(self) -> None: ...

    def update_cache(self, input_pos: Array, k_val: Array, v_val: Array, *args, **kwds) -> tuple[Array, Array]: ...

    def prefill_kv(self, k_val: Array, v_val: Array) -> None: ...

    def sync_cache(self, kv_cache: KVCacheProtocol) -> None: ...


class T2SDecoderProtocol(Protocol):
    max_seq_length: int
    EOS: int
    n_head: int

    def embed(self, x: list[Array], y: Array, bert_features: list[Array]) -> Array: ...


class T2SEngineProtocol(Protocol):
    def _handle_request(self, request: T2SRequest) -> tuple[list[Array], float]: ...

    def generate(self, request: T2SRequest) -> T2SResult: ...

    @staticmethod
    def load_decoder(
        weights_path: os.PathLike, max_batch_size: int = 1, implement: str = "MLX"
    ) -> T2SDecoderProtocol: ...


class T2SSessionMLX:
    def __init__(
        self,
        decoder: T2SDecoderProtocol,
        request_torch: T2SRequest,
        sample_func: Type[SampleProtocolMLX] = sample_naive,
        device: mx.Device = mx.Device(mx.cpu),
        dtype: mx.Dtype = mx.float32,
    ):
        with mx.stream(device):
            request = T2SRequestMLX.from_torch(request_torch)

            self.decoder = decoder
            self.request = request
            self.device = device
            self.dtype = dtype

            bsz = len(request.x)
            y_len: int = cast(tuple[int, ...], request.prompts.shape)[-1]
            self.bsz = bsz
            self.y_len = y_len

            # Cache
            self.kv_cache: MutableSequence[KVCacheProtocol]
            self.sample = sample_func()

            # Forward args
            self.x = [i.astype(mx.int32) for i in request.x]
            self.x_lens = request.x_lens.astype(mx.int32)
            self.y = mx.zeros((bsz, decoder.max_seq_length)).astype(mx.int32)
            self.y[:, : cast(tuple[int, ...], request.prompts.shape)[-1]] = request.prompts.astype(mx.int32)
            self.bert_feature = [i.astype(dtype) for i in request.bert_feature]

            self.prefill_len = self.x_lens + cast(tuple[int, ...], request.prompts.shape)[1]

            self.input_pos = mx.zeros_like(self.prefill_len)
            self.input_pos += self.prefill_len

            # EOS
            self.completed = mx.array([False] * len(self.x)).astype(mx.bool_)
            self.y_results: List[Array] = [None] * len(self.x)  # type: ignore

            self.xy_pos = decoder.embed(self.x, request.prompts, self.bert_feature)

            max_len = int(self.prefill_len.max(-1))
            attn_mask = mx.zeros(shape=(bsz, max_len, max_len), dtype=mx.bool_)

            for bs in range(bsz):
                pos = int(self.x_lens[bs])
                seq_len = pos + y_len

                attn_mask[bs, :seq_len, :pos] = True

                ar_mask = ~mx.triu(
                    x=mx.ones(
                        shape=(
                            y_len,
                            y_len,
                        ),
                        dtype=mx.bool_,
                    ),
                    k=1,
                )
                attn_mask[bs, pos:seq_len, pos:seq_len] = ar_mask

            attn_mask = mx.repeat(mx.expand_dims(attn_mask, 1), decoder.n_head, 1)
            self.attn_mask = attn_mask

            mx.eval(self.attn_mask)