import os
import torch
import gradio as gr
import spaces
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "Virtue-AI-HUB/VulnLLM-R-7B"

# 1. Tokenizer can be global because it only uses CPU RAM
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 2. This function ONLY boots the model when an API request comes in
@spaces.GPU(duration=120)
def generate_hacking_logic(prompt):
    # Quantization must happen INSIDE the GPU space
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4"
    )
    
    # Model loads dynamically, runs, and disappears to prevent crashes
    model = AutoModelForCausalLM.from_pretrained(
        model_id, 
        quantization_config=quantization_config,
        device_map="auto"
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=1024)
    return tokenizer.decode(outputs, skip_special_tokens=True)

# 3. Standard Gradio Interface to keep Hugging Face happy
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 VulnLLM-R-7B Backend Tunnel Bridge Running")
    gr.Markdown("FastAPI endpoint wrapper is actively listening for local BlackArch commands.")

app = demo.app 

# 4. The OpenAI-compatible API endpoint for HexStrike
@app.post("/v1/chat/completions")
async def openai_endpoint(request: Request):
    data = await request.json()
    messages = data.get("messages", [])
    user_prompt = messages[-1]["content"] if messages else ""
    
    # Triggers the dynamic GPU loading function
    ai_response = generate_hacking_logic(user_prompt)
    
    return JSONResponse(content={
        "choices": [{"message": {"role": "assistant", "content": ai_response}, "finish_reason": "stop"}]
    })

# 5. Launch using the proper Gradio loop required by the container platform
demo.launch(server_name="0.0.0.0", server_port=7860)