"""
forensic_scope_api.py — web-friendly version of forensic_scope.py's core
analysis (ELA, edge detection, local sharpness, noise map). Same exact
math/logic, just with the matplotlib GUI and CLI file-saving stripped out
and replaced with a function that returns base64-encoded PNGs + stats,
ready to be sent back as JSON from a Flask route.

This is a MANUAL VISUAL INSPECTION tool ("magnifier"), not an automated
scoring signal — unlike analysis_engine.py's checks, it doesn't return a
suspicion score. It hands the user visual maps of a region they select,
so THEY judge it, the same way forensic_scope.py's desktop GUI did.
"""

import base64
import io

import numpy as np
import cv2
from PIL import Image


def load_image_rgb_from_bytes(data):
    img = Image.open(io.BytesIO(data)).convert("RGB")
    return np.array(img)


def compute_ela_from_bytes(data, quality=90, scale=15):
    """Same ELA logic as forensic_scope.py's compute_ela(), but works
    directly from in-memory bytes (no temp file needed for the original,
    only for the JPEG-recompressed comparison)."""
    orig = Image.open(io.BytesIO(data)).convert("RGB")
    buf = io.BytesIO()
    orig.save(buf, "JPEG", quality=quality)
    buf.seek(0)
    resaved = Image.open(buf).convert("RGB")

    orig_arr = np.array(orig).astype(np.int16)
    resaved_arr = np.array(resaved).astype(np.int16)

    diff = np.abs(orig_arr - resaved_arr)
    diff = diff * scale
    diff = np.clip(diff, 0, 255).astype(np.uint8)
    return diff


def compute_edges(gray):
    return cv2.Canny(gray, 50, 150)


def compute_local_sharpness(gray, ksize=9):
    lap = cv2.Laplacian(gray, cv2.CV_64F)
    lap_abs = np.abs(lap)
    heat = cv2.GaussianBlur(lap_abs, (ksize, ksize), 0)
    heat = heat / (heat.max() + 1e-6) * 255
    return heat.astype(np.uint8)


def compute_noise_map(gray, ksize=5):
    blurred = cv2.medianBlur(gray, ksize)
    noise = cv2.absdiff(gray, blurred)
    noise = cv2.normalize(noise, None, 0, 255, cv2.NORM_MINMAX)
    return noise


def compute_stats(gray, sharpness_map, noise_map):
    return {
        "sharpness_mean": float(np.mean(sharpness_map)),
        "sharpness_std": float(np.std(sharpness_map)),
        "noise_mean": float(np.mean(noise_map)),
        "noise_std": float(np.std(noise_map)),
        "brightness_mean": float(np.mean(gray)),
    }


def _encode_png_base64(arr, colormap=None):
    """Encodes a numpy array (grayscale or RGB) as a base64 PNG data URL."""
    if arr is None:
        return None
    if colormap is not None:
        # Apply an OpenCV colormap for heatmap-style visuals (edges/
        # sharpness/noise look much clearer in color than raw grayscale).
        arr_uint8 = arr.astype(np.uint8) if arr.dtype != np.uint8 else arr
        colored = cv2.applyColorMap(arr_uint8, colormap)
        success, buf = cv2.imencode(".png", colored)
    elif arr.ndim == 2:
        success, buf = cv2.imencode(".png", arr)
    else:
        success, buf = cv2.imencode(".png", cv2.cvtColor(arr, cv2.COLOR_RGB2BGR))
    if not success:
        return None
    b64 = base64.b64encode(buf.tobytes()).decode("ascii")
    return f"data:image/png;base64,{b64}"


def analyze_region_for_web(image_bytes, box=None, ela_quality=90, ela_scale=15, zoom_factor=4):
    """
    Runs the same forensic maps as forensic_scope.py's analyze_region(),
    on an in-memory image, and returns everything as base64 PNG data URLs
    + stats — directly JSON-serializable for a Flask response.

    box: optional (x1, y1, x2, y2) in ORIGINAL image pixel coordinates.
         If None, analyzes the whole image.
    """
    rgb_full = load_image_rgb_from_bytes(image_bytes)
    full_ela = compute_ela_from_bytes(image_bytes, quality=ela_quality, scale=ela_scale)

    if box is not None:
        x1, y1, x2, y2 = [int(v) for v in box]
        x1, x2 = sorted((x1, x2))
        y1, y2 = sorted((y1, y2))
        crop = rgb_full[y1:y2, x1:x2].copy()
        ela_crop = full_ela[y1:y2, x1:x2]
    else:
        crop = rgb_full.copy()
        ela_crop = full_ela

    if crop.size == 0:
        raise ValueError("Selected region is empty — check the coordinates")

    gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)

    zoomed = cv2.resize(crop, None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_CUBIC)
    edges = cv2.resize(compute_edges(gray), None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_NEAREST)
    sharpness = cv2.resize(compute_local_sharpness(gray), None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_CUBIC)
    noise = cv2.resize(compute_noise_map(gray), None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_CUBIC)
    ela_zoomed = cv2.resize(ela_crop, None, fx=zoom_factor, fy=zoom_factor, interpolation=cv2.INTER_CUBIC)

    stats_gray = cv2.cvtColor(crop, cv2.COLOR_RGB2GRAY)
    sharpness_raw = compute_local_sharpness(stats_gray)
    noise_raw = compute_noise_map(stats_gray)
    stats = compute_stats(stats_gray, sharpness_raw, noise_raw)

    return {
        "zoomed": _encode_png_base64(zoomed),
        "edges": _encode_png_base64(edges, colormap=None),
        "sharpness": _encode_png_base64(sharpness, colormap=cv2.COLORMAP_INFERNO),
        "noise": _encode_png_base64(noise, colormap=cv2.COLORMAP_INFERNO),
        "ela": _encode_png_base64(ela_zoomed),
        "stats": stats,
        "box": list(box) if box else None,
    }