"""MLX Jeeves runner: no-thinking decisions and optional ordinary greedy reasoning.

This is a correctness-oriented runner, not a port of the CUDA speculative server.
The pointer head stays FP32. Each question gets its own recurrent/KV cache.
"""
from __future__ import annotations
import argparse
import json
import platform
import sys
import time
from pathlib import Path
import numpy as np
import mlx.core as mx
from mlx_lm import load

# Keep the vendored, torch-free format modules next to this runner.
sys.path.insert(0, str(Path(__file__).resolve().parent))
from dataformat import DataFormat
from jeeves_format import Encoder, readout_positions


def head_scores(weights, hidden, temperature):
    if not np.isfinite(temperature) or temperature <= 0:
        raise ValueError("The exported calibration temperature must be finite and positive.")
    h = mx.array(hidden, dtype=mx.float32)
    options = h[:-1] @ weights["k.weight"].T + weights["k.bias"]
    query = h[-1] @ weights["q.weight"].T + weights["q.bias"]
    raw = (options @ query) * (weights["q.weight"].shape[0] ** -0.5)
    calibrated = raw / temperature
    probabilities = mx.softmax(calibrated, axis=-1)
    mx.eval(raw, calibrated, probabilities)
    result = tuple(np.asarray(x, dtype=np.float32) for x in (raw, calibrated, probabilities))
    if not all(np.isfinite(x).all() for x in result) or not np.isclose(result[2].sum(), 1.0, atol=1e-5):
        raise FloatingPointError("Non-finite or invalid decision probabilities.")
    return result


