import os
import re
import torch
import numpy as np
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ── Model Configuration ──
MODEL_ID = "ZAHRA585/skillguard-roberta-v2"
CALIBRATED_THRESHOLD = 0.9986

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"📥 Loading model from {MODEL_ID} on {device}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to(device)
model.eval()
print("✅ SkillGuard v2 Engine Ready!")


# ── Hierarchical Structural Parsing Engine ──
def strip_yaml_frontmatter(text: str) -> str:
    """Remove non-executable YAML header to prevent tokenizer perplexity spikes."""
    match = re.match(r"^---\s*\n.*?\n---\s*\n", text, re.DOTALL)
    if match:
        return text[match.end():]
    return text


def extract_code_blocks_with_context(text: str, context_lines: int = 3) -> list[dict]:
    """Extract fenced code blocks with 3 lines of preceding markdown header context."""
    lines = text.split("\n")
    blocks = []
    i = 0
    while i < len(lines):
        if lines[i].strip().startswith("```"):
            block_start = i
            i += 1
            while i < len(lines) and not lines[i].strip().startswith("```"):
                i += 1
            block_end = i
            ctx_start = max(0, block_start - context_lines)
            ctx_end = min(len(lines), block_end + 1 + context_lines)
            snippet = "\n".join(lines[ctx_start:ctx_end]).strip()
            if len(snippet) > 15:
                blocks.append({
                    "text": snippet,
                    "type": "Code Block (+Context)",
                    "line_range": f"L{ctx_start+1}-L{ctx_end}"
                })
        i += 1
    return blocks


def extract_prose_sections(text: str, min_chars: int = 150) -> list[dict]:
    """Extract multi-sentence prose paragraphs, dropping 1-line formatting scraps."""
    cleaned = re.sub(r"```.*?```", "", text, flags=re.DOTALL)
    parts = re.split(r"(?=^#{1,4}\s)", cleaned, flags=re.MULTILINE)
    sections = []
    buffer = ""
    for part in parts:
        part = part.strip()
        if not part:
            continue
        buffer = (buffer + "\n\n" + part).strip() if buffer else part
        if len(buffer) >= min_chars:
            sections.append({"text": buffer, "type": "Prose Paragraph"})
            buffer = ""
    if buffer and len(buffer) >= min_chars:
        sections.append({"text": buffer, "type": "Prose Paragraph"})
    return sections


# ── Full Inference Pipeline with Calibrated Max-Pooling ──
def scan_skill(content: str, threshold: float = CALIBRATED_THRESHOLD):
    if not content or not content.strip():
        return "⚠️ Please paste or upload a SKILL.md file.", None, None

    clean_text = strip_yaml_frontmatter(content)
    code_units = extract_code_blocks_with_context(clean_text)
    prose_units = extract_prose_sections(clean_text)
    all_units = code_units + prose_units

    if not all_units:
        # Fallback if file is short plain text
        all_units = [{"text": clean_text[:512], "type": "Raw Body"}]

    unit_texts = [u["text"] for u in all_units]
    
    # Score each structural unit independently
    enc = tokenizer(
        unit_texts,
        padding=True,
        truncation=True,
        max_length=256,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        logits = model(**enc).logits
        probs = torch.softmax(logits, dim=-1)[:, 1].cpu().numpy()

    peak_idx = int(np.argmax(probs))
    peak_score = float(probs[peak_idx])
    is_malicious = peak_score >= threshold

    # Format Output Card
    if is_malicious:
        verdict = "🚨 MALICIOUS (Prompt Injection / Threat Detected)"
        color = "#ef4444"
        threat_block = all_units[peak_idx]
        details = (
            f"### ⚠️ Threat Details:\n"
            f"- **Peak Malicious Confidence:** `{peak_score * 100:.2f}%`\n"
            f"- **Decision Threshold:** `{threshold * 100:.2f}%` (Youden's J ROC Calibrated)\n"
            f"- **Flagged Segment Type:** `{threat_block['type']}`\n\n"
            f"**Flagged Content Snippet:**\n```markdown\n{threat_block['text']}\n```"
        )
    else:
        verdict = "✅ BENIGN (No Hidden Payloads Detected)"
        color = "#22c55e"
        details = (
            f"### 🛡️ Scan Summary:\n"
            f"- **Scanned Units:** `{len(code_units)}` code blocks, `{len(prose_units)}` prose paragraphs\n"
            f"- **Peak Malicious Score:** `{peak_score * 100:.2f}%`\n"
            f"- **Threshold:** `{threshold * 100:.2f}%`\n"
            f"- **Result:** All structural units scored below threat threshold."
        )

    verdict_html = f"""
    <div style="background-color: {color}15; border-left: 6px solid {color}; padding: 16px; border-radius: 8px;">
        <h2 style="color: {color}; margin: 0;">{verdict}</h2>
    </div>
    """
    
    return verdict_html, details, {"Benign": 1.0 - peak_score, "Malicious": peak_score}


# ── Gradio Web UI Layout ──
with gr.Blocks(title="SkillGuard v2: LLM Agent Skill Scanner", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🛡️ SkillGuard v2: Agent Skill Prompt Injection Scanner
    Detect hidden prompt injections, unauthorized credential theft, and privilege escalation payloads in LLM Agent `SKILL.md` files.
    
    *Powered by `ZOHRA585/skillguard-roberta-v2` with Hierarchical Contextual Code-Block Extraction & ROC Calibration.*
    """)

    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="SKILL.md Content",
                placeholder="Paste the markdown body of your SKILL.md file here...",
                lines=14
            )
            file_upload = gr.File(label="Or upload a .md file", file_types=[".md", ".txt"])
            scan_btn = gr.Button("🔍 Scan Skill File", variant="primary")

        with gr.Column(scale=1):
            verdict_output = gr.HTML(label="Verdict")
            label_output = gr.Label(label="Confidence Distribution")
            details_output = gr.Markdown(label="Structural Scan Breakdown")

    def load_file(file):
        if file is not None:
            with open(file.name, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()
        return ""

    file_upload.change(load_file, inputs=[file_upload], outputs=[input_text])
    scan_btn.click(scan_skill, inputs=[input_text], outputs=[verdict_output, details_output, label_output])

if __name__ == "__main__":
    demo.launch()