from collections.abc import Generator
from dataclasses import dataclass
import os
from typing import Any

from loguru import logger
import mlx.core as mx
from mlx_lm.generate import GenerationResponse, stream_generate
from mlx_lm.models.cache import can_trim_prompt_cache, make_prompt_cache
from mlx_lm.sample_utils import make_logits_processors, make_sampler
from mlx_lm.utils import load
from outlines.processors import JSONLogitsProcessor

from ..utils.debug_logging import log_debug_chat_template
from ..utils.outlines_transformer_tokenizer import OutlinesTransformerTokenizer

DEFAULT_TEMPERATURE = float(os.getenv("DEFAULT_TEMPERATURE", "0.7"))
DEFAULT_TOP_P = float(os.getenv("DEFAULT_TOP_P", "0.95"))
DEFAULT_TOP_K = int(os.getenv("DEFAULT_TOP_K", "20"))
DEFAULT_MIN_P = float(os.getenv("DEFAULT_MIN_P", "0.0"))
DEFAULT_XTC_PROBABILITY = float(os.getenv("DEFAULT_XTC_PROBABILITY", "0.0"))
DEFAULT_XTC_THRESHOLD = float(os.getenv("DEFAULT_XTC_THRESHOLD", "0.0"))
DEFAULT_SEED = int(os.getenv("DEFAULT_SEED", "0"))
DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "1000000"))
DEFAULT_REPETITION_CONTEXT_SIZE = int(os.getenv("DEFAULT_REPETITION_CONTEXT_SIZE", "20"))


@dataclass
class CompletionResponse:
    """
    The output of :func:`__call__` when stream is False.

    Args:
        text (str): The next segment of decoded text. This can be an empty string.
        tokens (List[int]): The list of tokens in the response.
        peak_memory (float): The peak memory used so far in GB.
        generation_tps (float): The tokens-per-second for generation.
        generation_tokens (int): The number of generated tokens.
        prompt_tps (float): The prompt processing tokens-per-second.
        prompt_tokens (int): The number of tokens in the prompt.
    """

    text: str = None
    tokens: list[int] = None
    peak_memory: float = None
    generation_tps: float = None
    prompt_tps: float = None
    prompt_tokens: int = None
    generation_tokens: int = None


