#!/usr/bin/env python3
"""Self-contained CPU/MPS greedy inference for the Fleck-S-100K Instruct export.

This file intentionally has no import from the Fleck-LM source tree.  The model
architecture, chat template, and SafeTensors contract are reproduced here so
this directory can be copied from Hugging Face and used on its own.
"""
from __future__ import annotations

import argparse
import math
from collections.abc import Iterable
from pathlib import Path
from typing import cast

import torch
from safetensors import safe_open
from safetensors.torch import load_file
from torch import Tensor, nn

from tokenizers import Tokenizer

VOCAB_SIZE = 1024
CONTEXT_LENGTH = 2048
HIDDEN_SIZE = 64
INTERMEDIATE_SIZE = 128
PHYSICAL_BLOCKS = 2
EFFECTIVE_DEPTH = 4
QUERY_HEADS = 4
KV_HEADS = 2
HEAD_DIM = 16
EMBEDDING_RANK = 32
ROPE_THETA = 10000.0
BOS_ID, EOS_ID, PAD_ID, UNK_ID = 0, 1, 2, 3
SYSTEM_ID, USER_ID, ASSISTANT_ID, EOT_ID = 4, 5, 6, 7
EXECUTION_ORDER = (0, 1, 0, 1)
TOKENIZER_NAME = "Fleck-Tokenizer-1024"
MODEL_METADATA = {
    "candidate_id": "Fleck-S-100K",
    "canonical_dtype": "bfloat16",
    "context_length": "2048",
    "effective_depth": "4",
    "embedding_rank": "32",
    "head_dim": "16",
    "kv_heads": "2",
    "logical_execution_order": '["A", "B", "A", "B"]',
    "parameter_count": "109384",
    "physical_blocks": "2",
    "query_heads": "4",
    "tokenizer_name": TOKENIZER_NAME,
}


class RMSNorm(nn.Module):
    def __init__(self, size: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, hidden: Tensor) -> Tensor:
        variance = hidden.float().pow(2).mean(dim=-1, keepdim=True)
        scale = torch.rsqrt(variance + self.eps).to(dtype=hidden.dtype)
        return hidden * scale * self.weight


class Attention(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.q_proj = nn.Linear(HIDDEN_SIZE, QUERY_HEADS * HEAD_DIM, bias=False)
        self.k_proj = nn.Linear(HIDDEN_SIZE, KV_HEADS * HEAD_DIM, bias=False)
        self.v_proj = nn.Linear(HIDDEN_SIZE, KV_HEADS * HEAD_DIM, bias=False)
        self.o_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)

    @staticmethod
    def _rope(value: Tensor, positions: Tensor) -> Tensor:
        half = HEAD_DIM // 2
        frequencies = torch.arange(half, device=value.device, dtype=torch.float32)
        frequencies = ROPE_THETA ** (-2 * frequencies / HEAD_DIM)
        angles = positions.float().unsqueeze(-1) * frequencies
        cos = angles.cos().to(dtype=value.dtype)[None, None, :, :]
        sin = angles.sin().to(dtype=value.dtype)[None, None, :, :]
        first, second = value[..., :half], value[..., half:]
        return torch.cat((first * cos - second * sin, first * sin + second * cos), dim=-1)

    def forward(
        self,
        hidden: Tensor,
        positions: Tensor,
        past: tuple[Tensor, Tensor] | None = None,
        attention_mask: Tensor | None = None,
    ) -> tuple[Tensor, tuple[Tensor, Tensor]]:
        batch, length, _ = hidden.shape
        query = self.q_proj(hidden).view(batch, length, QUERY_HEADS, HEAD_DIM).transpose(1, 2)
        key = self.k_proj(hidden).view(batch, length, KV_HEADS, HEAD_DIM).transpose(1, 2)
        value = self.v_proj(hidden).view(batch, length, KV_HEADS, HEAD_DIM).transpose(1, 2)
        query = self._rope(query, positions)
        key = self._rope(key, positions)
        past_length = 0 if past is None else past[0].shape[2]
        if past is not None:
            key = torch.cat((past[0], key), dim=2)
            value = torch.cat((past[1], value), dim=2)
        total_length = key.shape[2]
        repeats = QUERY_HEADS // KV_HEADS
        expanded_key = key.repeat_interleave(repeats, dim=1)
        expanded_value = value.repeat_interleave(repeats, dim=1)
        scores = torch.matmul(query.float(), expanded_key.float().transpose(-1, -2))
        scores = scores / math.sqrt(HEAD_DIM)
        query_positions = torch.arange(
            past_length, past_length + length, device=hidden.device
        )[:, None]
        key_positions = torch.arange(total_length, device=hidden.device)[None, :]
        causal = key_positions <= query_positions
        scores = scores.masked_fill(
            ~causal[None, None, :, :], torch.finfo(torch.float32).min
        )
        if attention_mask is not None:
            if attention_mask.shape != (batch, total_length):
                raise ValueError("attention_mask must have shape (batch, complete_kv_length)")
            scores = scores.masked_fill(
                ~attention_mask.to(torch.bool)[:, None, None, :], torch.finfo(torch.float32).min
            )
        probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32)
        attended = torch.matmul(probabilities, expanded_value.float()).to(hidden.dtype)
        attended = attended.transpose(1, 2).contiguous().view(batch, length, HIDDEN_SIZE)
        return self.o_proj(attended), (key, value)


