app.py

import os
import base64
import hashlib
import threading
from io import BytesIO
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor

from PIL import Image
from pydantic import BaseModel
from fastapi.responses import HTMLResponse
import gradio as gr
from gradio import Server

from xai_engine.model import get_model
from xai_engine.methods import run_all_methods, run_single_method, run_per_head_attention
from forensics.runner import run_forensic_tool
from logger import log_run_dataset

# Initialize Gradio Server app
app = Server(title="DeepfakeDetection-XAI")
demo = app

# ──────────────────────────────────────────────────────────────────────
# Result Cache — SHA256-keyed LRU, server-side RAM
# ──────────────────────────────────────────────────────────────────────
_result_cache: OrderedDict[str, dict] = OrderedDict()
_cache_lock = threading.Lock()
_CACHE_MAX = 50  # ~120MB max at 6 XAI + 5 forensics per entry


def _img_hash(image: Image.Image) -> str:
    """SHA256 hash of image pixel content for cache keying."""
    buf = BytesIO()
    image.save(buf, format="PNG")
    return hashlib.sha256(buf.getvalue()).hexdigest()


def _cache_get(key: str) -> dict | None:
    with _cache_lock:
        val = _result_cache.get(key)
        if val is not None:
            _result_cache.move_to_end(key)  # LRU touch
        return val


def _cache_set(key: str, val: dict):
    with _cache_lock:
        _result_cache[key] = val
        if len(_result_cache) > _CACHE_MAX:
            _result_cache.popitem(last=False)  # Evict oldest


# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────
def b64_to_pil(b64_str: str) -> Image.Image:
    """Convert base64 image data URL to PIL Image."""
    if not b64_str or not isinstance(b64_str, str):
        raise ValueError("Invalid base64 image string provided.")
    if "," in b64_str:
        _, data = b64_str.split(",", 1)
    else:
        data = b64_str
    image_data = base64.b64decode(data)
    return Image.open(BytesIO(image_data)).convert("RGB")


def _run_all_forensics_defaults(image: Image.Image) -> dict:
    """Run all 5 classic forensics tools at default parameters. Returns dict keyed by tool id."""
    tools = ["ela", "gradient", "bitplane", "minmax", "wavelet"]
    results = {}
    for t in tools:
        try:
            results[t] = run_forensic_tool(image, t)
        except Exception as e:
            print(f"[forensics] Error running {t}: {e}")
    return results


# ──────────────────────────────────────────────────────────────────────
# Pydantic Request Schemas
# ──────────────────────────────────────────────────────────────────────
class PredictReq(BaseModel):
    image_b64: str


class AnalyzeReq(BaseModel):
    image_b64: str
    alpha: float = 0.6
    cmap: str = "jet"


class MethodDetailReq(BaseModel):
    image_b64: str
    method_name: str = "rollout"
    steps: int = 20
    patch_size: int = 32
    stride: int = 16
    target_layer_idx: int = 11
    alpha: float = 0.6
    cmap: str = "jet"


class HeadsReq(BaseModel):
    image_b64: str
    layer_idx: int = 11
    cmap: str = "jet"


class ForensicReq(BaseModel):
    image_b64: str
    tool_name: str = "ela"
    quality: int = 75
    scale: int = 50
    contrast: int = 20
    gradient_intensity: int = 90
    blue_mode: str = "Abs"
    channel: str = "Luminance"
    bit: int = 0
    minmax_radius: int = 2


class SubmitRunReq(BaseModel):
    image_b64: str
    image_hash: str | None = None
    ground_truth: str = "unknown"
    opt_in_image: bool = True


# ──────────────────────────────────────────────────────────────────────
# REST API Routes
# ──────────────────────────────────────────────────────────────────────
@app.post("/api/predict")
def predict_post(req: PredictReq) -> dict:
    """Fast prediction — no XAI, just verdict + confidence."""
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    model = get_model()
    pred_data = model.predict(image)
    img_hash = _img_hash(image)

    return {
        "status": "success",
        "prediction": pred_data,
        "image_hash": img_hash,
    }


