from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import gradio as gr
import os

# Download model if not cached
model_path = hf_hub_download(
    "BugTraceAI/BugTraceAI-CORE-Fast",
    "bugtraceai-core-fast.gguf"
)

print(f"Model path: {model_path}")
print(f"Model size: {os.path.getsize(model_path) / 1e9:.1f} GB")

llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_batch=512,
    n_gpu_layers=-1,         # all layers to GPU
    flash_attn=True,          # flash attention — lower VRAM for KV cache
    use_mmap=False,           # load directly into RAM (faster init)
    use_mlock=True,           # lock in RAM, prevent swap
    chat_format="chatml",
    verbose=False,
)

print(f"Model loaded successfully. GPU layers: {llm.n_gpu_layers}")

def chat(prompt, max_tokens=512, temperature=0.7):
    response = llm.create_chat_completion(
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=0.9,
    )
    return response["choices"][0]["message"]["content"]

demo = gr.Interface(
    fn=chat,
    inputs=[
        gr.Textbox(label="Prompt", lines=5, placeholder="Enter your security analysis prompt..."),
        gr.Slider(64, 1024, value=512, label="Max tokens"),
        gr.Slider(0.0, 1.5, value=0.7, label="Temperature"),
    ],
    outputs=gr.Textbox(label="Response", lines=10),
    title="BugTraceAI Fast 7B — Bug Bounty Analysis",
    description="Security-focused LLM for vulnerability analysis, exploit reasoning, and finding triage.",
)

demo.launch(server_name="0.0.0.0", server_port=7860)