class MLX_LM:
    """
    A wrapper class for MLX Language Model that handles both streaming and non-streaming inference.

    This class provides a unified interface for generating text responses from text prompts,
    supporting both streaming and non-streaming modes.
    """

    def __init__(
        self,
        model_path: str,
        draft_model_path: str = None,
        num_draft_tokens: int = 2,
        context_length: int | None = None,
        trust_remote_code: bool = False,
        chat_template_file: str = None,
        debug: bool = False,
    ):
        try:
            self.model, self.tokenizer = load(
                model_path, lazy=False, tokenizer_config={"trust_remote_code": trust_remote_code}
            )
            self.context_length = context_length
            self.draft_model = None
            self.draft_tokenizer = None
            self.num_draft_tokens = num_draft_tokens
            if draft_model_path:
                self._load_draft_model(draft_model_path, trust_remote_code)
            self.pad_token_id = self.tokenizer.pad_token_id
            self.bos_token = self.tokenizer.bos_token
            self.model_type = self.model.model_type
            self.debug = debug
            self.outlines_tokenizer = OutlinesTransformerTokenizer(self.tokenizer)
            initial_cache = make_prompt_cache(self.model)
            self._cache_is_trimmable = can_trim_prompt_cache(initial_cache)
            self._num_model_cache_layers = len(initial_cache)
            if chat_template_file:
                if not os.path.exists(chat_template_file):
                    raise ValueError(f"Chat template file {chat_template_file} does not exist")
                with open(chat_template_file) as f:
                    template_content = f.read()
                    self.tokenizer.chat_template = template_content
                if self.debug:
                    log_debug_chat_template(
                        chat_template_file=chat_template_file, template_content=template_content
                    )
        except Exception as e:
            raise ValueError(f"Error loading model: {e!s}")

    def _load_draft_model(self, draft_model_path: str, trust_remote_code: bool) -> None:
        try:
            self.draft_model, self.draft_tokenizer = load(
                draft_model_path,
                lazy=False,
                tokenizer_config={"trust_remote_code": trust_remote_code},
            )
            self.context_length = (
                None  # speculative decoding does not support context length, should be set to None
            )
            self._validate_draft_tokenizer()
        except Exception as e:
            raise ValueError(f"Error loading draft model: {e!s}")

    def _validate_draft_tokenizer(self) -> None:
        if self.draft_tokenizer.vocab_size != self.tokenizer.vocab_size:
            logger.warning(
                "Draft model tokenizer does not match model tokenizer. "
                "Speculative decoding may not work as expected."
            )

    def create_prompt_cache(self) -> list[Any]:
        cache = make_prompt_cache(self.model, max_kv_size=self.context_length)
        if self.draft_model:
            cache += make_prompt_cache(self.draft_model, max_kv_size=self.context_length)
        return cache

    def get_model_type(self) -> str:
        return self.model_type

    def create_input_prompt(
        self, messages: list[dict[str, str]], chat_template_kwargs: dict[str, Any]
    ) -> str:
        use_partial = chat_template_kwargs.pop("_partial_mode", False)

        return self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=not use_partial,
            continue_final_message=use_partial,
            **chat_template_kwargs,
        )

    def encode_prompt(self, input_prompt: str) -> list[int]:
        """Encode a prompt string into token IDs.

        Parameters
        ----------
        input_prompt : str
            The prompt string to encode.

        Returns
        -------
        list[int]
            Token IDs for the prompt.
        """
        return self.tokenizer.encode(input_prompt)

    @property
    def cache_is_trimmable(self) -> bool:
        """Whether the model's prompt cache supports trimming.

        Pure-attention models return ``True``; hybrid models with
        ``ArraysCache`` (SSM/recurrent layers) return ``False``.
        """
        return self._cache_is_trimmable

    def _prefill_cache(
        self,
        token_ids: list[int],
        prompt_cache: list[Any],
        prefill_step_size: int = 2048,
    ) -> None:
        """Process tokens through the model to warm up the prompt cache.

        Parameters
        ----------
        token_ids : list[int]
            Token IDs to prefill into the cache.
        prompt_cache : list[Any]
            Prompt cache to update in-place.
        prefill_step_size : int, optional
            Maximum chunk size per forward pass, by default 2048.
        """
        tokens = mx.array(token_ids)
        n_model = self._num_model_cache_layers
        model_cache = prompt_cache[:n_model]

        remaining = tokens
        while remaining.size > 0:
            chunk = remaining[:prefill_step_size]
            self.model(chunk[None], cache=model_cache)
            mx.eval([c.state for c in model_cache])
            remaining = remaining[prefill_step_size:]
            if remaining.size > 0:
                mx.clear_cache()

        if self.draft_model:
            draft_cache = prompt_cache[n_model:]
            remaining = tokens
            while remaining.size > 0:
                chunk = remaining[:prefill_step_size]
                self.draft_model(chunk[None], cache=draft_cache)
                mx.eval([c.state for c in draft_cache])
                remaining = remaining[prefill_step_size:]
                if remaining.size > 0:
                    mx.clear_cache()

    def __call__(
        self, input_ids: list[int], prompt_cache: list[Any] = None, stream: bool = False, **kwargs
    ) -> CompletionResponse | Generator[GenerationResponse, None, None]:
        """Generate text response from the model.

        Parameters
        ----------
        input_ids : list[int]
            Token IDs for the input prompt.
        prompt_cache : list[Any], optional
            Pre-computed prompt cache for faster inference.
        stream : bool, optional
            Whether to stream the response, by default ``False``.
        **kwargs
            Additional generation parameters (temperature, max_tokens, etc.)
            and optional checkpoint control parameters:

            - ``checkpoint_position`` (int | None): Token index at which to
              split prefill and save a cache checkpoint.
            - ``checkpoint_callback`` (callable | None): Called with the
              prompt cache after processing the prefix so the caller can
              persist a checkpoint.

        Returns
        -------
        CompletionResponse | Generator[GenerationResponse, None, None]
            Complete response or streaming generator.
        """
        checkpoint_position: int | None = kwargs.pop("checkpoint_position", None)
        checkpoint_callback = kwargs.pop("checkpoint_callback", None)

        if (
            checkpoint_position is not None
            and checkpoint_callback is not None
            and prompt_cache is not None
            and 0 < checkpoint_position < len(input_ids)
        ):
            self._prefill_cache(input_ids[:checkpoint_position], prompt_cache)
            checkpoint_callback(prompt_cache)
            input_ids = input_ids[checkpoint_position:]

        def _get(key, default):
            v = kwargs.get(key)
            return default if v is None else v

        seed = _get("seed", DEFAULT_SEED)
        max_tokens = _get("max_tokens", None)
        if max_tokens is None:
            max_tokens = _get("max_completion_tokens", DEFAULT_MAX_TOKENS)
        sampler_kwargs = {
            "temp": _get("temperature", DEFAULT_TEMPERATURE),
            "top_p": _get("top_p", DEFAULT_TOP_P),
            "top_k": _get("top_k", DEFAULT_TOP_K),
            "min_p": _get("min_p", DEFAULT_MIN_P),
            "xtc_probability": _get("xtc_probability", DEFAULT_XTC_PROBABILITY),
            "xtc_threshold": _get("xtc_threshold", DEFAULT_XTC_THRESHOLD),
        }

        # Add XTC special tokens (EOS and newline) when XTC is enabled
        if sampler_kwargs["xtc_probability"] > 0:
            sampler_kwargs["xtc_special_tokens"] = [
                self.tokenizer.eos_token_id,
                *self.tokenizer.encode("\n"),
            ]

        repetition_penalty = kwargs.get("repetition_penalty")
        repetition_context_size = _get("repetition_context_size", DEFAULT_REPETITION_CONTEXT_SIZE)
        logit_bias = kwargs.get("logit_bias")

        # Convert string keys to int if logit_bias is provided (OpenAI API uses string keys)
        if logit_bias:
            logit_bias = {int(k): v for k, v in logit_bias.items()}

        logits_processors = make_logits_processors(
            logit_bias=logit_bias,
            repetition_penalty=repetition_penalty,
            repetition_context_size=repetition_context_size,
        )

        json_schema = kwargs.get("schema")
        if json_schema:
            logits_processors.append(
                JSONLogitsProcessor(
                    schema=json_schema, tokenizer=self.outlines_tokenizer, tensor_library_name="mlx"
                )
            )

        # Only seed RNG when an explicit non-negative seed is provided
        # None or negative values (e.g., -1) result in non-deterministic generation
        if seed and seed >= 0:
            mx.random.seed(seed)

        prompt_progress_callback = kwargs.get("prompt_progress_callback")

        sampler = make_sampler(**sampler_kwargs)

        stream_response = stream_generate(
            self.model,
            self.tokenizer,
            input_ids,
            draft_model=self.draft_model,
            sampler=sampler,
            max_tokens=max_tokens,
            num_draft_tokens=self.num_draft_tokens,
            prompt_cache=prompt_cache,
            logits_processors=logits_processors,
            prompt_progress_callback=prompt_progress_callback,
        )
        if stream:
            return stream_response

        text = ""
        tokens = []
        final_chunk = None
        for chunk in stream_response:
            text += chunk.text
            tokens.append(chunk.token)
            if chunk.finish_reason:
                final_chunk = chunk

        return CompletionResponse(
            text=text,
            tokens=tokens,
            peak_memory=final_chunk.peak_memory,
            generation_tps=final_chunk.generation_tps,
            prompt_tps=final_chunk.prompt_tps,
            prompt_tokens=final_chunk.prompt_tokens,
            generation_tokens=final_chunk.generation_tokens,
        )
