#!/usr/bin/env python3
# /// script
# requires-python = ">=3.9"
# dependencies = ["rawpy>=0.19", "numpy>=1.24", "pillow>=10"]
# ///
"""
DYELOT raw sidecar
==================

Decodes camera raw and 16-bit TIFF on this machine and hands DYELOT a 16-bit
buffer instead of an 8-bit screenshot, plus the capture metadata the browser
cannot see.

    pip install rawpy numpy          # Pillow optional, adds TIFF/PSD reading
    python dyelot_server.py          # then open http://127.0.0.1:8788

Binds to loopback only. Nothing is written to disk and nothing leaves the
machine; files are decoded from memory and discarded.

Why the decode is set up the way it is
--------------------------------------
no_auto_bright   LibRaw's auto-exposure would silently rescale every frame and
                 make the measurements meaningless.
gamma 2.4/12.92  sRGB transfer. LibRaw applies the transfer independently of
                 the primaries, so every working space comes back sRGB-encoded
                 and the front end needs only one linearisation path.
half_size        When the output is being downsampled anyway, reading each 2x2
                 CFA block directly beats demosaicing and then throwing the
                 interpolated pixels away. Fewer invented colours.
box downsample   Averaging is the correct operator for reducing a measurement.
                 Clipping is counted at full resolution first, before any
                 averaging can hide it.
"""

import io
import json
import os
import re
import struct
import subprocess
import sys
import threading
import webbrowser
import zlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import numpy as np

try:
    import rawpy
    RAWPY_ERR = None
except Exception as e:                                    # pragma: no cover
    rawpy = None
    RAWPY_ERR = str(e)

try:
    from PIL import Image
    PIL_OK = True
except Exception:
    Image = None
    PIL_OK = False

HERE = os.path.dirname(os.path.abspath(__file__))
# The published site lives in ./public so that `wrangler deploy` cannot pick up
# source files by accident. Look there first, then alongside this script.
APP_DIR = next((d for d in (os.path.join(HERE, "public"), HERE)
                if any(os.path.exists(os.path.join(d, n)) for n in ("index.html", "dyelot.html"))), HERE)
HTML = next((os.path.join(APP_DIR, n) for n in ("index.html", "dyelot.html")
             if os.path.exists(os.path.join(APP_DIR, n))), os.path.join(APP_DIR, "index.html"))
MAX_BODY = 768 * 1024 * 1024
ALLOWED_ORIGINS = set()

RAW_EXT = {
    "3fr", "arw", "cr2", "cr3", "crw", "dcr", "dng", "erf", "fff", "iiq", "kdc",
    "mef", "mos", "mrw", "nef", "nrw", "orf", "pef", "raf", "raw", "rw2", "rwl",
    "sr2", "srf", "srw", "x3f",
}
PIL_EXT = {"tif", "tiff", "psd", "webp", "png", "jpg", "jpeg"}

# RGB -> XYZ, all D65 so the front end can stay on one white point.
PRIMARIES = {
    "srgb": [0.4124564, 0.3575761, 0.1804375,
             0.2126729, 0.7151522, 0.0721750,
             0.0193339, 0.1191920, 0.9503041],
    "adobe": [0.5767309, 0.1855540, 0.1881852,
              0.2973769, 0.6273491, 0.0752741,
              0.0270343, 0.0706872, 0.9911085],
    "rec2020": [0.6369580, 0.1446169, 0.1688810,
                0.2627002, 0.6779981, 0.0593017,
                0.0000000, 0.0280727, 1.0609851],
}
SPACE_LABEL = {"srgb": "sRGB", "adobe": "Adobe RGB (1998)", "rec2020": "Rec. 2020"}


# ----------------------------------------------------------------- utilities
def png_encode(rgb8):
    """Minimal PNG writer so Pillow stays optional."""
    h, w, _ = rgb8.shape
    rows = np.hstack([np.zeros((h, 1), np.uint8), rgb8.reshape(h, w * 3)])
    comp = zlib.compress(rows.tobytes(), 6)

    def chunk(tag, data):
        c = tag + data
        return struct.pack(">I", len(data)) + c + struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)

    return (b"\x89PNG\r\n\x1a\n"
            + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
            + chunk(b"IDAT", comp)
            + chunk(b"IEND", b""))