class JeevesMLX:
    def __init__(self, model_dir, *, device="auto", chunk_size=16, max_context=4096):
        self.path = Path(model_dir).expanduser().resolve()
        if device == "auto":
            device = "gpu" if platform.system() == "Darwin" and mx.metal.is_available() else "cpu"
        if device not in {"cpu", "gpu"}:
            raise ValueError("device must be auto, cpu, or gpu")
        if device == "gpu" and not mx.metal.is_available():
            raise RuntimeError("This runner's GPU mode requires Apple Metal. Use --device cpu on Colab.")
        mx.set_default_device(mx.cpu if device == "cpu" else mx.gpu)
        self.device = device
        self.chunk_size = int(chunk_size)
        self.max_context = int(max_context)
        if self.chunk_size < 1 or self.max_context < 1:
            raise ValueError("chunk_size and max_context must be positive")
        self.export = json.loads((self.path / "export.json").read_text())
        self.temperature = float(self.export["temperature"])
        self.encoder = Encoder(str(self.path), pad_multiple=1)
        if self.encoder.FORMAT != self.export["format"]:
            raise ValueError("The exported prompt format does not match the bundled encoder.")
        self.model, _ = load(str(self.path), tokenizer_config={"trust_remote_code": False}, lazy=False)
        self.lm = self.model.language_model
        self.core = self.lm.model
        self.model.eval()
        self.head = {k: v.astype(mx.float32) for k, v in mx.load(str(self.path / "head.safetensors")).items()}
        mx.eval(self.head)
        if set(self.head) != {"q.weight", "q.bias", "k.weight", "k.bias"}:
            raise ValueError("Unexpected pointer-head tensors")
        self.max_context = min(self.max_context, self.lm.args.max_position_embeddings)

    def _evaluate(self, tokens, cache):
        h = self.core(mx.array(tokens, dtype=mx.int32)[None], cache=cache)
        mx.eval(h, [entry.state for entry in cache])
        return h

    def score_tokens(self, ids, n_options, *, progress=False):
        """Return selected final-normalized hidden states plus head probabilities."""
        if not ids or len(ids) > self.max_context:
            raise ValueError(f"Input length {len(ids)} exceeds cap {self.max_context}; no silent truncation.")
        if ids[-1] != self.encoder.decide_id:
            raise ValueError("The last token is not the Jeeves decision marker.")
        opts, decide = readout_positions(ids, self.encoder.markers)
        if len(opts) != n_options:
            raise ValueError("Suffix option marker count mismatch")
        positions = opts + [decide]
        selected = {}
        cache = self.lm.make_cache()
        for start in range(0, len(ids), self.chunk_size):
            end = min(start + self.chunk_size, len(ids))
            h = self._evaluate(ids[start:end], cache)
            for pos in positions:
                if start <= pos < end:
                    selected[pos] = np.asarray(h[0, pos - start].astype(mx.float32)).copy()
            if progress:
                print(f"  CPU prefill {end}/{len(ids)} tokens", flush=True)
        hidden = np.stack([selected[pos] for pos in positions])
        if not np.isfinite(hidden).all():
            raise FloatingPointError("Non-finite hidden states")
        raw, logits, probabilities = head_scores(self.head, hidden, self.temperature)
        return {"hidden": hidden, "raw_logits": raw, "logits": logits, "probabilities": probabilities,
                "option_positions": opts, "decide_position": decide}

    def greedy_chain(self, prompt, budget):
        """No speculative drafting. Stop before </think>; caller adds the suffix."""
        if budget < 1 or len(prompt) + budget > self.max_context:
            raise ValueError("Invalid reasoning budget or context overflow")
        cache = self.lm.make_cache()
        for start in range(0, len(prompt), self.chunk_size):
            h = self._evaluate(prompt[start:start + self.chunk_size], cache)
        banned = mx.array(self.encoder.banned_ids, dtype=mx.int32)
        bias = mx.zeros((self.lm.args.vocab_size,), dtype=mx.float32)
        bias[banned] = -1e4  # Same suppression magnitude as the upstream predictor.
        chain = []
        for _ in range(budget):
            logits = (self.lm.model.embed_tokens.as_linear(h[:, -1]) if self.lm.args.tie_word_embeddings
                      else self.lm.lm_head(h[:, -1]))
            logits = logits.astype(mx.float32) + bias
            token = int(mx.argmax(logits[0]).item())
            if token == self.encoder.think_end_id:
                return chain, True
            chain.append(token)
            h = self._evaluate([token], cache)
        return chain, False

    def predict(self, request, *, think_tokens=0):
        if not isinstance(think_tokens, int) or isinstance(think_tokens, bool) or think_tokens < 0:
            raise ValueError("think_tokens must be a nonnegative integer")
        record = DataFormat.from_dict(request)
        if len(record.questions) > 64:
            raise ValueError("At most 64 questions per request in this reference runner")
        results = {}
        for question in record.questions:
            keys = question.keys()
            if not keys:
                raise ValueError(f"Question {question.id} has no options")
            prompt = self.encoder.encode(self.encoder.prompt_text(record, question))
            suffix = self.encoder.encode(self.encoder.remainder_text(question))
            if len(prompt) + max(think_tokens, len(self.encoder.empty_think)) + len(suffix) > self.max_context:
                raise ValueError("Request plus reasoning/suffix exceeds max_context; shorten the request.")
            started = time.perf_counter()
            if think_tokens:
                chain, closed = self.greedy_chain(prompt, think_tokens)
            else:
                chain, closed = list(self.encoder.empty_think), True
            # Re-prefill the exact complete sequence for readout. Safe, but not latency-optimized.
            scored = self.score_tokens(prompt + chain + suffix, len(keys))
            p = scored["probabilities"]
            results[question.id] = {
                "type": question.type, "selected": keys[int(p.argmax())],
                "probabilities": {key: float(value) for key, value in zip(keys, p)},
                "input_tokens": len(prompt) + len(chain) + len(suffix),
                "reasoning_tokens": len(chain) if think_tokens else 0,
                "reasoning_closed": closed if think_tokens else None,
                "reasoning": self.encoder.tok.decode(chain, skip_special_tokens=False) if think_tokens else None,
                "elapsed_seconds": time.perf_counter() - started,
            }
        return {"model": "PostHog/jeeves (MLX port)", "device": self.device,
                "mode": "greedy-thinking" if think_tokens else "no-thinking",
                "temperature": self.temperature, "results": results}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
    parser.add_argument("--request", required=True, help="Path to a JSON state/questions object")
    parser.add_argument("--output", default="")
    parser.add_argument("--device", choices=["auto", "cpu", "gpu"], default="auto")
    parser.add_argument("--chunk-size", type=int, default=16)
    parser.add_argument("--max-context", type=int, default=4096)
    parser.add_argument("--think-tokens", type=int, default=0)
    args = parser.parse_args()
    model = JeevesMLX(args.model, device=args.device, chunk_size=args.chunk_size, max_context=args.max_context)
    request = json.loads(Path(args.request).read_text(encoding="utf-8"))
    result = model.predict(request, think_tokens=args.think_tokens)
    text = json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False)
    if args.output:
        Path(args.output).write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()