import os
import sys
import gradio as gr
import torch

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

WEIGHTS = os.environ.get("FELA_LLM_WEIGHTS", "..")
try:
    torch.set_num_threads(4)
    MODEL = load_model(WEIGHTS, threads=4)
except Exception as e:
    MODEL = None
    print(f"[Note] Model not loaded ({e}); the Space reports the input instead.")


def complete(prefix, suffix, max_tokens):
    if MODEL is None:
        return "Weights not found. Set FELA_LLM_WEIGHTS to the repo directory."
    if not (prefix or "").strip() and (not (suffix or "").strip()):
        return "Type some code before the cursor (prefix), and optionally after it (suffix)."
    r = MODEL.complete(prefix or "", suffix or "", max_tokens=int(max_tokens))
    tag = "fill in the middle" if r["used_fim"] else "continuation"
    return f"{r['middle']}\n\n[{tag}, {r['n_tokens']} tokens, {r['tok_per_s']} tok/s on CPU]"


with gr.Blocks(title="FELA LLM 1.5 code completion") as demo:
    gr.Markdown(
        "# FELA LLM 1.5 code completion\nA CPU native code autocomplete model. Type the code before your cursor in Prefix, and optionally the code after your cursor in Suffix, and it writes the line that goes in between. It runs on a plain CPU with no GPU. This is the final fill in the middle model: single line patterns land well, multi line blocks and novel logic are outside what it is built for. Every completion is a real forward of the model."
    )
    prefix = gr.Textbox(
        label="Prefix (code before the cursor)", value="def add(a, b):\n    ", lines=4
    )
    suffix = gr.Textbox(
        label="Suffix (code after the cursor, optional)",
        value="\nresult = add(2, 3)\n",
        lines=2,
    )
    max_tokens = gr.Slider(4, 64, value=16, step=1, label="Max tokens")
    out = gr.Textbox(label="Completed middle", lines=4)
    gr.Button("Complete").click(complete, [prefix, suffix, max_tokens], out)
if __name__ == "__main__":
    demo.launch()