def box_downsample(arr, factor):
    """Integer box average. arr is HxWx3, any numeric dtype."""
    if factor <= 1:
        return arr
    h, w, c = arr.shape
    h2, w2 = h // factor * factor, w // factor * factor
    a = arr[:h2, :w2].astype(np.float32)
    a = a.reshape(h2 // factor, factor, w2 // factor, factor, c).mean(axis=(1, 3))
    return a


def clip_stats(arr16):
    """Fraction of pixels pinned at either end, measured before downsampling."""
    hi = arr16 >= 65280            # top 1/256 of the range
    lo = arr16 <= 255
    n = arr16.shape[0] * arr16.shape[1]
    return {
        "highAllPct": float(100.0 * np.count_nonzero(hi.all(axis=2)) / n),
        "lowAllPct": float(100.0 * np.count_nonzero(lo.all(axis=2)) / n),
        "highAnyPct": float(100.0 * np.count_nonzero(hi.any(axis=2)) / n),
        "lowAnyPct": float(100.0 * np.count_nonzero(lo.any(axis=2)) / n),
        "perChannelHighPct": [float(100.0 * np.count_nonzero(hi[:, :, i]) / n) for i in range(3)],
    }


EXIF_KEYS = [
    ("Make", "make"), ("Model", "model"), ("LensModel", "lens"), ("Lens", "lens"),
    ("LensID", "lens"), ("ISO", "iso"), ("ExposureTime", "shutter"),
    ("FNumber", "aperture"), ("FocalLength", "focal"), ("DateTimeOriginal", "shot"),
    ("WhiteBalance", "wbMode"), ("ColorTemperature", "wbKelvin"),
    ("ColorSpace", "colorSpace"), ("Orientation", "orientation"),
]


def exif_via_exiftool(data, suffix):
    """exiftool if it happens to be installed; it is the only reliable CR3 reader."""
    try:
        p = subprocess.run(["exiftool", "-j", "-n", "-fast2", "-"],
                           input=data, capture_output=True, timeout=25)
        if p.returncode != 0 or not p.stdout:
            return {}
        tags = json.loads(p.stdout.decode("utf-8", "replace"))[0]
    except Exception:
        return {}
    out = {}
    for src, dst in EXIF_KEYS:
        if src in tags and dst not in out and tags[src] not in (None, ""):
            out[dst] = tags[src]
    return out


def exif_via_pil(data):
    if not PIL_OK:
        return {}
    try:
        img = Image.open(io.BytesIO(data))
        raw = img.getexif()
        if not raw:
            return {}
        from PIL.ExifTags import TAGS
        named = {TAGS.get(k, k): v for k, v in raw.items()}
        out = {}
        for src, dst in EXIF_KEYS:
            if src in named and dst not in out:
                out[dst] = named[src]
        return out
    except Exception:
        return {}


def fmt_shutter(v):
    try:
        v = float(v)
    except (TypeError, ValueError):
        return v
    return f"{v:g} s" if v >= 1 else f"1/{round(1 / v):g} s"


# -------------------------------------------------------------------- decode
def decode_raw(data, wb, space, maxedge):
    if rawpy is None:
        raise RuntimeError("rawpy is not installed: " + (RAWPY_ERR or "unknown"))

    kw = dict(no_auto_bright=True, output_bps=16, gamma=(2.4, 12.92),
              output_color=getattr(rawpy.ColorSpace,
                                   {"srgb": "sRGB", "adobe": "Adobe", "rec2020": "Rec2020"}[space]))
    if wb == "camera":
        kw["use_camera_wb"] = True
    elif wb == "auto":
        kw["use_auto_wb"] = True
    elif wb == "neutral":
        kw["user_wb"] = [1.0, 1.0, 1.0, 1.0]
    # "daylight" falls through to LibRaw's built-in daylight balance

    with rawpy.imread(io.BytesIO(data)) as r:
        nat_h, nat_w = r.sizes.height, r.sizes.width
        # Reading 2x2 CFA blocks directly is both faster and free of demosaic
        # invention, so use it whenever we are downsampling past 2x anyway.
        kw["half_size"] = max(nat_w, nat_h) >= 2 * maxedge
        meta = {
            "cameraWhitebalance": [float(v) for v in r.camera_whitebalance],
            "daylightWhitebalance": [float(v) for v in r.daylight_whitebalance],
            "blackLevelPerChannel": [int(v) for v in r.black_level_per_channel],
            "whiteLevel": int(r.white_level),
            "cfaDescription": r.color_desc.decode("ascii", "replace"),
            "rawSize": [int(r.sizes.raw_width), int(r.sizes.raw_height)],
        }
        try:
            m = np.asarray(r.rgb_xyz_matrix, dtype=float)
            if np.any(m):
                meta["rgbXyzMatrix"] = [[float(x) for x in row] for row in m[:3]]
        except Exception:
            pass
        rgb = r.postprocess(**kw)

    meta["halfSize"] = bool(kw["half_size"])
    meta["nativeSize"] = [int(nat_w), int(nat_h)]
    return rgb, meta


def decode_pil(data):
    if not PIL_OK:
        raise RuntimeError("Pillow is not installed, so TIFF and PSD cannot be read")
    img = Image.open(io.BytesIO(data))
    mode16 = img.mode in ("I;16", "I;16B", "I;16L", "I", "F")
    if img.mode not in ("RGB", "RGBA") and not mode16:
        img = img.convert("RGB")
    a = np.asarray(img)
    if a.ndim == 2:
        a = np.dstack([a, a, a])
    a = a[:, :, :3]
    if a.dtype == np.uint8:
        a = a.astype(np.uint16) * 257
    elif a.dtype != np.uint16:
        a = np.clip(a.astype(np.float32), 0, 65535).astype(np.uint16)
    return a, {"nativeSize": [int(a.shape[1]), int(a.shape[0])], "halfSize": False}


def build_payload(data, filename, wb, space, maxedge):
    ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    if ext in RAW_EXT:
        rgb16, meta = decode_raw(data, wb, space, maxedge)
        meta["decoder"] = "LibRaw %s" % (".".join(str(x) for x in rawpy.libraw_version))
    elif ext in PIL_EXT:
        rgb16, meta = decode_pil(data)
        meta["decoder"] = "Pillow"
        space = "srgb"
    else:
        raise RuntimeError("Unrecognised extension .%s" % ext)

    full_h, full_w = rgb16.shape[:2]
    meta["decodedSize"] = [int(full_w), int(full_h)]
    meta["clipping"] = clip_stats(rgb16)

    factor = max(1, -(-max(full_w, full_h) // maxedge))
    small = box_downsample(rgb16, factor)
    small16 = np.clip(np.rint(small), 0, 65535).astype(np.uint16) if factor > 1 else rgb16
    h, w = small16.shape[:2]

    preview = png_encode((small16 >> 8).astype(np.uint8))

    exif = exif_via_exiftool(data, ext) or exif_via_pil(data)
    if "shutter" in exif:
        exif["shutter"] = fmt_shutter(exif["shutter"])

    meta.update({
        "ok": True,
        "filename": filename,
        "workingSize": [int(w), int(h)],
        "downsampleFactor": int(factor),
        "bits": 16,
        "whiteBalanceMode": wb,
        "space": space,
        "spaceLabel": SPACE_LABEL.get(space, space),
        "rgbToXyz": PRIMARIES.get(space, PRIMARIES["srgb"]),
        # The encoding travels with the pixels. LibRaw applies dcraw's
        # parameterised curve, which is not bit-identical to piecewise sRGB, so
        # say exactly what it is rather than letting the client assume.
        "transfer": {"kind": "dcraw", "power": 2.4, "slope": 12.92,
                     "label": "sRGB via dcraw curve (2.4 / 12.92)"},
        "exif": exif,
        "exiftool": bool(exif) and "make" in exif,
    })

    mj = json.dumps(meta).encode("utf-8")
    hdr = small16.tobytes()
    return (b"DYL1" + struct.pack("<III", len(mj), len(preview), len(hdr))
            + mj + preview + hdr)


# ---------------------------------------------------------------------- HTTP
class Handler(BaseHTTPRequestHandler):
    server_version = "DYELOT/2.0"
    protocol_version = "HTTP/1.1"

    def log_message(self, fmt, *args):
        sys.stderr.write("  %s\n" % (fmt % args))

    # Loopback-only, and only pages served from here, opened from disk, or
    # explicitly allowed with --allow-origin may call in. The server never reads
    # from the filesystem, so the worst a stray page could do is have its own
    # bytes decoded, but keeping the list short costs nothing.
    def _cors(self):
        origin = self.headers.get("Origin")
        if not origin:
            return
        ok = (origin == "null"
              or re.match(r"^https?://(127\.0\.0\.1|localhost|\[::1\])(:\d+)?$", origin)
              or origin in ALLOWED_ORIGINS)
        if ok:
            self.send_header("Access-Control-Allow-Origin", origin)
            self.send_header("Vary", "Origin")

    def _send(self, code, body, ctype, extra=None):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        for k, v in (extra or {}).items():
            self.send_header(k, v)
        self._cors()
        self.end_headers()
        self.wfile.write(body)

    def _json(self, code, obj):
        self._send(code, json.dumps(obj).encode(), "application/json")

    def do_OPTIONS(self):
        self.send_response(204)
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, X-Filename")
        self.send_header("Content-Length", "0")
        self._cors()
        self.end_headers()

    def do_GET(self):
        path = self.path.split("?", 1)[0]
        if path in ("/", "/index.html", "/dyelot.html"):
            if not os.path.exists(HTML):
                self._send(500, b"index.html not found in ./public or next to dyelot_server.py", "text/plain")
                return
            self._send(200, open(HTML, "rb").read(), "text/html; charset=utf-8")
        elif path == "/favicon.ico":
            self._send(204, b"", "image/x-icon")
        elif path == "/api/health":
            self._json(200, {
                "ok": True,
                "service": "dyelot-sidecar",
                "version": "2.0",
                "rawpy": getattr(rawpy, "__version__", None),
                "libraw": ".".join(str(x) for x in rawpy.libraw_version) if rawpy else None,
                "rawpyError": RAWPY_ERR,
                "pillow": PIL_OK,
                "exiftool": _exiftool_present(),
                "rawExtensions": sorted(RAW_EXT),
                "otherExtensions": sorted(PIL_EXT),
                "spaces": [{"id": k, "label": SPACE_LABEL[k]} for k in PRIMARIES],
            })
        else:
            self._static(path)

    # Only these names, resolved inside the app folder. No directory listing and
    # no traversal: an unlisted path is a 404 regardless of what is on disk.
    STATIC = {
        "/sw.js": "application/javascript",
        "/manifest.webmanifest": "application/manifest+json",
        "/icon-192.png": "image/png",
        "/icon-512.png": "image/png",
        "/icon-maskable.png": "image/png",
    }

    def _static(self, path):
        ctype = self.STATIC.get(path)
        if ctype is None and re.fullmatch(r"/fonts/[A-Za-z0-9._-]+\.woff2", path or ""):
            ctype = "font/woff2"
        if ctype is None:
            self._send(404, b"not found", "text/plain")
            return
        full = os.path.normpath(os.path.join(APP_DIR, path.lstrip("/")))
        if not full.startswith(APP_DIR + os.sep) or not os.path.isfile(full):
            self._send(404, b"not found", "text/plain")
            return
        self._send(200, open(full, "rb").read(), ctype)

    def do_POST(self):
        if self.path.split("?", 1)[0] != "/api/decode":
            self._send(404, b"not found", "text/plain")
            return
        try:
            n = int(self.headers.get("Content-Length") or 0)
        except ValueError:
            n = 0
        if n <= 0:
            self._json(400, {"ok": False, "error": "Empty request body"})
            return
        if n > MAX_BODY:
            self._json(413, {"ok": False, "error": "File exceeds %d MB" % (MAX_BODY // 1048576)})
            return

        q = dict(p.split("=", 1) for p in self.path.partition("?")[2].split("&") if "=" in p)
        wb = q.get("wb", "camera")
        space = q.get("space", "srgb")
        if space not in PRIMARIES:
            space = "srgb"
        if wb not in ("camera", "auto", "neutral", "daylight"):
            wb = "camera"
        try:
            maxedge = max(512, min(6000, int(q.get("maxedge", 2560))))
        except ValueError:
            maxedge = 2560

        filename = self.headers.get("X-Filename") or "upload.raw"
        filename = os.path.basename(filename)[:200]

        body = bytearray()
        while len(body) < n:
            chunk = self.rfile.read(min(1 << 20, n - len(body)))
            if not chunk:
                break
            body += chunk

        try:
            payload = build_payload(bytes(body), filename, wb, space, maxedge)
        except Exception as e:
            self._json(422, {"ok": False, "error": explain(e, filename)})
            return
        self._send(200, payload, "application/octet-stream",
                   {"X-Dyelot-Filename": re.sub(r"[^\x20-\x7e]", "?", filename)})


def explain(e, filename):
    """Turn LibRaw's terse failures into something actionable."""
    msg = str(e) or e.__class__.__name__
    low = msg.lower()
    if "unsupported file format" in low or "fileunsupported" in low:
        return (f"LibRaw {'.'.join(str(x) for x in rawpy.libraw_version)} does not recognise "
                f"{filename}. Very new bodies often need a newer LibRaw: try "
                f"'pip install -U rawpy', or convert the file to DNG with Adobe DNG Converter, "
                f"which LibRaw always reads.")
    if "rawpy is not installed" in low:
        return msg + ". Run: pip install rawpy numpy"
    if "cannot identify image file" in low or "unrecognised extension" in low:
        return f"Could not read {filename}. Supported: raw files, TIFF, PSD, PNG, JPEG."
    if isinstance(e, MemoryError):
        return "Ran out of memory decoding this file. Lower the working resolution and retry."
    return f"{filename}: {msg}"


def _exiftool_present():
    try:
        subprocess.run(["exiftool", "-ver"], capture_output=True, timeout=5)
        return True
    except Exception:
        return False


def main():
    port = 8788
    host = "127.0.0.1"
    args = sys.argv[1:]
    if args and args[0].isdigit():
        port = int(args[0])
    for i, a in enumerate(args):
        if a == "--allow-origin" and i + 1 < len(args):
            ALLOWED_ORIGINS.add(args[i + 1].rstrip("/"))
        elif a.startswith("--allow-origin="):
            ALLOWED_ORIGINS.add(a.split("=", 1)[1].rstrip("/"))
    if "--help" in args or "-h" in args:
        print(__doc__)
        print("usage: dyelot_server.py [PORT] [--no-browser] [--allow-origin URL]\n")
        return
    if "--no-browser" not in args:
        threading.Timer(0.8, lambda: webbrowser.open(f"http://{host}:{port}/")).start()

    srv = ThreadingHTTPServer((host, port), Handler)
    print("DYELOT sidecar")
    print(f"  serving      http://{host}:{port}/   from {APP_DIR}")
    print(f"  rawpy        {getattr(rawpy, '__version__', 'NOT INSTALLED')}"
          + (f"  ·  LibRaw {'.'.join(str(x) for x in rawpy.libraw_version)}" if rawpy else ""))
    print(f"  pillow       {'yes' if PIL_OK else 'no  (TIFF and PSD reading disabled)'}")
    print(f"  exiftool     {'yes' if _exiftool_present() else 'no  (capture metadata will be sparse)'}")
    if ALLOWED_ORIGINS:
        print(f"  also serving {', '.join(sorted(ALLOWED_ORIGINS))}")
    if rawpy is None:
        print(f"  !! raw decoding unavailable: {RAWPY_ERR}")
    print("  loopback only. Ctrl-C to stop.\n")
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        print("\nstopped")


if __name__ == "__main__":
    main()