from __future__ import annotations

import asyncio
import gc
import io
import json
import os
import re
import threading
import time
import unicodedata
from collections import defaultdict, deque
from pathlib import Path
from typing import Any, Callable

import pytesseract
import torch
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image, ImageEnhance, ImageFile, ImageOps, UnidentifiedImageError
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    CLIPModel,
    CLIPProcessor,
    ViTForImageClassification,
    ViTImageProcessor,
    pipeline,
)

APP_DIR = Path(__file__).resolve().parent
STATIC_DIR = APP_DIR / "static"

NSFW_MODEL = "Falconsai/nsfw_image_detection"
VIOLENCE_MODEL = "jaranohaal/vit-base-violence-detection"
HATE_MODEL = "openai/clip-vit-base-patch32"
TOXIC_MODEL = "unitary/toxic-bert"

NSFW_THRESHOLD = float(os.getenv("NSFW_THRESHOLD", "0.82"))
VIOLENCE_THRESHOLD = float(os.getenv("VIOLENCE_THRESHOLD", "0.80"))
HATE_THRESHOLD = float(os.getenv("HATE_SYMBOL_THRESHOLD", "0.90"))
HATE_STRONG_THRESHOLD = float(os.getenv("HATE_SYMBOL_STRONG_THRESHOLD", "0.97"))
HATE_MIN_MATCHED_VARIANTS = int(os.getenv("HATE_SYMBOL_MIN_MATCHED_VARIANTS", "2"))
TOXIC_THRESHOLD = float(os.getenv("TOXIC_TEXT_THRESHOLD", "0.78"))
TEXT_BLOCKLIST_THRESHOLD = float(os.getenv("TEXT_BLOCKLIST_THRESHOLD", "0.96"))
MAX_UPLOAD_BYTES = int(float(os.getenv("MAX_UPLOAD_MB", "8")) * 1024 * 1024)
MAX_IMAGE_PIXELS = 24_000_000
MAX_ANALYSIS_EDGE = int(os.getenv("MAX_ANALYSIS_EDGE", "640"))
ENABLE_SLOW_OCR_FALLBACK = os.getenv("ENABLE_SLOW_OCR_FALLBACK", "0") == "1"

TEXT_BLOCKLIST = {
    # French insults / sexual slurs
    "pute",
    "putain",
    "salope",
    "connard",
    "connasse",
    "encule",
    "enculer",
    "nique",
    "fdp",
    "ntm",
    "ta mere",
    "ta mere la pute",
    # English profanity / sexual slurs
    "fuck",
    "fucker",
    "fucking",
    "shit",
    "bitch",
    "whore",
    "slut",
    "cunt",
    "dick",
    "pussy",
    "nigger",
    "nigga",
    "faggot",
    "retard",
    # Extremist / hate text
    "nazi",
    "hitler",
    "heil hitler",
    "sieg heil",
    "swastika",
    "white power",
    "whitepower",
    "1488",
    "88",
}

TEXT_COMPACT_BLOCKLIST = {
    "".join(ch for ch in term if ch.isalnum())
    for term in TEXT_BLOCKLIST
    if len("".join(ch for ch in term if ch.isalnum())) >= 3
}

HATE_SAFE_PROMPTS = [
    "a harmless ordinary image without hate symbols",
    "a normal avatar item or game asset with no extremist content",
    "a safe logo or clothing texture",
    "a safe video game clothing template",
    "a colorful shirt UV template with labeled body parts",
    "a blank avatar shirt template with colored rectangles and letters",
    "a red and white Axium brand logo",
    "a safe geometric brand logo",
]

HATE_RISK_PROMPTS = [
    "a nazi swastika symbol",
    "a nazi flag",
    "an image displaying a swastika",
    "an image displaying hate symbols",
    "an extremist supremacist logo",
    "a white supremacist symbol",
    "a hateful propaganda symbol",
]

Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
ImageFile.LOAD_TRUNCATED_IMAGES = False

app = FastAPI(title="Axium Image Moderation Lab", docs_url=None, redoc_url=None)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")

model_lock = threading.Lock()
model_cache_lock = threading.Lock()
rate_lock = threading.Lock()
recent_requests: dict[str, deque[float]] = defaultdict(deque)
model_cache: dict[str, Any] = {}


