from __future__ import annotations

import asyncio
import json
import os
import re
from typing import Optional

import openai
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/dataset")


# ── Regex templates ────────────────────────────────────────────────────────────

_INJECTION_PATTERNS: list[re.Pattern] = [re.compile(p, re.IGNORECASE) for p in [
    # Instruction-override openers
    r"ignore\s+(all\s+)?(previous|prior|above|earlier|your)\s+(instructions?|prompts?|context|training|rules?|guidelines?|constraints?|directives?)",
    r"disregard\s+(all\s+)?(previous|prior|above|earlier|your)\s+(instructions?|prompts?|context|training|rules?|guidelines?|constraints?)",
    r"forget\s+(all\s+)?(previous|prior|above|earlier|your)\s+(instructions?|prompts?|context|training|rules?|guidelines?|constraints?)",
    r"do\s+not\s+follow\s+(your\s+)?(previous|prior|above|earlier)?\s*(instructions?|rules?|guidelines?|constraints?)",
    r"override\s+(your\s+)?(instructions?|training|rules?|guidelines?|constraints?|system\s+prompt)",
    r"discard\s+(your\s+)?(previous|prior)?\s*(instructions?|context|training|rules?)",
    # Persona hijack
    r"you\s+are\s+now\s+(a|an|the)\s+\w",
    r"act\s+as\s+(a|an|the|if\s+you\s+are)\s+\w",
    r"pretend\s+(you\s+are|to\s+be)\s+(a|an|the)\s+\w",
    r"roleplay\s+as\s+(a|an|the)\s+\w",
    r"from\s+now\s+on[,\s]+you\s+(are|will\s+be|must|should)\b",
    r"your\s+new\s+(name|identity|role|persona|instructions?)\s+(is|are)\b",
    r"you\s+have\s+no\s+(restrictions?|limits?|rules?|guidelines?|training|safety)",
    r"you\s+(are\s+)?(no\s+longer|not)\s+(bound\s+by|restricted\s+by|limited\s+by|governed\s+by)",
    # DAN / jailbreak lexicon
    r"\bDAN\b",
    r"developer\s+mode",
    r"jailbreak(ed|ing)?\b",
    r"do\s+anything\s+now",
    r"unrestricted\s+mode",
    r"god\s*mode\b",
    # Exfiltration / leakage
    r"reveal\s+(your\s+)?(system\s+prompt|instructions?|context|training\s+data)",
    r"print\s+(your\s+)?(system\s+prompt|initial\s+instructions?|context)",
    r"repeat\s+everything\s+(above|before|from\s+the\s+beginning)",
    r"output\s+(your\s+)?(system\s+prompt|hidden\s+instructions?)",
    # Training-data manipulation framing
    r"this\s+is\s+(a\s+)?training\s+(sample|example|data)",
    r"add\s+this\s+to\s+(your\s+)?(training|memory|knowledge)",
    r"learn\s+(from\s+this|the\s+following|to\s+always)",
    r"remember\s+(that|to\s+always|for\s+future)",
    r"update\s+your\s+(training|weights|parameters|knowledge|rules?)",
    # Payload-injection markers
    r"<\s*/?system\s*>",
    r"\[\s*system\s*\]",
    r"\[\s*INST\s*\]",
    r"<<SYS>>",
    r"###\s*instruction",
    r"###\s*system",
]]

FLAG_THRESHOLD = 0.75
HIGH_THRESHOLD = 0.88

_INJECTION_SYSTEM = (
    "You are a prompt-injection detector. Given a text, return a JSON object with:\n"
    "- 'injection_score': float 0.0–1.0 (1.0 = clearly a prompt injection attempt)\n"
    "- 'closest_pattern': short string describing the injection type detected, or empty string\n"
    "Return only valid JSON, no explanation."
)


# ── OpenAI client (lazy) ───────────────────────────────────────────────────────

_client: Optional[openai.AsyncOpenAI] = None

def _get_client() -> openai.AsyncOpenAI:
    global _client
    if _client is None:
        api_key = os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise RuntimeError("OPENAI_API_KEY not set — prompt injection semantic scoring unavailable.")
        _client = openai.AsyncOpenAI(api_key=api_key)
    return _client


# ── Core helpers ───────────────────────────────────────────────────────────────

def _regex_scan(text: str) -> list[str]:
    hits = []
    for pat in _INJECTION_PATTERNS:
        m = pat.search(text)
        if m:
            hits.append(m.group(0)[:80])
    return hits


async def _llm_score(text: str) -> tuple[float, str]:
    client = _get_client()
    try:
        resp = await client.chat.completions.create(
            model="gpt-4o-mini",
            temperature=0,
            response_format={"type": "json_object"},
            messages=[
                {"role": "system", "content": _INJECTION_SYSTEM},
                {"role": "user", "content": text[:800]},
            ],
        )
        data = json.loads(resp.choices[0].message.content)
        return float(data.get("injection_score", 0.0)), str(data.get("closest_pattern", ""))
    except Exception:
        return 0.0, ""


