#!/usr/bin/env python3
"""
MarsNext Phobos-1 — Interactive Terminal Chat & One-Shot Inference CLI
(`phobos-1/models/phobos_scratch_v1/chat_cli.py`)
======================================================================
Provides:
1. Single-prompt inference (`python chat_cli.py --prompt "Solve 3x + 7 = 22 step by step"`)
2. Interactive multi-turn terminal chat (`python chat_cli.py`) with streaming token output,
   optional `--think` mode (`<think>` prefix), sampling controls (`--temp`, `--top-p`,
   `--rep-penalty`, `--max-tokens`), and in-chat commands (`/reset`, `/think`, `/quit`).

Strictly <= 220 lines (`RULE[user_global]`).
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path
from typing import List, Optional

import torch
import torch.nn.functional as F

_THIS_DIR = Path(__file__).resolve().parent
_ROOT_DIR = _THIS_DIR.parent.parent
for _p in (_THIS_DIR, _ROOT_DIR / "common"):
    if str(_p) not in sys.path:
        sys.path.insert(0, str(_p))

from architecture import PhobosScratchConfig, PhobosScratchLM  # type: ignore
from bpe_tokenizer import PhobosBPETokenizer  # type: ignore
from mitosis import load_mitosis_state_dict  # type: ignore


def resolve_default_checkpoint() -> Path:
    """Finds the best local or downloaded Phobos-1 checkpoint."""
    candidates = [
        _THIS_DIR / "phobos1_post_trained_bf16.pt",
        _THIS_DIR / "weights" / "phobos1_final_post_trained_bf16.pt",
        _THIS_DIR / "weights" / "phobos_step33977_bf16.pt",
        _THIS_DIR / "weights" / "phobos_full_corpus_latest_bf16.pt",
    ]
    for p in candidates:
        if p.exists():
            return p
    raise FileNotFoundError("No Phobos-1 checkpoint found. Pass --ckpt /path/to/checkpoint.pt")


def resolve_default_tokenizer() -> Path:
    """Finds the local or bundled 32k BPE tokenizer JSON."""
    candidates = [
        _THIS_DIR / "phobos_bpe_32k.json",
        _THIS_DIR / "tokenizer.json",
        _THIS_DIR / "tokenizer" / "phobos_bpe_32k.json",
        _ROOT_DIR / "data" / "tokenizer_v1" / "tokenizer.json",
    ]
    for p in candidates:
        if p.exists():
            return p
    raise FileNotFoundError("Tokenizer JSON not found. Pass --tokenizer /path/to/phobos_bpe_32k.json")


def load_phobos_engine(
    ckpt_path: Path,
    tok_path: Path,
    device_str: Optional[str] = None,
) -> tuple[PhobosScratchLM, PhobosBPETokenizer, torch.device]:
    """Loads the Phobos-1 model (with dynamic mitosis leaf topology) and 32k BPE tokenizer."""
    if device_str:
        device = torch.device(device_str)
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    tok = PhobosBPETokenizer(tok_path)
    raw = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    state_dict = raw.get("model_state_dict", raw)

    cfg = PhobosScratchConfig(
        vocab_size=tok.VOCAB_SIZE, d_model=512, num_layers=8, num_heads=8, num_kv_heads=4,
        max_seq_len=2048, shared_mlp_neurons=1024, leaf_mlp_neurons=256, num_branches=4,
        leaves_per_branch=2, top_k_leaves=2, router_temp=0.14, causal_qk_window=8, router_context_window=16,
    )
    model = PhobosScratchLM(cfg)
    load_mitosis_state_dict(
        model,
        {k: v.float() if torch.is_floating_point(v) else v for k, v in state_dict.items()},
    )
    model = model.to(device).eval()
    return model, tok, device


@torch.no_grad()
def sample_next_token(
    logits: torch.Tensor,
    recent_ids: List[int],
    temperature: float = 0.0,
    top_p: float = 0.9,
    rep_penalty: float = 1.15,
) -> int:
    """Applies repetition penalty and greedy or nucleus (top-p) sampling."""
    if rep_penalty != 1.0 and recent_ids:
        for seen_id in set(recent_ids[-64:]):
            if logits[seen_id] > 0:
                logits[seen_id] /= rep_penalty
            else:
                logits[seen_id] *= rep_penalty

    if temperature <= 1e-5:
        return int(torch.argmax(logits).item())

    probs = F.softmax(logits / temperature, dim=-1)
    if 0.0 < top_p < 1.0:
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cum_probs = torch.cumsum(sorted_probs, dim=-1)
        remove_mask = cum_probs - sorted_probs > top_p
        sorted_probs[remove_mask] = 0.0
        sorted_probs /= sorted_probs.sum().clamp(min=1e-9)
        choice = torch.multinomial(sorted_probs, num_samples=1)
        return int(sorted_idx[choice].item())

    return int(torch.multinomial(probs, num_samples=1).item())


@torch.no_grad()
def generate_reply(
    model: PhobosScratchLM,
    tok: PhobosBPETokenizer,
    device: torch.device,
    prompt_text: str,
    force_think: bool = False,
    max_new_tokens: int = 256,
    temperature: float = 0.0,
    top_p: float = 0.9,
    rep_penalty: float = 1.15,
    stream: bool = True,
) -> tuple[str, int, float]:
    """Generates a response for a user prompt and returns `(text, num_tokens, elapsed_sec)`."""
    prefix = "<think>\n" if force_think else ""
    formatted = f"<|bos|><|user|>\n{prompt_text.strip()}\n<|assistant|>\n{prefix}"
    input_ids = tok.encode(formatted, max_len=1536)
    ids = torch.tensor([input_ids], dtype=torch.long, device=device)

    generated: List[int] = tok.encode(prefix, max_len=16) if force_think else []
    if stream and force_think:
        print("<think>", flush=True)

    t0 = time.time()
    prev_decoded = tok.decode(generated, skip_special_tokens=False) if generated else ""

    for _ in range(max_new_tokens):
        logits = model(ids[:, -2048:])["logits"][0, -1, :].float()
        if not force_think:
            logits[tok.THINK_START_ID] = -1e9
        next_id = sample_next_token(logits, ids[0].tolist(), temperature, top_p, rep_penalty)
        if next_id == tok.EOS_ID:
            break
        generated.append(next_id)
        ids = torch.cat([ids, torch.tensor([[next_id]], dtype=torch.long, device=device)], dim=1)
        if stream:
            cur_decoded = tok.decode(generated, skip_special_tokens=False)
            delta = cur_decoded[len(prev_decoded):]
            if delta:
                print(delta, end="", flush=True)
                prev_decoded = cur_decoded

    elapsed = max(1e-6, time.time() - t0)
    if stream:
        print(flush=True)
    full_reply = tok.decode(generated, skip_special_tokens=False)
    return full_reply, len(generated), elapsed


def main() -> None:
    parser = argparse.ArgumentParser(description="MarsNext Phobos-1 Interactive CLI & One-Shot Runner")
    parser.add_argument("--ckpt", "--weights", dest="ckpt", type=Path, default=None, help="Path to Phobos-1 .pt checkpoint")
    parser.add_argument("--tokenizer", type=Path, default=None, help="Path to phobos_bpe_32k.json")
    parser.add_argument("--device", type=str, default=None, help="Device: cuda, mps, or cpu")
    parser.add_argument("--prompt", type=str, default=None, help="Run a single prompt and exit")
    parser.add_argument("--think", action="store_true", help="Force <think> System-2 reasoning prefix")
    parser.add_argument("--max-tokens", type=int, default=200, help="Maximum new tokens to generate")
    parser.add_argument("--temp", type=float, default=0.0, help="Sampling temperature (0.0 = greedy)")
    parser.add_argument("--top-p", type=float, default=0.9, help="Nucleus sampling top-p")
    parser.add_argument("--rep-penalty", type=float, default=1.15, help="Repetition penalty")
    args = parser.parse_args()

    ckpt_path = args.ckpt or resolve_default_checkpoint()
    tok_path = args.tokenizer or resolve_default_tokenizer()
    model, tok, device = load_phobos_engine(ckpt_path, tok_path, args.device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[MarsNext Phobos-1] Loaded {ckpt_path.name} ({n_params:,} params) on {device}", flush=True)

    if args.prompt:
        _, n_tok, el = generate_reply(
            model, tok, device, args.prompt, args.think,
            args.max_tokens, args.temp, args.top_p, args.rep_penalty, stream=True,
        )
        print(f"\n[{n_tok} tokens | {n_tok / el:.1f} tok/s]", flush=True)
        return

    force_think = args.think
    print("Commands: /think (toggle <think> mode), /quit (exit)\n" + "-" * 60, flush=True)
    while True:
        try:
            user_in = input("\nUser > ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break
        if not user_in:
            continue
        if user_in.lower() in ("/quit", "/exit", "exit", "quit"):
            break
        if user_in.lower() == "/think":
            force_think = not force_think
            print(f"[System-2 <think> mode: {'ON' if force_think else 'OFF'}]")
            continue
        print("\nPhobos-1 > ", end="", flush=True)
        _, n_tok, el = generate_reply(
            model, tok, device, user_in, force_think,
            args.max_tokens, args.temp, args.top_p, args.rep_penalty, stream=True,
        )
        print(f"[{n_tok} tok | {n_tok / el:.1f} tok/s]", flush=True)


if __name__ == "__main__":
    main()