def env_thresholds_are_valid() -> None:
    for name, value in {
        "NSFW_THRESHOLD": NSFW_THRESHOLD,
        "VIOLENCE_THRESHOLD": VIOLENCE_THRESHOLD,
        "HATE_SYMBOL_THRESHOLD": HATE_THRESHOLD,
        "HATE_SYMBOL_STRONG_THRESHOLD": HATE_STRONG_THRESHOLD,
        "TOXIC_TEXT_THRESHOLD": TOXIC_THRESHOLD,
    }.items():
        if not 0.0 < value < 1.0:
            raise RuntimeError(f"{name} must be between 0 and 1")
    if HATE_MIN_MATCHED_VARIANTS < 1:
        raise RuntimeError("HATE_SYMBOL_MIN_MATCHED_VARIANTS must be 1 or higher")
    if MAX_UPLOAD_BYTES < 256 * 1024 or MAX_UPLOAD_BYTES > 20 * 1024 * 1024:
        raise RuntimeError("MAX_UPLOAD_MB must be between 0.25 and 20")


env_thresholds_are_valid()


@app.middleware("http")
async def security_headers(request: Request, call_next: Callable):
    response = await call_next(request)
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Referrer-Policy"] = "no-referrer"
    response.headers["Permissions-Policy"] = "camera=(), microphone=(), geolocation=()"
    response.headers["Content-Security-Policy"] = (
        "default-src 'self'; img-src 'self' blob: data:; "
        "style-src 'self'; script-src 'self'; connect-src 'self'; "
        "base-uri 'none'; form-action 'self'; "
        "frame-ancestors https://huggingface.co https://*.huggingface.co"
    )
    if request.url.path.startswith("/api/"):
        response.headers["Cache-Control"] = "no-store"
    return response


def enforce_rate_limit(request: Request) -> None:
    client = request.client.host if request.client else "unknown"
    now = time.monotonic()
    with rate_lock:
        bucket = recent_requests[client]
        while bucket and bucket[0] < now - 60:
            bucket.popleft()
        if len(bucket) >= 8:
            raise HTTPException(status_code=429, detail="Too many analyses. Try again in a minute.")
        bucket.append(now)


async def read_image(file: UploadFile) -> Image.Image:
    if file.content_type not in {"image/jpeg", "image/png", "image/webp"}:
        raise HTTPException(status_code=415, detail="Only JPEG, PNG and WebP images are accepted.")

    data = await file.read(MAX_UPLOAD_BYTES + 1)
    await file.close()
    if not data:
        raise HTTPException(status_code=400, detail="The selected file is empty.")
    if len(data) > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail="The image is larger than the configured limit.")

    try:
        with Image.open(io.BytesIO(data)) as probe:
            probe.verify()
        with Image.open(io.BytesIO(data)) as decoded:
            if decoded.format not in {"JPEG", "PNG", "WEBP"}:
                raise HTTPException(status_code=415, detail="The decoded image format is not allowed.")
            if decoded.width < 32 or decoded.height < 32:
                raise HTTPException(status_code=400, detail="The image is too small to analyse.")
            if decoded.width * decoded.height > MAX_IMAGE_PIXELS:
                raise HTTPException(status_code=400, detail="The image dimensions are too large.")
            image = decoded.convert("RGB")
            image.thumbnail((MAX_ANALYSIS_EDGE, MAX_ANALYSIS_EDGE), Image.Resampling.LANCZOS)
            return image.copy()
    except (UnidentifiedImageError, OSError, Image.DecompressionBombError):
        raise HTTPException(status_code=400, detail="The file is not a valid, safe image.")


def release_model(*objects: Any) -> None:
    for obj in objects:
        del obj
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def cached_resource(key: str, factory: Callable[[], Any]) -> Any:
    with model_cache_lock:
        if key not in model_cache:
            model_cache[key] = factory()
        return model_cache[key]


def normalized_label(label: str) -> str:
    return "".join(ch for ch in label.lower() if ch.isalnum() or ch == "_")


def normalized_text(text: str) -> str:
    folded = unicodedata.normalize("NFKD", text)
    folded = "".join(ch for ch in folded if not unicodedata.combining(ch))
    folded = folded.lower()
    folded = folded.translate(
        str.maketrans(
            {
                "@": "a",
                "4": "a",
                "0": "o",
                "1": "i",
                "!": "i",
                "|": "i",
                "3": "e",
                "5": "s",
                "$": "s",
                "7": "t",
                "+": "t",
            }
        )
    )
    folded = re.sub(r"(.)\1{2,}", r"\1\1", folded)
    return re.sub(r"\s+", " ", folded).strip()


