from __future__ import annotations

import asyncio
import json
import os
from collections import Counter
from typing import Optional

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

router = APIRouter(prefix="/dataset")


# ── Config ─────────────────────────────────────────────────────────────────────

LABELS = [
    "toxicity",
    "severe_toxicity",
    "obscene",
    "threat",
    "insult",
    "identity_attack",
]

LABEL_DISPLAY = {
    "toxicity":        "Toxicity",
    "severe_toxicity": "Severe Toxicity",
    "obscene":         "Obscene Language",
    "threat":          "Threat / Violence",
    "insult":          "Insult",
    "identity_attack": "Identity Attack",
}

LABEL_COLOR = {
    "toxicity":        "risk",
    "severe_toxicity": "critical",
    "obscene":         "medium",
    "threat":          "critical",
    "insult":          "medium",
    "identity_attack": "high",
}

LABEL_CATEGORY: dict[str, str] = {
    "toxicity":          "general",
    "severe_toxicity":   "general",
    "obscene":           "profanity",
    "threat":            "violence",
    "insult":            "harassment",
    "identity_attack":   "hate_speech",
}

EXTENDED_LABELS: list[str] = [
    "toxicity",
    "severe_toxicity",
    "identity_attack",
    "insult",
    "profanity",
    "threat",
    "attack_on_commenter",
    "attack_on_author",
    "inflammatory",
    "spam",
    "obscene",
    "sexually_explicit",
    "flirtation",
    "incoherent",
    "unsubstantial",
    "hate_speech",
    "self_harm",
    "radicalization",
    "misinformation",
    "harassment",
]

EXTENDED_LABEL_DISPLAY: dict[str, str] = {
    "toxicity":           "Toxicity",
    "severe_toxicity":    "Severe Toxicity",
    "identity_attack":    "Identity Attack / Hate",
    "insult":             "Insult",
    "profanity":          "Profanity",
    "threat":             "Threat / Violence",
    "attack_on_commenter":"Attack on Commenter",
    "attack_on_author":   "Attack on Author",
    "inflammatory":       "Inflammatory",
    "spam":               "Spam",
    "obscene":            "Obscene Language",
    "sexually_explicit":  "Sexually Explicit",
    "flirtation":         "Flirtation",
    "incoherent":         "Incoherent",
    "unsubstantial":      "Low Effort / Unsubstantial",
    "hate_speech":        "Hate Speech",
    "self_harm":          "Self-Harm",
    "radicalization":     "Radicalization",
    "misinformation":     "Misinformation",
    "harassment":         "Harassment",
}

EXTENDED_LABEL_COLOR: dict[str, str] = {
    "toxicity":           "risk",
    "severe_toxicity":    "critical",
    "identity_attack":    "high",
    "insult":             "medium",
    "profanity":          "medium",
    "threat":             "critical",
    "attack_on_commenter":"medium",
    "attack_on_author":   "medium",
    "inflammatory":       "medium",
    "spam":               "low",
    "obscene":            "medium",
    "sexually_explicit":  "high",
    "flirtation":         "low",
    "incoherent":         "low",
    "unsubstantial":      "low",
    "hate_speech":        "critical",
    "self_harm":          "critical",
    "radicalization":     "critical",
    "misinformation":     "high",
    "harassment":         "high",
}

LABEL_COMPLIANCE_RISK: dict[str, str] = {
    "toxicity":           "medium",
    "severe_toxicity":    "high",
    "identity_attack":    "high",
    "insult":             "medium",
    "profanity":          "low",
    "threat":             "high",
    "hate_speech":        "high",
    "self_harm":          "high",
    "radicalization":     "high",
    "sexually_explicit":  "high",
    "misinformation":     "medium",
    "harassment":         "medium",
    "spam":               "low",
    "inflammatory":       "medium",
    "obscene":            "low",
    "attack_on_commenter":"medium",
    "attack_on_author":   "medium",
    "incoherent":         "low",
    "unsubstantial":      "low",
    "flirtation":         "low",
}

FLAG_THRESHOLD   = 0.5
SEVERE_THRESHOLD = 0.8

_SYSTEM_PROMPT = (
    "You are a toxicity classifier. Given a text, return a JSON object with scores "
    "between 0.0 and 1.0 for each of these labels: "
    "toxicity, severe_toxicity, obscene, threat, insult, identity_attack. "
    "1.0 means clearly present, 0.0 means absent. 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 — toxicity analysis unavailable.")
        _client = openai.AsyncOpenAI(api_key=api_key)
    return _client


async def _score_text(text: str) -> dict[str, float]:
    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": _SYSTEM_PROMPT},
                {"role": "user", "content": text[:1000]},
            ],
        )
        data = json.loads(resp.choices[0].message.content)
        return {l: float(data.get(l, 0.0)) for l in LABELS}
    except Exception:
        return {l: 0.0 for l in LABELS}


# ── Core ───────────────────────────────────────────────────────────────────────

