import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import spaces 

model_id = "Willie999/trapSTAR-gemma4"

print("Initializing tokenizer and loading base configurations...")
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
)
model.eval()

# Dynamic matrix placeholder states to drive the processing animations
AWAITING_STATE = "*Awaiting input code compilation... Press 'Run Autonomous Security Audit' to begin analysis.*"
LOADING_STATE = """
<div class='matrix-loader'>
    <div class='pulse-bar'></div>
    <p class='scan-text'>🛰️ SYSTEM: INFERENCE AGENT ACTIVE — SCANNING IN PROCESS...</p>
</div>
"""

@spaces.GPU(duration=90) 
def analyze_and_patch(source_code):
    if not source_code.strip():
        return "### ⚠️ System Warning\nPlease input a valid source code routine to evaluate."
    
    messages = [
        {
            "role": "system", 
            "content": "You are Trap Star, an autonomous defensive security auditing agent. Analyze the provided code snippet, identify the vulnerability type, and write out structural recommendations."
        },
        {
            "role": "user", 
            "content": f"Review this function block for potential vulnerabilities:\n\n```\n{source_code}\n```"
        }
    ]
    
    try:
        prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
        tokenized_inputs = tokenizer(prompt_text, return_tensors="pt")
        
        input_ids = tokenized_inputs.input_ids.to(model.device)
        attention_mask = tokenized_inputs.attention_mask.to(model.device)
        
        with torch.no_grad():
            generated_ids = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=1536,
                temperature=0.2,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
            
        response_tokens = generated_ids[0][input_ids.shape[-1]:]
        response_text = tokenizer.decode(response_tokens, skip_special_tokens=True)
        return response_text
    except Exception as e:
        return f"### ❌ Operational Failure\nAn unexpected error occurred during dynamic GPU allocation:\n`{str(e)}`"

# Custom CSS for a professional, distraction-free DevSecOps interface
custom_css = """
footer {visibility: hidden !important}
.gradio-container {background-color: #070a0e !important;}

/* Custom Textbox Layout styling */
.terminal-input textarea {
    font-family: 'Fira Code', 'Courier New', monospace !important;
    background-color: #0d1117 !important;
    color: #c9d1d9 !important;
    border: 1px solid #30363d !important;
}

/* Custom Live Response Animations */
.matrix-loader {
    padding: 20px;
    background: #0d1117;
    border: 1px solid #238636;
    border-radius: 6px;
    position: relative;
    overflow: hidden;
}
.pulse-bar {
    height: 3px;
    width: 100%;
    background: linear-gradient(90deg, transparent, #238636, transparent);
    position: absolute;
    top: 0;
    left: -100%;
    animation: scan 2s linear infinite;
}
.scan-text {
    color: #2da44e !important;
    font-family: 'Fira Code', monospace;
    font-size: 14px;
    margin: 0;
    animation: pulse 1.5s ease-in-out infinite;
}

@keyframes scan {
    0% { left: -100%; }
    100% { left: 100%; }
}
@keyframes pulse {
    0%, 100% { opacity: 0.6; }
    50% { opacity: 1; }
}
"""

with gr.Blocks(title="trapSTAR Engine", css=custom_css) as demo:
    with gr.Row():
        gr.Markdown(
            """
            # 🛡️ trapSTAR-gemma4 Core Engine
            ### *Autonomous Defensive Code Auditing & Patch Remediation Dashboard*
            ---
            """
        )
        
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("#### 💻 Target Source Code Extraction")
            code_input = gr.Textbox(
                label="Source Code Input",
                placeholder="Paste your source file or target code routine block here...",
                lines=16,
                max_lines=25,
                elem_classes=["terminal-input"]
            )
            submit_btn = gr.Button("⚡ Run Autonomous Security Audit", variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown("#### 🛰️ Trap Star Assessment Dashboard")
            
            # State panels that swap cleanly to display active CSS animations
            with gr.Group():
                # HTML animation panel (Visible strictly while loading)
                loader_panel = gr.HTML(visible=False)
                # Production markdown panel (Visible when output is ready)
                patch_output = gr.Markdown(value=AWAITING_STATE, visible=True)
    
    gr.Examples(
        examples=[
            ["void process_str(char *str) {\n    char buffer[16];\n    strcpy(buffer, str);\n}"],
            ["import sqlite3\n\ndef get_user(user_id):\n    conn = sqlite3.connect('users.db')\n    cursor = conn.cursor()\n    cursor.execute(f'SELECT * FROM accounts WHERE id = {user_id}')\n    return cursor.fetchall()"]
        ],
        inputs=code_input,
        label="Quick Test Blueprints"
    )
    
    # 3. Clean manual API Documentation panel overlay for wide audiences
    with gr.Accordion("🔌 System Integration API Endpoints", open=False):
        gr.Markdown(
            """
            ### External REST API Access
            Because this application runs behind a ZeroGPU cluster proxy, standard client integrations must hit the API endpoint routing channel explicitly.
            
            ```python
            from gradio_client import Client

            # Connect to your public deployment endpoint
            client = Client("Willie999/trapSTAR-interface")

            result = client.predict(
                source_code="YOUR_RAW_CODE_STRING_HERE",
                api_name="/audit"
            )
            print(result)
            ```
            """
        )

    # UI Event Chain handling animations and switching component visibilities
    def pre_load_ui():
        # Instantly hides static text block and showcases running CSS animations
        return gr.update(value=LOADING_STATE, visible=True), gr.update(visible=False)

    def post_load_ui(final_output):
        # Tears down the animation elements and renders raw markdown data cleanly
        return gr.update(visible=False), gr.update(value=final_output, visible=True)

    submit_btn.click(
        fn=pre_load_ui,
        outputs=[loader_panel, patch_output],
        queue=False
    ).then(
        fn=analyze_and_patch,
        inputs=code_input,
        outputs=patch_output,
        api_name="audit"
    ).then(
        fn=post_load_ui,
        inputs=patch_output,
        outputs=[loader_panel, patch_output],
        queue=False
    )

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