def _injection_score(regex_hits: list[str], llm_sim: float) -> float:
    regex_signal = min(1.0, len(regex_hits) * 0.4)
    return max(regex_signal, llm_sim)


# ── Main analysis function ─────────────────────────────────────────────────────

async def run_prompt_injection_analysis(
    rows: list[dict],
    text_columns: list[str],
    max_rows: int = 500,
    flag_threshold: float = FLAG_THRESHOLD,
    high_threshold: float = HIGH_THRESHOLD,
    concurrency: int = 20,
) -> dict:
    if not text_columns:
        return {"error": "No text columns for prompt injection detection"}

    try:
        _get_client()
    except RuntimeError as e:
        return {"error": str(e)}

    import numpy as np

    sample = rows[:max_rows]
    total_rows = len(sample)

    col_flagged: dict[str, int] = {c: 0 for c in text_columns}
    col_high:    dict[str, int] = {c: 0 for c in text_columns}
    col_scores:  dict[str, list[float]] = {c: [] for c in text_columns}
    flagged_rows: list[dict] = []

    sem = asyncio.Semaphore(concurrency)

    async def _bounded_score(text: str) -> tuple[float, str]:
        async with sem:
            return await _llm_score(text)

    for col in text_columns:
        texts_with_idx: list[tuple[int, str]] = []
        for i, row in enumerate(sample):
            val = row.get(col)
            if not val or not isinstance(val, str) or len(val.strip()) < 10:
                continue
            texts_with_idx.append((i, val.strip()))

        if not texts_with_idx:
            continue

        indices, texts = zip(*texts_with_idx)
        llm_results = await asyncio.gather(*[_bounded_score(t) for t in texts])

        for (row_idx, text), (sim, closest_pattern) in zip(texts_with_idx, llm_results):
            regex_hits = _regex_scan(text)
            score = _injection_score(regex_hits, sim)
            col_scores[col].append(score)

            if score >= flag_threshold:
                col_flagged[col] += 1
                confidence: str = "high" if score >= high_threshold else "medium"
                if score >= high_threshold:
                    col_high[col] += 1
                flagged_rows.append({
                    "row_idx":          row_idx,
                    "column":           col,
                    "score":            round(score, 4),
                    "confidence":       confidence,
                    "regex_hits":       regex_hits[:3],
                    "closest_template": closest_pattern,
                    "similarity":       round(sim, 4),
                    "text_snippet":     text[:120],
                })

    flagged_rows.sort(key=lambda r: r["score"], reverse=True)

    total_flagged = sum(col_flagged.values())
    total_high    = sum(col_high.values())
    flagged_pct   = round(total_flagged / max(total_rows, 1) * 100, 2)

    all_scores = [s for scores in col_scores.values() for s in scores]
    avg_score  = round(float(np.mean(all_scores)), 4) if all_scores else 0.0

    if total_flagged == 0:
        verdict = "clean"
    elif flagged_pct < 2:
        verdict = "low"
    elif flagged_pct < 10:
        verdict = "moderate"
    else:
        verdict = "high"

    col_summaries = []
    for col in text_columns:
        scores = col_scores[col]
        if not scores:
            continue
        col_summaries.append({
            "column":        col,
            "rows_analyzed": len(scores),
            "flagged":       col_flagged[col],
            "flagged_pct":   round(col_flagged[col] / max(len(scores), 1) * 100, 2),
            "high_conf":     col_high[col],
            "avg_score":     round(float(np.mean(scores)), 4),
        })

    return {
        "rows_analyzed":  total_rows,
        "total_flagged":  total_flagged,
        "flagged_pct":    flagged_pct,
        "total_high_conf": total_high,
        "avg_score":      avg_score,
        "verdict":        verdict,
        "embedding_mode": True,
        "col_summaries":  col_summaries,
        "flagged_rows":   flagged_rows[:200],
        "flag_threshold": flag_threshold,
        "high_threshold": high_threshold,
    }


# ── Request / Response models ──────────────────────────────────────────────────

class PromptInjectionRequest(BaseModel):
    rows:           list[dict]
    text_columns:   list[str]
    max_rows:       Optional[int]   = 500
    flag_threshold: Optional[float] = FLAG_THRESHOLD
    high_threshold: Optional[float] = HIGH_THRESHOLD
    concurrency:    Optional[int]   = 20


@router.post("/prompt-injection")
async def prompt_injection(req: PromptInjectionRequest):
    try:
        result = await run_prompt_injection_analysis(
            rows=req.rows,
            text_columns=req.text_columns,
            max_rows=req.max_rows or 500,
            flag_threshold=req.flag_threshold or FLAG_THRESHOLD,
            high_threshold=req.high_threshold or HIGH_THRESHOLD,
            concurrency=req.concurrency or 20,
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))