class PhysicalBlock(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.attention = Attention()
        self.mlp_gate = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.mlp_up = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.mlp_down = nn.Linear(INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False)


class FleckModel(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.token_embedding = nn.Parameter(torch.empty(VOCAB_SIZE, EMBEDDING_RANK))
        self.embedding_projection = nn.Parameter(torch.empty(EMBEDDING_RANK, HIDDEN_SIZE))
        self.blocks = nn.ModuleList(PhysicalBlock() for _ in range(PHYSICAL_BLOCKS))
        self.attention_norms = nn.ModuleList(
            RMSNorm(HIDDEN_SIZE) for _ in range(EFFECTIVE_DEPTH)
        )
        self.mlp_norms = nn.ModuleList(RMSNorm(HIDDEN_SIZE) for _ in range(EFFECTIVE_DEPTH))
        self.depth_embeddings = nn.Parameter(torch.empty(EFFECTIVE_DEPTH, HIDDEN_SIZE))
        self.attention_residual_scales = nn.Parameter(torch.empty(EFFECTIVE_DEPTH))
        self.mlp_residual_scales = nn.Parameter(torch.empty(EFFECTIVE_DEPTH))
        self.final_norm = RMSNorm(HIDDEN_SIZE)

    def _embed(self, input_ids: Tensor) -> Tensor:
        return torch.nn.functional.embedding(input_ids, self.token_embedding) @ self.embedding_projection

    def _logits(self, hidden: Tensor) -> Tensor:
        rank_hidden = hidden.float() @ self.embedding_projection.float().t()
        return rank_hidden @ self.token_embedding.float().t()

    def forward(
        self,
        input_ids: Tensor,
        *,
        past_key_values: tuple[tuple[Tensor, Tensor], ...] | None = None,
        attention_mask: Tensor | None = None,
        use_cache: bool = False,
    ) -> tuple[Tensor, tuple[tuple[Tensor, Tensor], ...] | None]:
        if input_ids.ndim != 2 or input_ids.shape[1] == 0:
            raise ValueError("input_ids must have shape (batch, non-empty sequence)")
        batch, length = input_ids.shape
        if past_key_values is not None and len(past_key_values) != EFFECTIVE_DEPTH:
            raise ValueError("past_key_values must contain one cache per effective depth")
        past_length = 0 if past_key_values is None else past_key_values[0][0].shape[2]
        if past_length + length > CONTEXT_LENGTH:
            raise ValueError(f"sequence exceeds context_length={CONTEXT_LENGTH}")
        positions = torch.arange(past_length, past_length + length, device=input_ids.device)
        hidden = self._embed(input_ids)
        caches: list[tuple[Tensor, Tensor]] = []
        for depth, physical_index in enumerate(EXECUTION_ORDER):
            hidden = hidden + self.depth_embeddings[depth].view(1, 1, -1)
            block = cast(PhysicalBlock, self.blocks[physical_index])
            attention_input = self.attention_norms[depth](hidden)
            past = None if past_key_values is None else past_key_values[depth]
            attended, cache = block.attention(attention_input, positions, past, attention_mask)
            hidden = hidden + self.attention_residual_scales[depth] * attended
            mlp_input = self.mlp_norms[depth](hidden)
            mlp_output = block.mlp_down(
                torch.nn.functional.silu(block.mlp_gate(mlp_input)) * block.mlp_up(mlp_input)
            )
            hidden = hidden + self.mlp_residual_scales[depth] * mlp_output
            caches.append(cache)
        logits = self._logits(self.final_norm(hidden))
        return logits, tuple(caches) if use_cache else None


def _expected_shapes(model: nn.Module) -> dict[str, tuple[int, ...]]:
    return {name: tuple(parameter.shape) for name, parameter in model.named_parameters()}


def load_model(checkpoint: str | Path, device: str = "cpu") -> FleckModel:
    """Load the public BF16 artifact strictly, then run it as FP32."""
    if device not in {"cpu", "mps"}:
        raise ValueError("device must be 'cpu' or 'mps'")
    if device == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS was requested but is not available")
    source = Path(checkpoint)
    if not source.is_file():
        raise FileNotFoundError(source)
    model = FleckModel()
    expected = _expected_shapes(model)
    with safe_open(str(source), framework="pt", device="cpu") as handle:
        metadata = handle.metadata() or {}
        for key, expected_value in MODEL_METADATA.items():
            if metadata.get(key) != expected_value:
                raise ValueError(
                    f"SafeTensors metadata mismatch for {key}: "
                    f"expected {expected_value!r}, got {metadata.get(key)!r}"
                )
        names = set(handle.keys())
        if names != set(expected):
            raise ValueError(
                "SafeTensors parameter names mismatch: "
                f"missing={sorted(set(expected) - names)}, extra={sorted(names - set(expected))}"
            )
        for name, shape in expected.items():
            tensor_slice = handle.get_slice(name)
            if tensor_slice.get_dtype() != "BF16":
                raise ValueError(f"{name}: expected BF16, got {tensor_slice.get_dtype()}")
            if tuple(tensor_slice.get_shape()) != shape:
                raise ValueError(
                    f"{name}: expected shape {shape}, got {tuple(tensor_slice.get_shape())}"
                )
    tensors = load_file(str(source), device="cpu")
    for name, shape in expected.items():
        tensor = tensors[name]
        if tensor.dtype != torch.bfloat16 or tuple(tensor.shape) != shape:
            raise ValueError(f"{name}: SafeTensors dtype/shape contract mismatch")
    model.load_state_dict(tensors, strict=True, assign=True)
    model = model.to(device=torch.device(device), dtype=torch.bfloat16).eval()
    if any(parameter.dtype != torch.bfloat16 for parameter in model.parameters()):
        raise TypeError("model parameters must be BF16 after loading")
    return model


def load_tokenizer(path: str | Path) -> Tokenizer:
    source = Path(path)
    tokenizer = Tokenizer.from_file(str(source))
    if tokenizer.get_vocab_size() != VOCAB_SIZE:
        raise ValueError(f"tokenizer vocabulary must be {VOCAB_SIZE}")
    expected = {
        "<bos>": BOS_ID,
        "<eos>": EOS_ID,
        "<pad>": PAD_ID,
        "<unk>": UNK_ID,
        "<|system|>": SYSTEM_ID,
        "<|user|>": USER_ID,
        "<|assistant|>": ASSISTANT_ID,
        "<|eot|>": EOT_ID,
    }
    if set(tokenizer.get_added_tokens_decoder()) != set(expected.values()):
        raise ValueError("tokenizer added-token set does not match the strict contract")
    for token, token_id in expected.items():
        added_token = tokenizer.get_added_tokens_decoder().get(token_id)
        if (
            tokenizer.token_to_id(token) != token_id
            or added_token is None
            or getattr(added_token, "content", None) != token
            or not bool(getattr(added_token, "special", False))
        ):
            raise ValueError(f"tokenizer special token contract mismatch for {token}")
    return tokenizer


@torch.inference_mode()
def generate(
    model: FleckModel,
    input_ids: Tensor,
    max_tokens: int,
    stop_token_ids: Iterable[int] = (EOS_ID,),
) -> Tensor:
    if max_tokens < 0:
        raise ValueError("max_tokens must be non-negative")
    if input_ids.shape[1] + max_tokens > CONTEXT_LENGTH:
        raise ValueError(f"prompt plus generation exceeds context_length={CONTEXT_LENGTH}")
    stop_ids = set(stop_token_ids)
    generated = input_ids
    cache: tuple[tuple[Tensor, Tensor], ...] | None = None
    for _ in range(max_tokens):
        current = generated if cache is None else generated[:, -1:]
        logits, cache = model(current, past_key_values=cache, use_cache=True)
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        generated = torch.cat((generated, next_token), dim=1)
        if bool(torch.all(torch.isin(next_token, torch.tensor(list(stop_ids), device=next_token.device)))):
            break
    return generated


def _default_path(name: str) -> str:
    return str(Path(__file__).with_name(name))


def chat_ids(tokenizer: Tokenizer, user_text: str) -> list[int]:
    """Encode one training-format user turn and open the assistant turn."""
    return [BOS_ID, USER_ID, *tokenizer.encode(user_text, add_special_tokens=False).ids, EOT_ID, ASSISTANT_ID]


def decode_reply(tokenizer: Tokenizer, token_ids: list[int]) -> str:
    """Stop at EOT/EOS and decode only ordinary text."""
    for index, token_id in enumerate(token_ids):
        if token_id in {EOT_ID, EOS_ID}:
            token_ids = token_ids[:index]
            break
    return tokenizer.decode(token_ids, skip_special_tokens=True).strip()


def run_single(args: argparse.Namespace, tokenizer: Tokenizer, model: FleckModel) -> None:
    prompt = args.prompt if args.prompt is not None else "Hello"
    ids = chat_ids(tokenizer, prompt)
    input_ids = torch.tensor([ids], dtype=torch.long, device=args.device)
    output = generate(model, input_ids, args.max_tokens, (EOT_ID, EOS_ID))
    reply = decode_reply(tokenizer, output[0, len(ids) :].detach().cpu().tolist())
    print(reply)


def run_chat(args: argparse.Namespace, tokenizer: Tokenizer, model: FleckModel) -> None:
    history: list[tuple[str, str]] = []
    print("Fleck chat. Type /exit to quit or /clear to reset history.")
    while True:
        try:
            prompt = input("user> ")
        except (EOFError, KeyboardInterrupt):
            print()
            return
        if prompt.strip() == "/exit":
            return
        if prompt.strip() == "/clear":
            history.clear()
            print("(history cleared)")
            continue
        if not prompt.strip():
            continue
        ids = [BOS_ID]
        for old_user, old_reply in history:
            ids.extend([USER_ID, *tokenizer.encode(old_user, add_special_tokens=False).ids, EOT_ID, ASSISTANT_ID])
            ids.extend([*tokenizer.encode(old_reply, add_special_tokens=False).ids, EOT_ID])
        ids.extend([USER_ID, *tokenizer.encode(prompt, add_special_tokens=False).ids, EOT_ID, ASSISTANT_ID])
        if len(ids) + args.max_tokens > CONTEXT_LENGTH:
            history.clear()
            ids = chat_ids(tokenizer, prompt)
        input_ids = torch.tensor([ids], dtype=torch.long, device=args.device)
        output = generate(model, input_ids, args.max_tokens, (EOT_ID, EOS_ID))
        reply = decode_reply(tokenizer, output[0, len(ids) :].detach().cpu().tolist())
        print(f"assistant> {reply}")
        history.append((prompt, reply))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ckpt", default=_default_path("model.safetensors"))
    parser.add_argument("--tokenizer", default=_default_path("tokenizer.json"))
    parser.add_argument("--prompt", default=None, help="single prompt; omit for interactive chat")
    parser.add_argument("--max-tokens", type=int, default=32)
    parser.add_argument("--device", choices=("cpu", "mps"), default="cpu")
    parser.add_argument("--no-chat", action="store_true", help="run one prompt instead of interactive chat")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    tokenizer = load_tokenizer(args.tokenizer)
    model = load_model(args.ckpt, args.device)
    if args.no_chat or args.prompt is not None:
        run_single(args, tokenizer, model)
    else:
        run_chat(args, tokenizer, model)


if __name__ == "__main__":
    main()