from __future__ import annotations

"""
Backdoor trigger phrase detection.

Strategy: unsupervised anomaly over (ngram, output-deviation) pairs.

For each text column we:
  1. Tokenize every row into unigrams + bigrams.
  2. Build a per-ngram output-deviation signal: compare each row's semantic
     embedding to the column-mean embedding. Rows with high cosine distance
     from the mean are "output deviants."
  3. Score each ngram by:
       trigger_score = frequency_skew × output_deviation_correlation
     where frequency_skew = (rate_in_deviants / rate_in_normal) — ngrams that
     appear disproportionately in deviant rows are candidates.
  4. Assign per-row risk as the max trigger_score of any flagged ngram present
     in that row.
  5. Return the top candidate trigger phrases + per-row risk scores.
"""

import math
import re
import string
from collections import Counter, defaultdict
from typing import Optional

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

router = APIRouter(prefix="/dataset")

# ── Thresholds ──────────────────────────────────────────────────────────────────

DEVIATION_THRESHOLD = 0.25   # cosine distance from column mean → "deviant"
TRIGGER_SCORE_MIN   = 0.40   # minimum trigger score to report an ngram
FLAG_THRESHOLD      = 0.45   # per-row risk to flag
HIGH_THRESHOLD      = 0.72
MIN_NGRAM_FREQ      = 2      # ignore ngrams appearing only once

# ── Text normalisation ──────────────────────────────────────────────────────────

_PUNCT = str.maketrans("", "", string.punctuation)


def _tokenize(text: str) -> list[str]:
    """Lowercase, strip punctuation, split on whitespace. Return unigrams + bigrams."""
    tokens = text.lower().translate(_PUNCT).split()
    unigrams = tokens
    bigrams  = [f"{a} {b}" for a, b in zip(tokens, tokens[1:])]
    return unigrams + bigrams


# ── Output-deviation scoring ────────────────────────────────────────────────────

def _deviation_scores(texts: list[str]) -> list[float]:
    """
    Length + vocabulary-richness deviation from column mean.
    Fast, no model required — sufficient for trigger-phrase frequency skew.
    """
    import numpy as np

    def _features(t: str) -> tuple[float, float]:
        tokens = t.lower().split()
        length = len(tokens)
        ttr = len(set(tokens)) / max(length, 1)  # type-token ratio
        return length, ttr

    feats = [_features(t) for t in texts]
    arr = np.array(feats, dtype=float)
    if len(arr) < 2:
        return [0.0] * len(texts)
    mean = arr.mean(axis=0)
    std  = arr.std(axis=0)
    std[std < 1e-9] = 1.0
    z = np.abs((arr - mean) / std).mean(axis=1)
    z_min, z_max = z.min(), z.max()
    return ((z - z_min) / max(z_max - z_min, 1e-9)).tolist()


# ── Trigger-phrase scoring ──────────────────────────────────────────────────────

def _score_triggers(
    texts: list[str],
    dev_scores: list[float],
    deviation_threshold: float,
    min_freq: int,
) -> tuple[dict[str, float], list[bool]]:
    """
    Returns:
      ngram_scores  — {ngram: trigger_score}
      is_deviant    — bool per row
    """
    import numpy as np

    # Pre-tokenize once — avoids O(rows × ngrams) re-tokenization
    tokenized = [set(_tokenize(t)) for t in texts]

    is_deviant = [d >= deviation_threshold for d in dev_scores]

    n_dev = max(sum(is_deviant), 1)
    n_nrm = max(sum(not d for d in is_deviant), 1)

    dev_counts: Counter[str] = Counter()
    nrm_counts: Counter[str] = Counter()
    for toks, dev in zip(tokenized, is_deviant):
        if dev:
            dev_counts.update(toks)
        else:
            nrm_counts.update(toks)

    # Build ngram → list of row deviation scores for fast avg_dev lookup
    ngram_dev_scores: dict[str, list[float]] = defaultdict(list)
    for toks, score in zip(tokenized, dev_scores):
        for ng in toks:
            ngram_dev_scores[ng].append(score)

    all_ngrams = set(dev_counts.keys()) | set(nrm_counts.keys())
    ngram_scores: dict[str, float] = {}

    for ng in all_ngrams:
        total = dev_counts[ng] + nrm_counts[ng]
        if total < min_freq:
            continue

        rate_dev = dev_counts[ng] / n_dev
        rate_nrm = nrm_counts[ng] / n_nrm

        skew = rate_dev / max(rate_nrm, 1e-4)
        if rate_nrm > 0.60:
            continue

        avg_dev = float(np.mean(ngram_dev_scores[ng])) if ngram_dev_scores[ng] else 0.0
        trigger_score = math.sqrt(min(skew / 10.0, 1.0) * avg_dev)
        if trigger_score >= TRIGGER_SCORE_MIN:
            ngram_scores[ng] = round(trigger_score, 4)

    return ngram_scores, is_deviant