@app.post("/api/analyze")
def analyze_post(req: AnalyzeReq) -> dict:
    """Full analysis: XAI + forensics. Results are cached."""
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    key = _img_hash(image)

    # Check cache
    cached = _cache_get(key)
    if cached:
        return {**cached, "cache_hit": True}

    model = get_model()
    pred_data = model.predict(image)

    # Run XAI methods and forensics in parallel
    with ThreadPoolExecutor(max_workers=2) as executor:
        xai_future = executor.submit(run_all_methods, model, image, alpha=req.alpha, cmap=req.cmap)
        forensics_future = executor.submit(_run_all_forensics_defaults, image)

        xai_results = xai_future.result()
        forensics_results = forensics_future.result()

    res_dicts = [
        {
            "name": r.name,
            "description": r.description,
            "overlay_b64": r.overlay_b64,
            "raw_heatmap_b64": r.raw_heatmap_b64,
            "metadata": r.metadata,
            "compute_time_ms": r.compute_time_ms,
        }
        for r in xai_results
    ]

    response = {
        "status": "success",
        "prediction": pred_data,
        "results": res_dicts,
        "forensics": forensics_results,
        "image_hash": key,
        "cache_hit": False,
    }

    _cache_set(key, response)
    return response


@app.post("/api/method_detail")
def method_detail_post(req: MethodDetailReq) -> dict:
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    model = get_model()
    res = run_single_method(
        model, image, req.method_name,
        alpha=req.alpha, cmap=req.cmap, steps=req.steps,
        patch_size=req.patch_size, stride=req.stride,
        target_layer_idx=req.target_layer_idx,
    )

    return {
        "status": "success",
        "result": {
            "name": res.name,
            "description": res.description,
            "overlay_b64": res.overlay_b64,
            "raw_heatmap_b64": res.raw_heatmap_b64,
            "metadata": res.metadata,
            "compute_time_ms": res.compute_time_ms,
        }
    }


@app.post("/api/heads")
def heads_post(req: HeadsReq) -> dict:
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    model = get_model()
    head_data = run_per_head_attention(model, image, layer_idx=req.layer_idx, cmap=req.cmap)
    head_data["status"] = "success"
    return head_data


@app.post("/api/forensic")
def forensic_post(req: ForensicReq) -> dict:
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    res = run_forensic_tool(
        image, req.tool_name,
        quality=req.quality, scale=req.scale, contrast=req.contrast,
        gradient_intensity=req.gradient_intensity, blue_mode=req.blue_mode,
        channel=req.channel, bit=req.bit, minmax_radius=req.minmax_radius,
    )
    res["status"] = "success"
    return res


@app.post("/api/submit_run")
def submit_run_post(req: SubmitRunReq) -> dict:
    """Submit ground truth feedback. Reads XAI/forensics from cache — zero re-inference."""
    try:
        image = b64_to_pil(req.image_b64)
    except Exception as e:
        return {"status": "error", "message": f"Image decoding failed: {e}"}

    key = req.image_hash or _img_hash(image)
    cached = _cache_get(key)

    if cached:
        pred_data = cached["prediction"]
        xai_results = cached["results"]
        forensics_results = cached.get("forensics", {})
    else:
        # Fallback: run analysis if cache miss (shouldn't happen in normal flow)
        model = get_model()
        pred_data = model.predict(image)
        xai_result_objs = run_all_methods(model, image)
        xai_results = [
            {
                "name": r.name,
                "description": r.description,
                "overlay_b64": r.overlay_b64,
                "raw_heatmap_b64": r.raw_heatmap_b64,
                "metadata": r.metadata,
                "compute_time_ms": r.compute_time_ms,
            }
            for r in xai_result_objs
        ]
        forensics_results = _run_all_forensics_defaults(image)

    run_id = log_run_dataset(
        image, pred_data, xai_results,
        forensics_results=forensics_results,
        ground_truth=req.ground_truth,
        opt_in_image=req.opt_in_image,
    )

    return {
        "status": "success",
        "run_id": run_id,
        "message": f"Run {run_id} successfully submitted to dataset!",
    }


@app.get("/api/config")
def client_config():
    """Return model and space configuration for frontend initialization."""
    return {
        "model_name": "buildborderless/CommunityForensics-DeepfakeDet-ViT",
        "model_architecture": "ViT-Small (384x384, 12 layers, 6 heads)",
        "xai_methods": [
            {"id": "rollout", "name": "Attention Rollout"},
            {"id": "chefer", "name": "Gradient Attention Rollout (Chefer)"},
            {"id": "gradcam", "name": "GradCAM"},
            {"id": "integrated", "name": "Integrated Gradients"},
            {"id": "deeplift", "name": "DeepLIFT"},
            {"id": "occlusion", "name": "Occlusion Sensitivity"},
        ],
        "forensic_tools": [
            {"id": "ela", "name": "Error Level Analysis (ELA)"},
            {"id": "gradient", "name": "Spatial Edge Gradient"},
            {"id": "bitplane", "name": "Bit Plane Extractor"},
            {"id": "minmax", "name": "MinMax Deviation"},
            {"id": "wavelet", "name": "Wavelet Noise Estimation"},
        ]
    }


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


app.setup()


if __name__ == "__main__":
    app.launch(show_error=True)