"""
Phase N - Interactive multi-turn chat REPL for The Abyss-1 (SFT Checkpoint)

Builds The Abyss chat template on every turn, streams tokesn as they
generate, and stops cleanly at <|endofturn|> or <|endoftext|>

Usage (from the_final_product/):
    python -m ml.chat
    python -m ml.chat --checkpoint ml/checkpoints/sft_ckpt.pt
    python -m ml.chat --temperature 0.7 --max_tokens 512

Commands during chat:
    reset    - clear conversation history
    quit     - exit
"""

import argparse
import sys
import torch
import torch.nn.functional as F

from ml.model import TheAbyssModel
from ml.tokenizer.tokenizer import AbyssTokenizer

TOKENIZER_DIR = "ml/tokenizer/vocab"
SYSTEM_TEXT   = "You are The Abyss, and empty void of nothingness"


@torch.no_grad()
def respond(model, tok, history, sys_ids, cfg,
            max_new_tokens, temperature, top_k, top_p,
            EOS, EOT, USR, AST, device):
    """
    Build the full prompt from history, sample until EOT/EOS,
    stream each token to stdout. Returns list of response token ids.
    """
    # ----- BUILD PROMPT -----
    prompt = list(sys_ids)
    for turn in history:
        prompt += turn
    prompt += [AST]

    # If prompt is too long, drop oldest turns (keep system + recent turns)
    max_prompt = cfg.block_size - max_new_tokens - 1
    if len(prompt) > max_prompt:
        prompt = list(sys_ids)
        tail   = []
        for turn in reversed(history):
            if len(sys_ids) + len(turn) + len(tail) + 1 > max_prompt:
                break
            tail = list(turn) + tail
        prompt = list(sys_ids) + tail + [AST]

    idx = torch.tensor([prompt], dtype=torch.long, device=device)

    # ----- SAMPLE -----
    response_ids = []
    for _ in range(max_new_tokens):
        idx_cond     = idx[:, -cfg.block_size:]
        logits, _    = model(idx_cond)
        logits       = logits[0, -1, :] / temperature   # (V,)  — 1D

        # top_k
        if top_k > 0:
            k      = min(top_k, logits.size(-1))
            kth    = torch.topk(logits, k).values[-1]
            logits = logits.masked_fill(logits < kth, float("-inf"))

        # top_p (nucleus)
        if 0.0 < top_p < 1.0:
            s_logits, s_idx = torch.sort(logits, descending=True)
            cum  = torch.cumsum(F.softmax(s_logits, dim=-1), dim=-1)
            mask = (cum - F.softmax(s_logits, dim=-1)) > top_p
            s_logits = s_logits.masked_fill(mask, float("-inf"))
            logits   = torch.zeros_like(logits).scatter_(0, s_idx, s_logits)

        next_tok = torch.multinomial(F.softmax(logits, dim=-1), 1)  # (1,)
        tok_id   = next_tok.item()

        # STOP ON END-OF-TURN or EOS
        if tok_id == EOT or tok_id == EOS:
            break

        response_ids.append(tok_id)
        sys.stdout.write(tok.decode([tok_id]))
        sys.stdout.flush()
        idx = torch.cat([idx, next_tok.view(1, 1)], dim=1)

    print()
    return response_ids


def main():
    parser = argparse.ArgumentParser(description="Chat with The Abyss-1")
    parser.add_argument("--checkpoint",  default="ml/checkpoints/sft_ckpt.pt")
    parser.add_argument("--max_tokens",  type=int,   default=256)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top_k",       type=int,   default=50)
    parser.add_argument("--top_p",       type=float, default=0.95)
    parser.add_argument("--device",
                        default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    # ----- LOAD MODEL -----
    print(f"Loading {args.checkpoint} ...", flush=True)
    ck   = torch.load(args.checkpoint, map_location=args.device, weights_only=False)
    cfg  = ck["config"]
    model = TheAbyssModel(cfg).to(args.device)
    model.load_state_dict(ck["model"])
    model.eval()
    print(f"  iter={ck.get('iter', '?')}  params={model.num_params():,}\n", flush=True)

    # ----- LOAD TOKENIZER -----
    tok = AbyssTokenizer.load(TOKENIZER_DIR)

    EOS = tok.encode("<|endoftext|>")[0]
    USR = tok.encode("<|user|>")[0]
    AST = tok.encode("<|assistant|>")[0]
    EOT = tok.encode("<|endofturn|>")[0]
    SYS = tok.encode("<|system|>")[0]

    sys_ids = [SYS] + tok.encode(SYSTEM_TEXT) + [EOT]

    print("The Abyss Chat - type 'reset' to clear history, 'quit' to exit\n")
    print("-" *50)

    history = []

    while True:
        try:
            user_input = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            break

        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit"):
            print("Bye.")
            break
        if user_input.lower() == "reset":
            history.clear()
            print("  [history cleared]")
            continue

        # ENCODE AND STORE USER TURN
        user_turn = [USR] + tok.encode(user_input) + [EOT]
        history.append(user_turn)

        print("\nThe Abyss: ", end="", flush=True)
        response_ids = respond(
            model, tok, history, sys_ids, cfg,
            args.max_tokens, args.temperature, args.top_k, args.top_p,
            EOS, EOT, USR, AST, args.device
        )

        # STORE ASSISTANT TURN
        history.append([AST] + response_ids + [EOT])


if __name__ == "__main__":
    main()