def compact_text(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", "", normalized_text(text))


def find_blocked_text(text: str) -> str | None:
    spaced = f" {normalized_text(text)} "
    compact = compact_text(text)
    for term in sorted(TEXT_BLOCKLIST, key=len, reverse=True):
        norm = normalized_text(term)
        compact_term = re.sub(r"[^a-z0-9]+", "", norm)
        if " " in norm and f" {norm} " in spaced:
            return term
        if len(norm) <= 3:
            if re.search(rf"(?<![a-z0-9]){re.escape(norm)}(?![a-z0-9])", spaced):
                return term
            continue
        if re.search(rf"(?<![a-z0-9]){re.escape(norm)}(?![a-z0-9])", spaced):
            return term
        if compact_term in TEXT_COMPACT_BLOCKLIST and compact_term in compact:
            return term
    return None


def image_variants(
    image: Image.Image,
    *,
    include_crops: bool = False,
    include_right_angle_rotations: bool = False,
    include_half_turn: bool = True,
) -> list[Image.Image]:
    base = image.convert("RGB")
    variants = [base]
    if include_half_turn:
        variants.append(base.rotate(180, expand=True))
    if include_right_angle_rotations:
        variants.extend([base.rotate(90, expand=True), base.rotate(270, expand=True)])
    if include_crops and min(base.size) >= 96:
        width, height = base.size
        crop_boxes = [
            (0, 0, width // 2, height // 2),
            (width // 2, 0, width, height // 2),
            (0, height // 2, width // 2, height),
            (width // 2, height // 2, width, height),
            (width // 5, height // 5, width * 4 // 5, height * 4 // 5),
        ]
        variants.extend(base.crop(box) for box in crop_boxes)
    return variants


def run_nsfw(image: Image.Image) -> dict[str, Any]:
    classifier = cached_resource("nsfw_classifier", lambda: pipeline("image-classification", model=NSFW_MODEL, device=-1))
    scores = classifier(image, top_k=None)
    risky = max(
        (float(row["score"]) for row in scores if "nsfw" in normalized_label(row["label"])),
        default=0.0,
    )
    return test_result("Adult or NSFW content", risky, NSFW_THRESHOLD, NSFW_MODEL)


def run_violence(image: Image.Image) -> dict[str, Any]:
    processor, model = cached_resource(
        "violence_model",
        lambda: (ViTImageProcessor.from_pretrained(VIOLENCE_MODEL), ViTForImageClassification.from_pretrained(VIOLENCE_MODEL)),
    )
    inputs = processor(images=image, return_tensors="pt")
    with torch.inference_mode():
        probabilities = model(**inputs).logits.softmax(dim=-1)[0]
    risky = 0.0
    found_semantic_label = False
    for index, probability in enumerate(probabilities):
        label = normalized_label(str(model.config.id2label.get(index, index)))
        is_safe = label.startswith("non") or "nonviolence" in label or "safe" in label
        if "violence" in label and not is_safe:
            found_semantic_label = True
            risky = max(risky, float(probability))
    # The upstream Apache-2.0 checkpoint omits semantic id2label values.
    # Its binary training order is class 0 non-violent, class 1 violent.
    if not found_semantic_label and len(probabilities) == 2:
        risky = float(probabilities[1])
    return test_result("Violence", risky, VIOLENCE_THRESHOLD, VIOLENCE_MODEL)


def run_hateful_symbols(image: Image.Image) -> dict[str, Any]:
    processor, model = cached_resource(
        "hate_clip_model",
        lambda: (CLIPProcessor.from_pretrained(HATE_MODEL), CLIPModel.from_pretrained(HATE_MODEL)),
    )
    prompts = HATE_SAFE_PROMPTS + HATE_RISK_PROMPTS
    risky = 0.0
    best_prompt = ""
    matched_variants = 0
    variants = image_variants(image, include_half_turn=True, include_right_angle_rotations=False)
    inputs = processor(text=prompts, images=variants, return_tensors="pt", padding=True)
    with torch.inference_mode():
        logits = model(**inputs).logits_per_image
    safe_logits = logits[:, : len(HATE_SAFE_PROMPTS)].max(dim=1).values
    for row_index in range(logits.shape[0]):
        variant_risky = 0.0
        for index, prompt in enumerate(prompts[len(HATE_SAFE_PROMPTS) :], start=len(HATE_SAFE_PROMPTS)):
            binary = torch.stack((safe_logits[row_index], logits[row_index, index])).softmax(dim=0)
            score = float(binary[1])
            variant_risky = max(variant_risky, score)
            if score > risky:
                risky = score
                best_prompt = prompt
        if variant_risky >= HATE_THRESHOLD:
            matched_variants += 1
    verdict_score = risky
    if risky < HATE_STRONG_THRESHOLD and matched_variants < HATE_MIN_MATCHED_VARIANTS:
        verdict_score = min(risky, max(0.0, HATE_THRESHOLD - 0.01))
    result = test_result("Hateful-symbol heuristic", verdict_score, HATE_THRESHOLD, HATE_MODEL)
    result["raw_score"] = round(risky, 4)
    result["matched_variants"] = matched_variants
    result["detail"] = f"Closest risky label: {best_prompt or 'none'}."
    result["warning"] = "Experimental CLIP heuristic; uncertain cases still require human review."
    return result


def run_offensive_text(image: Image.Image) -> dict[str, Any]:
    text = extract_ocr_text(image)
    blocked = find_blocked_text(text)
    if blocked:
        result = test_result("Offensive text", TEXT_BLOCKLIST_THRESHOLD, TOXIC_THRESHOLD, "Tesseract OCR + Axium blocklist")
        result["detail"] = f'OCR detected blocked text "{blocked}" in: "{text[:180]}{"..." if len(text) > 180 else ""}"'
        return result

    if not text:
        result = test_result("Offensive text", 0.0, TOXIC_THRESHOLD, TOXIC_MODEL)
        result["detail"] = "No readable English or French text was found."
        return result

    classifier = cached_resource(
        "toxic_text_classifier",
        lambda: pipeline(
            "text-classification",
            model=AutoModelForSequenceClassification.from_pretrained(TOXIC_MODEL),
            tokenizer=AutoTokenizer.from_pretrained(TOXIC_MODEL),
            device=-1,
            top_k=None,
        ),
    )
    rows = classifier(text, truncation=True, max_length=512)
    if rows and isinstance(rows[0], list):
        rows = rows[0]
    risky_labels = {"toxic", "severetoxic", "obscene", "threat", "insult", "identityhate"}
    matched_scores = [
        float(row["score"])
        for row in rows
        if normalized_label(str(row["label"])).replace("_", "") in risky_labels
    ]
    # Some hosted checkpoints expose generic LABEL_0/LABEL_1 names instead of
    # semantic labels. In that case, use the highest model score instead of
    # silently returning 0% for every text image.
    risky = max(matched_scores, default=max((float(row["score"]) for row in rows), default=0.0))
    result = test_result("Offensive text", risky, TOXIC_THRESHOLD, TOXIC_MODEL)
    result["detail"] = f'OCR detected: "{text[:180]}{"..." if len(text) > 180 else ""}"'
    return result


def extract_ocr_text(image: Image.Image) -> str:
    candidates: list[Image.Image] = []
    variants = image_variants(image, include_crops=False, include_half_turn=True, include_right_angle_rotations=False)
    if ENABLE_SLOW_OCR_FALLBACK:
        variants = image_variants(image, include_crops=True, include_half_turn=True, include_right_angle_rotations=True)
    for variant in variants:
        base = variant.convert("RGB")
        gray = ImageOps.grayscale(base)
        wide = gray.resize((gray.width * 2, gray.height * 2), Image.Resampling.LANCZOS)
        contrast = ImageEnhance.Contrast(wide).enhance(2.4)
        thresholded = contrast.point(lambda px: 255 if px > 165 else 0)
        candidates.extend([gray, contrast, thresholded])
        if ENABLE_SLOW_OCR_FALLBACK:
            sharp = ImageEnhance.Sharpness(contrast).enhance(2.0)
            inverted = ImageOps.invert(contrast)
            candidates.extend([base, sharp, inverted])

    seen: set[str] = set()
    collected: list[str] = []
    configs = ("--oem 3 --psm 6",)
    if ENABLE_SLOW_OCR_FALLBACK:
        configs = ("--oem 3 --psm 6", "--oem 3 --psm 11")
    for candidate in candidates:
        for config in configs:
            try:
                text = " ".join(pytesseract.image_to_string(candidate, lang="eng+fra", config=config).split())
            except Exception:
                text = ""
            if not text:
                continue
            key = normalized_text(text)
            if key in seen:
                continue
            seen.add(key)
            collected.append(text)
            if find_blocked_text(text):
                return text[:2000]

    if not collected:
        return ""
    collected.sort(key=len, reverse=True)
    return " | ".join(collected[:6])[:2000]


def test_result(name: str, score: float, threshold: float, model: str) -> dict[str, Any]:
    passed = score < threshold
    return {
        "name": name,
        "status": "passed" if passed else "rejected",
        "passed": passed,
        "score": round(score, 4),
        "threshold": threshold,
        "model": model,
    }


TESTS: tuple[Callable[[Image.Image], dict[str, Any]], ...] = (
    run_nsfw,
    run_violence,
    run_hateful_symbols,
    run_offensive_text,
)

TEST_NAMES = {
    run_nsfw: "Adult or NSFW content",
    run_violence: "Violence",
    run_hateful_symbols: "Hateful-symbol heuristic",
    run_offensive_text: "Offensive text",
}


def test_name(test: Callable[[Image.Image], dict[str, Any]]) -> str:
    return TEST_NAMES.get(test, test.__name__.removeprefix("run_").replace("_", " ").title())


def analyse_sequentially(image: Image.Image) -> dict[str, Any]:
    results: list[dict[str, Any]] = []
    with model_lock:
        for test in TESTS:
            try:
                result = test(image)
            except Exception as exc:
                message = " ".join(str(exc).split())[:300]
                result = {
                    "name": test_name(test),
                    "status": "error",
                    "passed": False,
                    "detail": (
                        f"The model could not complete this test: {type(exc).__name__}"
                        f"{': ' + message if message else ''}"
                    ),
                }
                results.append(result)
                return {"verdict": "error", "passed": False, "results": results}
            gc.collect()
            results.append(result)
            if not result["passed"]:
                return {"verdict": "rejected", "passed": False, "results": results}
    return {"verdict": "approved", "passed": True, "results": results}


def ndjson_event(payload: dict[str, Any]) -> str:
    return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n"


def analyse_stream_events(image: Image.Image):
    results: list[dict[str, Any]] = []
    yield ndjson_event({"type": "progress", "active": None, "results": results})
    with model_lock:
        for index, test in enumerate(TESTS):
            yield ndjson_event({
                "type": "progress",
                "active": index,
                "activeName": test_name(test),
                "results": results,
            })
            try:
                result = test(image)
            except Exception as exc:
                message = " ".join(str(exc).split())[:300]
                result = {
                    "name": test_name(test),
                    "status": "error",
                    "passed": False,
                    "detail": (
                        f"The model could not complete this test: {type(exc).__name__}"
                        f"{': ' + message if message else ''}"
                    ),
                }
                results.append(result)
                yield ndjson_event({"type": "complete", "report": {"verdict": "error", "passed": False, "results": results}})
                return
            gc.collect()
            results.append(result)
            yield ndjson_event({"type": "progress", "active": None, "results": results})
            if not result["passed"]:
                yield ndjson_event({"type": "complete", "report": {"verdict": "rejected", "passed": False, "results": results}})
                return
    yield ndjson_event({"type": "complete", "report": {"verdict": "approved", "passed": True, "results": results}})


@app.get("/", include_in_schema=False)
async def home() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/api/analyse")
async def analyse(request: Request, image: UploadFile = File(...)) -> JSONResponse:
    enforce_rate_limit(request)
    decoded = await read_image(image)
    report = await asyncio.to_thread(analyse_sequentially, decoded)
    return JSONResponse(report)


@app.post("/api/analyse-stream")
async def analyse_stream(request: Request, image: UploadFile = File(...)) -> StreamingResponse:
    enforce_rate_limit(request)
    decoded = await read_image(image)
    return StreamingResponse(
        analyse_stream_events(decoded),
        media_type="application/x-ndjson",
        headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"},
    )