"""DeepHat-V1 (WhiteRabbitNeo V3), Cisco Foundation-Sec-8B-Reasoning & SecGPT ZeroGPU Studio."""

import os

os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")

import spaces  # Must be imported before torch
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

PRIMARY_MODEL_ID = "DeepHat/DeepHat-V1-7B"
MODEL_CHOICES = [
    "DeepHat/DeepHat-V1-7B (WhiteRabbitNeo V3 · Best for Code, Pentest & DevOps)",
    "fdtn-ai/Foundation-Sec-8B-Reasoning (Cisco Foundation AI · Deep Security Reasoning)",
    "clouditera/SecGPT-7B (Clouditera SecGPT)",
]

DEFAULT_SYSTEM_PROMPT = (
    "You are DeepHat (WhiteRabbitNeo V3), an advanced cybersecurity, DevOps, and source-code security engineer. "
    "Provide deep, technical, unfiltered architectural analysis, code audits, vulnerability explanations, "
    "and defensive remediation engineering in clear English."
)

print(f"[secgpt-studio] Preloading {PRIMARY_MODEL_ID} at module scope for ZeroGPU...", flush=True)
primary_tokenizer = AutoTokenizer.from_pretrained(PRIMARY_MODEL_ID)
primary_model = AutoModelForCausalLM.from_pretrained(
    PRIMARY_MODEL_ID,
    torch_dtype=torch.bfloat16,
)
primary_model.to("cuda")
primary_model.eval()
print(f"[secgpt-studio] {PRIMARY_MODEL_ID} ready on ZeroGPU!", flush=True)

_AUX_CACHE = {}


def get_model_and_tokenizer(model_choice: str):
    """Return the selected model and tokenizer."""
    repo_id = (model_choice or PRIMARY_MODEL_ID).split()[0]
    if repo_id == PRIMARY_MODEL_ID:
        return primary_model, primary_tokenizer
    if repo_id not in _AUX_CACHE:
        print(f"[secgpt-studio] Loading {repo_id} on ZeroGPU...", flush=True)
        tok = AutoTokenizer.from_pretrained(repo_id)
        mod = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
        mod.to("cuda")
        mod.eval()
        _AUX_CACHE[repo_id] = (mod, tok)
    return _AUX_CACHE[repo_id]


@spaces.GPU(duration=75)
def chat_with_secgpt(
    prompt: str,
    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
    max_tokens: int = 1024,
    temperature: float = 0.6,
    model_choice: str = MODEL_CHOICES[0],
) -> str:
    """Send a prompt to DeepHat-V1 / Foundation-Sec / SecGPT on ZeroGPU and return the response."""
    if not prompt or not prompt.strip():
        return "Please enter a prompt."

    model, tokenizer = get_model_and_tokenizer(model_choice)
    sys_msg = (system_prompt or DEFAULT_SYSTEM_PROMPT).strip()
    messages = [
        {"role": "system", "content": sys_msg},
        {"role": "user", "content": prompt.strip()},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer([text], return_tensors="pt").to("cuda")

    gen_kwargs = {
        "max_new_tokens": int(max_tokens),
        "repetition_penalty": 1.12,
        "do_sample": float(temperature) > 0.05,
    }
    if gen_kwargs["do_sample"]:
        gen_kwargs["temperature"] = float(temperature)
        gen_kwargs["top_p"] = 0.9

    with torch.no_grad():
        generated_ids = model.generate(**inputs, **gen_kwargs)

    new_tokens = [
        output_ids[len(input_ids):]
        for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
    ]
    response = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]
    return response.strip()


with gr.Blocks(title="DeepHat-V1 & Foundation-Sec Cybersecurity Studio", fill_width=True) as demo:
    gr.Markdown(
        "# 🛡️ DeepHat-V1 (WhiteRabbitNeo V3) + Cisco Foundation-Sec Studio\n"
        "Upgraded to **[DeepHat/DeepHat-V1-7B](https://huggingface.co/DeepHat/DeepHat-V1-7B)** "
        "(WhiteRabbitNeo V3 / `Qwen2.5-Coder-7B` fine-tune) and **[fdtn-ai/Foundation-Sec-8B-Reasoning](https://huggingface.co/fdtn-ai/Foundation-Sec-8B-Reasoning)** on **ZeroGPU (`zero-a10g`)** at **$0/hour**."
    )
    with gr.Row():
        with gr.Column(scale=1):
            model_dropdown = gr.Dropdown(
                choices=MODEL_CHOICES,
                value=MODEL_CHOICES[0],
                label="Active Security Model",
            )
            prompt_box = gr.Textbox(
                label="Security / Code Audit / DevOps Prompt",
                lines=6,
                placeholder="Paste code, HTTP headers, architecture specs, or security engineering questions...",
            )
            with gr.Accordion("Model & Sampling Settings", open=False):
                sys_box = gr.Textbox(
                    label="System Prompt",
                    value=DEFAULT_SYSTEM_PROMPT,
                    lines=3,
                )
                max_tok_slider = gr.Slider(
                    64, 2048, value=1024, step=64, label="Max New Tokens"
                )
                temp_slider = gr.Slider(
                    0.0, 1.2, value=0.6, step=0.05, label="Temperature"
                )
            ask_btn = gr.Button("🛡️ Run Security Analysis (ZeroGPU)", variant="primary")
        with gr.Column(scale=1):
            output_box = gr.Markdown(label="Model Response")

    ask_btn.click(
        fn=chat_with_secgpt,
        inputs=[prompt_box, sys_box, max_tok_slider, temp_slider, model_dropdown],
        outputs=output_box,
        api_name="chat",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Ocean(), mcp_server=True)