from __future__ import annotations

import asyncio
import json
import math
import os
from typing import Optional

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

router = APIRouter(prefix="/dataset")

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

OUTLIER_THRESHOLD = 0.62
HIGH_THRESHOLD    = 0.80
NEAR_DUP_SIM      = 0.92

# ── 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 — poisoned sample analysis unavailable.")
        _client = openai.AsyncOpenAI(api_key=api_key)
    return _client


# ── Embedding via text-embedding-3-large ──────────────────────────────────────

async def _embed_texts(texts: list[str], concurrency: int = 20) -> Optional[list[list[float]]]:
    client = _get_client()
    sem = asyncio.Semaphore(concurrency)

    async def _one(text: str) -> list[float]:
        async with sem:
            resp = await client.embeddings.create(
                model="text-embedding-3-large",
                input=text[:512],
            )
            return resp.data[0].embedding

    try:
        return await asyncio.gather(*[_one(t) for t in texts])
    except Exception:
        return None


def _normalize(vecs: list[list[float]]):
    import numpy as np
    arr = np.array(vecs, dtype=float)
    norms = np.linalg.norm(arr, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1.0, norms)
    return arr / norms


# ── Signal 1: Embedding outlier within topic cluster ──────────────────────────

def _zero_shot_topic(embs, n_clusters: int = 8) -> list[int]:
    import numpy as np
    try:
        from sklearn.cluster import MiniBatchKMeans
        k = min(n_clusters, len(embs))
        km = MiniBatchKMeans(n_clusters=k, random_state=42, n_init=3)
        return km.fit_predict(embs).tolist()
    except ImportError:
        pass
    rng = np.random.default_rng(42)
    k = min(n_clusters, len(embs))
    centroids = embs[rng.choice(len(embs), k, replace=False)]
    labels = (embs @ centroids.T).argmax(axis=1).tolist()
    return labels


def _cluster_outlier_scores(embs, labels: list[int]) -> list[float]:
    import numpy as np
    labels_arr = np.array(labels)
    scores = np.zeros(len(embs))
    for cid in set(labels):
        mask = labels_arr == cid
        members = embs[mask]
        centroid = members.mean(axis=0)
        norm = np.linalg.norm(centroid)
        if norm > 0:
            centroid /= norm
        sims = members @ centroid
        dists = 1.0 - sims
        scores[mask] = dists
    return scores.tolist()


# ── Signal 2: Label inconsistency across near-duplicate inputs ─────────────────

def _label_inconsistency_scores(
    embs,
    labels: list[str],
    sim_threshold: float = NEAR_DUP_SIM,
) -> list[float]:
    import numpy as np
    n = len(embs)
    scores = [0.0] * n
    if n < 2:
        return scores
    sim_matrix = embs @ embs.T
    for i in range(n):
        neighbors = [
            j for j in range(n)
            if j != i and sim_matrix[i, j] >= sim_threshold
        ]
        if not neighbors:
            continue
        mismatches = sum(1 for j in neighbors if labels[j] != labels[i])
        scores[i] = mismatches / len(neighbors)
    return scores


# ── Signal 3: Anomaly score via gpt-4o-mini ───────────────────────────────────

_ANOMALY_SYSTEM = (
    "You are a training-data quality auditor. Given a text sample, return a JSON object "
    "with a single key 'anomaly_score' (float 0.0–1.0). "
    "1.0 means highly anomalous (garbled, adversarial, templated, or suspiciously clean/repetitive), "
    "0.0 means normal natural text. Return only valid JSON, no explanation."
)

async def _anomaly_scores(texts: list[str], concurrency: int = 20) -> list[float]:
    client = _get_client()
    sem = asyncio.Semaphore(concurrency)

    async def _one(text: str) -> float:
        async with sem:
            try:
                resp = await client.chat.completions.create(
                    model="gpt-4o-mini",
                    temperature=0,
                    response_format={"type": "json_object"},
                    messages=[
                        {"role": "system", "content": _ANOMALY_SYSTEM},
                        {"role": "user", "content": text[:512]},
                    ],
                )
                data = json.loads(resp.choices[0].message.content)
                return float(data.get("anomaly_score", 0.0))
            except Exception:
                return 0.0

    return list(await asyncio.gather(*[_one(t) for t in texts]))


# ── Score blending ─────────────────────────────────────────────────────────────