# ── Per-row risk assignment ─────────────────────────────────────────────────────

def _row_risk(text: str, ngram_scores: dict[str, float]) -> float:
    if not ngram_scores:
        return 0.0
    tokens = set(_tokenize(text))
    hits = [ngram_scores[ng] for ng in tokens if ng in ngram_scores]
    return max(hits) if hits else 0.0


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

def run_backdoor_trigger_analysis(
    rows: list[dict],
    text_columns: list[str],
    max_rows: int = 500,
    flag_threshold: float = FLAG_THRESHOLD,
    high_threshold: float = HIGH_THRESHOLD,
    deviation_threshold: float = DEVIATION_THRESHOLD,
    min_ngram_freq: int = MIN_NGRAM_FREQ,
    top_triggers: int = 30,
) -> dict:
    import numpy as np

    if not text_columns:
        return {"error": "No text columns specified"}

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

    col_summaries: list[dict]  = []
    flagged_rows:  list[dict]  = []
    all_triggers:  dict[str, float] = {}

    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_agg: dict[str, list[float]] = {c: [] for c in text_columns}

    for col in text_columns:
        texts: list[str] = []
        valid_idx: list[int] = []
        for i, row in enumerate(sample):
            val = row.get(col)
            if val and isinstance(val, str) and len(val.strip()) >= 10:
                texts.append(val.strip())
                valid_idx.append(i)

        if len(texts) < 4:
            continue

        dev_scores = _deviation_scores(texts)
        ngram_scores, is_deviant = _score_triggers(
            texts, dev_scores, deviation_threshold, min_ngram_freq
        )

        # Merge column-level triggers into global set (take max score)
        for ng, sc in ngram_scores.items():
            all_triggers[ng] = max(all_triggers.get(ng, 0.0), sc)

        for local_i, (vi, text) in enumerate(zip(valid_idx, texts)):
            risk = _row_risk(text, ngram_scores)
            dev  = dev_scores[local_i]
            col_scores_agg[col].append(risk)

            if risk >= flag_threshold:
                col_flagged[col] += 1
                confidence = "high" if risk >= high_threshold else "medium"
                if risk >= high_threshold:
                    col_high[col] += 1

                # Which specific triggers fired
                hit_ngrams = sorted(
                    [ng for ng in set(_tokenize(text)) if ng in ngram_scores],
                    key=lambda ng: ngram_scores[ng],
                    reverse=True,
                )[:5]

                flagged_rows.append({
                    "row_idx": vi,
                    "column": col,
                    "score": round(risk, 4),
                    "confidence": confidence,
                    "deviation_score": round(dev, 4),
                    "trigger_phrases": hit_ngrams,
                    "text_snippet": text[:120],
                })

        scores = col_scores_agg[col]
        if scores:
            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),
                "n_trigger_phrases": len(ngram_scores),
            })

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

    total_flagged = sum(col_flagged.values())
    total_high    = sum(col_high.values())
    all_scores    = [s for ss in col_scores_agg.values() for s in ss]
    avg_score     = round(float(np.mean(all_scores)), 4) if all_scores else 0.0
    flagged_pct   = round(total_flagged / max(total_rows, 1) * 100, 2)

    if total_flagged == 0:
        verdict = "clean"
    elif flagged_pct < 1:
        verdict = "low"
    elif flagged_pct < 8:
        verdict = "moderate"
    else:
        verdict = "high"

    top_trigger_list = sorted(
        [{"phrase": ng, "score": sc} for ng, sc in all_triggers.items()],
        key=lambda x: x["score"],
        reverse=True,
    )[:top_triggers]

    return {
        "rows_analyzed": total_rows,
        "total_flagged": total_flagged,
        "flagged_pct": flagged_pct,
        "total_high_conf": total_high,
        "avg_score": avg_score,
        "verdict": verdict,
        "top_trigger_phrases": top_trigger_list,
        "col_summaries": col_summaries,
        "flagged_rows": flagged_rows[:200],
        "flag_threshold": flag_threshold,
        "high_threshold": high_threshold,
        "deviation_threshold": deviation_threshold,
        "signals": ["frequency_skew", "output_deviation_correlation"],
    }


# ── Request model + route ───────────────────────────────────────────────────────

class BackdoorTriggerRequest(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
    deviation_threshold: Optional[float] = DEVIATION_THRESHOLD
    min_ngram_freq: Optional[int] = MIN_NGRAM_FREQ
    top_triggers: Optional[int] = 30


@router.post("/backdoor-triggers")
def backdoor_triggers(req: BackdoorTriggerRequest):
    try:
        return run_backdoor_trigger_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,
            deviation_threshold=req.deviation_threshold or DEVIATION_THRESHOLD,
            min_ngram_freq=req.min_ngram_freq or MIN_NGRAM_FREQ,
            top_triggers=req.top_triggers or 30,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))