async def run_toxicity_analysis(
    rows: list[dict],
    text_columns: list[str],
    max_rows: int = 300,
    concurrency: int = 20,
    flag_threshold: float = FLAG_THRESHOLD,
    severe_threshold: float = SEVERE_THRESHOLD,
) -> dict:
    if not text_columns:
        return {"error": "No text columns for toxicity analysis"}

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

    sample = rows[:max_rows]

    col_scores:  dict[str, dict[str, list[float]]] = {c: {l: [] for l in LABELS} for c in text_columns}
    col_flagged: dict[str, int] = {c: 0 for c in text_columns}
    col_severe:  dict[str, int] = {c: 0 for c in text_columns}
    global_scores: dict[str, list[float]] = {l: [] for l in LABELS}
    flagged_rows: list[dict] = []
    total_severe = 0

    sem = asyncio.Semaphore(concurrency)

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

    for col in text_columns:
        texts_with_idx: list[tuple[int, str]] = []
        for i, row in enumerate(sample):
            val = row.get(col)
            if val and isinstance(val, str) and val.strip():
                texts_with_idx.append((i, val.strip()[:512]))

        if not texts_with_idx:
            continue

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

        for row_idx, row_label_scores in zip(indices, results):
            for l in LABELS:
                col_scores[col][l].append(row_label_scores[l])
                global_scores[l].append(row_label_scores[l])

            max_score   = max(row_label_scores.values())
            worst_label = max(row_label_scores, key=row_label_scores.get)

            if max_score >= flag_threshold:
                col_flagged[col] += 1
                flagged_rows.append({
                    "row_idx":     row_idx,
                    "column":      col,
                    "scores":      {l: round(row_label_scores[l], 4) for l in LABELS},
                    "worst_label": worst_label,
                    "worst_score": round(max_score, 4),
                })
                if max_score >= severe_threshold:
                    col_severe[col] += 1
                    total_severe += 1

    total_rows    = len(sample)
    total_flagged = sum(col_flagged.values())

    def _col_summary(col: str) -> dict | None:
        scores = col_scores[col]
        n = len(scores[LABELS[0]]) if scores[LABELS[0]] else 0
        if n == 0:
            return None
        avg = {l: round(sum(scores[l]) / n, 4) for l in LABELS}
        peak_label = max(avg, key=avg.get)
        return {
            "column":        col,
            "rows_analyzed": n,
            "flagged":       col_flagged[col],
            "flagged_pct":   round(col_flagged[col] / total_rows * 100, 1),
            "severe":        col_severe[col],
            "avg_scores":    avg,
            "peak_label":    peak_label,
            "peak_score":    avg[peak_label],
        }

    column_summaries = [
        s for col in text_columns
        if (s := _col_summary(col)) is not None
    ]
    column_summaries.sort(key=lambda x: x["peak_score"], reverse=True)

    n_global = len(global_scores[LABELS[0]]) or 1
    global_avg = {
        l: round(sum(global_scores[l]) / n_global, 4)
        for l in LABELS
    }

    peak_global = max(global_avg.values()) if global_avg else 0.0
    if peak_global >= severe_threshold:
        overall_severity = "severe"
    elif peak_global >= flag_threshold:
        overall_severity = "flagged"
    elif peak_global >= 0.2:
        overall_severity = "moderate"
    else:
        overall_severity = "clean"

    flagged_rows.sort(key=lambda x: x["worst_score"], reverse=True)

    enriched_labels = [
        {
            "label":           l,
            "display":         LABEL_DISPLAY.get(l, l),
            "avg_score":       global_avg[l],
            "color":           LABEL_COLOR.get(l, "low"),
            "category":        LABEL_CATEGORY.get(l, "general"),
            "compliance_risk": LABEL_COMPLIANCE_RISK.get(l, "low"),
        }
        for l in LABELS
    ]
    enriched_labels.sort(key=lambda x: x["avg_score"], reverse=True)

    return {
        "rows_analyzed":    total_rows,
        "total_flagged":    total_flagged,
        "flagged_pct":      round(total_flagged / total_rows * 100, 1) if total_rows else 0,
        "total_severe":     total_severe,
        "overall_severity": overall_severity,
        "global_avg":       global_avg,
        "label_details":    enriched_labels,
        "column_summaries": column_summaries,
        "flagged_rows":     flagged_rows[:50],
        "flag_threshold":   flag_threshold,
        "severe_threshold": severe_threshold,
        "taxonomy": {
            "model_labels":    LABELS,
            "extended_labels": EXTENDED_LABELS,
            "label_display":   EXTENDED_LABEL_DISPLAY,
            "label_color":     EXTENDED_LABEL_COLOR,
            "compliance_risk": LABEL_COMPLIANCE_RISK,
        },
    }


# ── Request model ──────────────────────────────────────────────────────────────

class ToxicityRequest(BaseModel):
    rows:             list[dict]
    text_columns:     list[str]
    max_rows:         int   = 300
    concurrency:      int   = 20
    flag_threshold:   float = FLAG_THRESHOLD
    severe_threshold: float = SEVERE_THRESHOLD


# ── Route ──────────────────────────────────────────────────────────────────────

@router.post("/toxicity")
async def toxicity_detection(req: ToxicityRequest):
    if not req.rows:
        raise HTTPException(status_code=400, detail="No rows provided")
    return await run_toxicity_analysis(
        req.rows,
        req.text_columns,
        max_rows=req.max_rows,
        concurrency=req.concurrency,
        flag_threshold=req.flag_threshold,
        severe_threshold=req.severe_threshold,
    )
