import json
import os
from dataclasses import dataclass
from importlib.util import find_spec
from typing import Any, Callable
from urllib import request

from zerogpu import gpu


@dataclass(frozen=True)
class LocalChatClient:
    endpoint: str
    model: str
    timeout_seconds: int = 60
    temperature: float = 0.2
    max_tokens: int = 256
    enable_thinking: bool | None = None

    # Complete one chat prompt through an OpenAI-compatible local endpoint.
    def complete(self, system: str, user: str) -> str:
        payload = chat_payload(self.model, system, user, self.temperature, self.max_tokens, self.enable_thinking)
        req = request.Request(
            self.endpoint,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with request.urlopen(req, timeout=self.timeout_seconds) as response:
            return parse_chat_response(json.loads(response.read().decode("utf-8")))


@dataclass(frozen=True)
class LocalJsonChatClient:
    endpoint: str
    model: str
    timeout_seconds: int = 60
    temperature: float = 0.0
    max_tokens: int = 256

    # Complete one chat prompt through a JSON-constrained local endpoint.
    def complete(self, system: str, user: str) -> str:
        payload = json_chat_payload(self.model, system, user, self.temperature, self.max_tokens)
        req = request.Request(
            self.endpoint,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with request.urlopen(req, timeout=self.timeout_seconds) as response:
            return parse_chat_response(json.loads(response.read().decode("utf-8")))


ChatCompleter = LocalChatClient | LocalJsonChatClient


@dataclass(frozen=True)
class LocalCompletionClient:
    endpoint: str
    model: str
    prompt_template: Callable[[str, str], str]
    timeout_seconds: int = 60
    temperature: float = 0.2
    max_tokens: int = 256

    # Complete one prompt through an OpenAI-compatible local completion endpoint.
    def complete(self, system: str, user: str) -> str:
        prompt = self.prompt_template(system, user)
        payload = completion_payload(self.model, prompt, self.temperature, self.max_tokens)
        req = request.Request(
            self.endpoint,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with request.urlopen(req, timeout=self.timeout_seconds) as response:
            return parse_completion_response(json.loads(response.read().decode("utf-8")))


@dataclass(frozen=True)
class NemotronTransformersChatClient:
    model: Any
    tokenizer: Any
    max_new_tokens: int = 256
    temperature: float = 0.2

    # Set the active model/tokenizer globals, then run inference on ZeroGPU.
    # The @gpu worker reads the model from the global (inherited via fork) instead
    # of receiving it as an argument, which a model object cannot survive (pickle).
    def complete(self, system: str, user: str) -> str:
        global _nemotron_model, _nemotron_tokenizer
        _nemotron_model, _nemotron_tokenizer = self.model, self.tokenizer
        return _nemotron_generate(system, user, self.max_new_tokens, self.temperature)

    # Load Nemotron with its required tokenizer chat template (in the main process).
    @classmethod
    def load(
        cls,
        model_path: str,
        max_new_tokens: int = 256,
        temperature: float = 0.2,
    ) -> "NemotronTransformersChatClient":  # pragma: no cover
        from transformers import AutoModelForCausalLM, AutoTokenizer

        # Load on CPU (no device_map="auto"): on ZeroGPU the move to CUDA must
        # happen inside the @gpu call, in _nemotron_generate.
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto")
        return cls(model.eval(), tokenizer, max_new_tokens=max_new_tokens, temperature=temperature)


_nemotron_model: Any = None
_nemotron_tokenizer: Any = None


# Run Nemotron generation on a ZeroGPU allocation, reading the model from module
# globals so the forked GPU worker inherits it (only strings cross the boundary).
@gpu
def _nemotron_generate(system: str, user: str, max_new_tokens: int, temperature: float) -> str:
    _nemotron_model.to(best_device())
    messages = chat_messages(system, user)
    inputs = _nemotron_tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
    )
    outputs = _nemotron_model.generate(
        tensor_to_model_device(inputs, _nemotron_model),
        max_new_tokens=max_new_tokens,
        do_sample=temperature > 0,
        temperature=temperature,
    )
    return decode_generated_text(_nemotron_tokenizer, inputs, outputs)


@dataclass(frozen=True)
class MLXChatClient:
    model: Any
    tokenizer: Any
    prompt_template: Callable[[str, str], str]
    generate_func: Callable[..., str]
    max_tokens: int = 128

    # Complete one prompt through a local MLX model.
    def complete(self, system: str, user: str) -> str:
        prompt = self.prompt_template(system, user)
        return self.generate_func(self.model, self.tokenizer, prompt, verbose=False, max_tokens=self.max_tokens)

    # Load one MLX model for Apple Silicon inference.
    @classmethod
    def load(
        cls,
        model_path: str,
        prompt_template: Callable[[str, str], str],
        max_tokens: int = 128,
    ) -> "MLXChatClient":  # pragma: no cover
        from mlx_lm import generate, load

        model, tokenizer = load(model_path)
        return cls(model, tokenizer, prompt_template, generate, max_tokens=max_tokens)


@dataclass(frozen=True)
class MiniCPMTransformersChatClient:
    model: Any
    tokenizer: Any
    max_new_tokens: int = 512
    temperature: float = 0.7

    # Set the active model/tokenizer globals, then run inference on ZeroGPU.
    # Sampling is on by default so card authoring is not deterministically repetitive.
    def complete(self, system: str, user: str) -> str:
        global _minicpm_model, _minicpm_tokenizer
        _minicpm_model, _minicpm_tokenizer = self.model, self.tokenizer
        return _minicpm_generate(system, user, self.max_new_tokens, self.temperature)

    # Load MiniCPM with trust_remote_code for its custom chat method (main process).
    @classmethod
    def load(cls, model_path: str, max_new_tokens: int = 512, temperature: float = 0.7) -> "MiniCPMTransformersChatClient":  # pragma: no cover
        import torch
        from transformers import AutoModel, AutoTokenizer

        model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=True,
            attn_implementation="sdpa",
            torch_dtype=local_torch_dtype(torch),
        )
        # Stay on CPU here: on ZeroGPU the GPU only exists inside @gpu calls,
        # so the move to CUDA happens in _minicpm_generate.
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        return cls(model.eval(), tokenizer, max_new_tokens=max_new_tokens, temperature=temperature)


_minicpm_model: Any = None
_minicpm_tokenizer: Any = None


# Run MiniCPM's chat() on a ZeroGPU allocation, reading the model from module
# globals so the forked GPU worker inherits it (only strings cross the boundary).
@gpu
def _minicpm_generate(system: str, user: str, max_new_tokens: int, temperature: float) -> str:
    _minicpm_model.to(best_device())
    return str(
        _minicpm_model.chat(
            msgs=[{"role": "user", "content": user}],
            image=None,
            tokenizer=_minicpm_tokenizer,
            system_prompt=system,
            sampling=temperature > 0,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
        )
    )


# Build an OpenAI-compatible chat completion payload.
def chat_payload(
    model: str,
    system: str,
    user: str,
    temperature: float,
    max_tokens: int | None = None,
    enable_thinking: bool | None = None,
) -> dict[str, Any]:
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        "temperature": temperature,
    }
    if max_tokens is not None:
        payload["max_tokens"] = max_tokens
    if enable_thinking is not None:
        payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking}
    return payload


