import os
import time
from pathlib import Path


def _load_env(path: str = ".env"):
    p = Path(path)
    if not p.exists():
        return
    for line in p.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip())


_load_env()

import gradio as gr
from huggingface_hub import hf_hub_download

try:
    import spaces
    gpu = spaces.GPU
except Exception:
    def gpu(*args, **kwargs):
        def deco(fn):
            return fn
        return deco

MODEL_REPO = os.environ.get("MODEL_REPO", "AfkaraLP/rustlean-gguf")
MODEL_FILE = os.environ.get("MODEL_FILE", "rustlean-v4.Q8_0.gguf")
CTX = int(os.environ.get("RUSTLEAN_CTX", "8192"))
N_GPU_LAYERS = int(os.environ.get("N_GPU_LAYERS", "-1"))

FIM_PREFIX = "<|fim_prefix|>"
FIM_SUFFIX = "<|fim_suffix|>"
FIM_MIDDLE = "<|fim_middle|>"
CURSOR = "<|cursor|>"
STOPS = ["<|endoftext|>", FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE, "<|fim_pad|>"]

DEFAULT_CODE = """use std::collections::HashMap;

fn word_count(text: &str) -> HashMap<&str, u32> {
    let mut counts = HashMap::new();
<|cursor|>
    counts
}
"""

EXAMPLES = [
    ("""fn fibonacci(n: u32) -> u32 {
    match n {
        0 => 0,
        1 => 1,
<|cursor|>
    }
}
""", 0.2, 96, 0.9),
    ("""struct Point { x: f64, y: f64 }

impl Point {
    fn distance(&self, other: &Point) -> f64 {
<|cursor|>
    }
}
""", 0.2, 96, 0.9),
    ("""fn is_prime(n: u64) -> bool {
    if n < 2 {
        return false;
    }
<|cursor|>
}
""", 0.2, 128, 0.9),
]

_llm = None

_model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    token=os.environ.get("HF_TOKEN"),
)


def _preload_cuda():
    import ctypes
    import glob
    import site

    dirs = []
    for sp in site.getsitepackages():
        dirs += glob.glob(sp + "/nvidia/*/lib")
    for d in sorted(dirs):
        for so in sorted(glob.glob(d + "/*.so*")):
            try:
                ctypes.CDLL(so, mode=ctypes.RTLD_GLOBAL)
            except OSError:
                pass


def get_llm():
    global _llm
    if _llm is None:
        _preload_cuda()
        from llama_cpp import Llama
        _llm = Llama(
            model_path=_model_path,
            n_gpu_layers=N_GPU_LAYERS,
            n_ctx=CTX,
            n_batch=512,
            verbose=False,
            logits_all=False,
        )
    return _llm


def build_prompt(prefix: str, suffix: str) -> str:
    return f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}"


def clean(text: str) -> str:
    for s in [FIM_MIDDLE, FIM_SUFFIX, FIM_PREFIX, "<|endoftext|>", "<|fim_pad|>"]:
        text = text.replace(s, "")
    return text.rstrip()


def _complete_impl(code, temperature, max_tokens, top_p):
    if code.strip() == "":
        return "", code, "*Empty input.*"
    if CURSOR in code:
        prefix, suffix = code.split(CURSOR, 1)
    else:
        prefix, suffix = code, ""
    prompt = build_prompt(prefix, suffix)

    def run():
        llm = get_llm()
        return llm.create_completion(
            prompt=prompt,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            top_p=float(top_p),
            stop=STOPS,
            stream=False,
        )

    t0 = time.time()
    try:
        out = run()
    except Exception:
        global _llm
        _llm = None
        out = run()
    dt = time.time() - t0

    text = clean(out["choices"][0]["text"])
    assembled = f"{prefix}{text}{suffix}"
    n_tok = out.get("usage", {}).get("completion_tokens", 0)
    tps = (n_tok / dt) if dt > 0 else 0.0
    mode = "prefix" if suffix == "" else "FIM (prefix+suffix)"
    stats = (
        f"**{n_tok}** tokens · **{dt:.2f}s** · **{tps:.1f} tok/s**"
        f"  ·  mode: {mode}"
        f"  ·  gpu_layers: {N_GPU_LAYERS}"
    )
    return text, assembled, stats


@gpu(duration=45)
def complete(code, temperature, max_tokens, top_p):
    import traceback
    try:
        return _complete_impl(code, temperature, max_tokens, top_p)
    except Exception:
        return "", code, "```\n" + traceback.format_exc() + "\n```"


with gr.Blocks(title="RustLean v4 FIM Playground") as demo:
    gr.Markdown(
        "# RustLean v4: Rust FIM Completion\n"
        "Fill-in-the-middle completion from "
        "[`AfkaraLP/rustlean-gguf`](https://huggingface.co/AfkaraLP/rustlean-gguf) "
        "(Qwen2.5-Coder-1.5B, Q8_0). The v4 release improves repository-disjoint "
        "Rust completion while retaining the base model's compiler-tested pass rate. "
        "Put **`<|cursor|>`** where you want the model to "
        "write code. Everything before it is the prefix, everything after is the suffix."
    )

    with gr.Row():
        with gr.Column(scale=1):
            code_in = gr.Code(
                value=DEFAULT_CODE,
                language=None,
                label="Rust source (mark the hole with <|cursor|>)",
                interactive=True,
            )
            with gr.Accordion("Sampling", open=False):
                temperature = gr.Slider(minimum=0.0, maximum=1.5, value=0.2, step=0.05, label="temperature")
                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="top_p")
                max_tokens = gr.Slider(minimum=16, maximum=256, value=96, step=8, label="max tokens")
            with gr.Row():
                btn = gr.Button("Complete at cursor", variant="primary")
                clear = gr.Button("Reset")
        with gr.Column(scale=1):
            stats = gr.Markdown("_Not run yet._")
            completion = gr.Code(
                label="Generated middle",
                language=None,
                interactive=False,
            )
            result = gr.Code(
                label="Assembled result (prefix + middle + suffix)",
                language=None,
                interactive=True,
            )

    gr.Examples(examples=[list(e) for e in EXAMPLES], inputs=[code_in, temperature, max_tokens, top_p])

    btn.click(
        complete,
        inputs=[code_in, temperature, max_tokens, top_p],
        outputs=[completion, result, stats],
    )
    clear.click(lambda: (DEFAULT_CODE, "", "", "_Reset._"), None,
                outputs=[code_in, completion, result, stats])

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())