import argparse
import os
import sys
import torch

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))
from modeling import load_model


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--weights",
        default=os.environ.get("FELA_LLM_WEIGHTS", ".."),
        help="Dir with model.safetensors, config.json, tokenizer.json",
    )
    ap.add_argument("--threads", type=int, default=4)
    args = ap.parse_args()
    torch.set_num_threads(args.threads)
    print("Loading FELA LLM 1.5 on CPU (no GPU needed)...")
    m = load_model(args.weights, threads=args.threads)
    print(
        f"Loaded {m.cfg_json.get('n_params', 0) / 1000000000.0:.2f}B params, fim={m.fim_ok}\n"
    )
    print("Fill in the middle: prefix + suffix -> the model writes the middle")
    prefix = "def add(a, b):\n    "
    suffix = "\nresult = add(2, 3)\n"
    r = m.complete(prefix, suffix, max_tokens=16)
    print(f"  Prefix: {prefix!r}")
    print(f"  Suffix: {suffix!r}")
    print(
        f"  Middle: {r['middle']!r}  (fim={r['used_fim']}, {r['n_tokens']} tokens, {r['tok_per_s']} tok/s)\n"
    )
    print("Plain autocomplete: continue a single line")
    for prefix in ["import numpy as ", "from fastapi import ", "for i in range("]:
        r = m.complete(prefix, "", max_tokens=12)
        print(f"  {prefix!r} -> {r['middle']!r}")
    print(
        "\nThis is the final fill in the middle model. Single line completions land well; multi line blocks and novel logic are outside what it is built for. Every completion is a real forward of the model."
    )


if __name__ == "__main__":
    main()