# Build an OpenAI-compatible JSON-constrained chat payload.
def json_chat_payload(model: str, system: str, user: str, temperature: float, max_tokens: int) -> dict[str, Any]:
    payload = chat_payload(model, system, user, temperature, max_tokens)
    payload["response_format"] = {"type": "json_object"}
    return payload


# Build an OpenAI-compatible text completion payload.
def completion_payload(model: str, prompt: str, temperature: float, max_tokens: int) -> dict[str, Any]:
    return {
        "model": model,
        "prompt": prompt,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }


# Build standard system/user chat messages.
def chat_messages(system: str, user: str) -> list[dict[str, str]]:
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


# Render Nemotron's documented single-turn prompt markers.
def nemotron_prompt(system: str, user: str) -> str:
    return f"<extra_id_0>System\n{system}\n\n<extra_id_1>User\n{user}\n<extra_id_1>Assistant\n"


# Render a text-only MiniCPM prompt for its model.chat API.
def minicpm_text_prompt(system: str, user: str) -> str:
    return f"System:\n{system}\n\nUser:\n{user}"


# Parse text content from an OpenAI-compatible chat response.
def parse_chat_response(raw: dict[str, Any]) -> str:
    return str(raw["choices"][0]["message"]["content"])


# Parse text content from an OpenAI-compatible completion response.
def parse_completion_response(raw: dict[str, Any]) -> str:
    if "content" in raw:
        return str(raw["content"])
    choice = raw["choices"][0]
    if "text" in choice:
        return str(choice["text"])
    return str(choice["message"]["content"])


# Move generated inputs onto the model device when tensors support it.
def tensor_to_model_device(inputs: Any, model: Any) -> Any:
    device = getattr(model, "device", None)
    if device is not None and hasattr(inputs, "to"):
        return inputs.to(device)
    return inputs


# Decode only tokens generated after the input prompt.
def decode_generated_text(tokenizer: Any, inputs: Any, outputs: Any) -> str:
    output = outputs[0]
    prompt_length = token_length(inputs)
    generated = output[prompt_length:]
    return str(tokenizer.decode(generated, skip_special_tokens=True)).strip()


# Return the final token dimension for tensor-like input ids.
def token_length(inputs: Any) -> int:
    if hasattr(inputs, "shape"):
        return int(inputs.shape[-1])
    if inputs and isinstance(inputs[0], list):
        return len(inputs[0])
    return len(inputs)


# Return practical local loading kwargs for causal Transformers models.
def transformers_model_kwargs() -> dict[str, Any]:
    kwargs: dict[str, Any] = {"torch_dtype": "auto"}
    if find_spec("accelerate") is not None:
        kwargs["device_map"] = "auto"
        kwargs["offload_folder"] = os.environ.get("TABRAS_MODEL_OFFLOAD", "/tmp/tabras-model-offload")
    return kwargs


# Return the best dtype available for MiniCPM local inference.
def local_torch_dtype(torch: Any) -> Any:
    if torch.cuda.is_available():
        return torch.bfloat16
    return "auto"


# Return the best available torch device string (CUDA on a Space/Linux GPU,
# MPS on Apple Silicon, else CPU), so inference works wherever it runs.
def best_device() -> str:
    import torch

    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return "mps"
    return "cpu"