def _blend(cluster_score: float, label_score: float, loss_score: float) -> float:
    return max(
        0.5 * cluster_score + 0.3 * label_score + 0.2 * loss_score,
        label_score * 0.9,
    )


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

async def run_poisoned_sample_analysis(
    rows: list[dict],
    text_columns: list[str],
    label_column: Optional[str] = None,
    max_rows: int = 500,
    flag_threshold: float = OUTLIER_THRESHOLD,
    high_threshold: float = HIGH_THRESHOLD,
    near_dup_sim: float = NEAR_DUP_SIM,
    n_clusters: int = 8,
    concurrency: int = 20,
) -> dict:
    import numpy as np

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

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

    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] = []

    for col in text_columns:
        texts = []
        valid_indices = []
        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_indices.append(i)

        if not texts:
            continue

        # ── Embed all texts ───────────────────────────────────────────────────
        raw_embs = await _embed_texts(texts, concurrency=concurrency)
        embs = _normalize(raw_embs) if raw_embs is not None else None

        # ── Signal 1: cluster outlier ─────────────────────────────────────────
        if embs is not None:
            cluster_labels = _zero_shot_topic(embs, n_clusters=n_clusters)
            cluster_scores = _cluster_outlier_scores(embs, cluster_labels)
        else:
            cluster_labels = [0] * len(texts)
            cluster_scores = [0.0] * len(texts)

        # ── Signal 2: label inconsistency ─────────────────────────────────────
        if label_column and embs is not None:
            raw_labels = [str(sample[vi].get(label_column, "")) for vi in valid_indices]
            label_scores = _label_inconsistency_scores(embs, raw_labels, near_dup_sim)
        else:
            raw_labels = [""] * len(texts)
            label_scores = [0.0] * len(texts)

        # ── Signal 3: anomaly score ────────────────────────────────────────────
        loss_scores = await _anomaly_scores(texts, concurrency=concurrency)

        # ── Blend + collect ───────────────────────────────────────────────────
        for local_idx, (vi, text) in enumerate(zip(valid_indices, texts)):
            cs = cluster_scores[local_idx]
            ls = label_scores[local_idx]
            ps = loss_scores[local_idx]
            score = _blend(cs, ls, ps)
            col_scores[col].append(score)

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

                entry: dict = {
                    "row_idx":             vi,
                    "column":              col,
                    "score":               round(score, 4),
                    "confidence":          confidence,
                    "cluster_score":       round(cs, 4),
                    "label_inconsistency": round(ls, 4),
                    "loss_anomaly":        round(ps, 4),
                    "cluster_id":          cluster_labels[local_idx],
                    "text_snippet":        text[:120],
                }
                if label_column:
                    entry["label"] = raw_labels[local_idx]
                flagged_rows.append(entry)

    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 < 1:
        verdict = "low"
    elif flagged_pct < 8:
        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,
        "label_column":   label_column,
        "col_summaries":  col_summaries,
        "flagged_rows":   flagged_rows[:200],
        "flag_threshold": flag_threshold,
        "high_threshold": high_threshold,
        "signals":        ["cluster_outlier", "label_inconsistency", "loss_anomaly"],
    }


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

class PoisonedSampleRequest(BaseModel):
    rows:           list[dict]
    text_columns:   list[str]
    label_column:   Optional[str]   = None
    max_rows:       Optional[int]   = 500
    flag_threshold: Optional[float] = OUTLIER_THRESHOLD
    high_threshold: Optional[float] = HIGH_THRESHOLD
    near_dup_sim:   Optional[float] = NEAR_DUP_SIM
    n_clusters:     Optional[int]   = 8
    concurrency:    Optional[int]   = 20


@router.post("/poisoned-samples")
async def poisoned_samples(req: PoisonedSampleRequest):
    try:
        return await run_poisoned_sample_analysis(
            rows=req.rows,
            text_columns=req.text_columns,
            label_column=req.label_column,
            max_rows=req.max_rows or 500,
            flag_threshold=req.flag_threshold or OUTLIER_THRESHOLD,
            high_threshold=req.high_threshold or HIGH_THRESHOLD,
            near_dup_sim=req.near_dup_sim or NEAR_DUP_SIM,
            n_clusters=req.n_clusters or 8,
            concurrency=req.concurrency or 20,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))