"""Full-resolution still capture: DNG + 16-bit TIFF + JPEG.

The IMX477 is a 12-bit sensor and the Pi's ISP only outputs 8-bit processed
images, so a genuinely 16-bit TIFF cannot come from the ISP. Instead the TIFF is
built from the raw Bayer frame through the same chain a raw converter would use:

    unpack 12-bit MIPI -> black-level subtract -> white balance
        -> colour-correction matrix -> bilinear demosaic -> sRGB gamma -> 16-bit

The result is a normal, viewable 16-bit image rather than linear sensor data.
The untouched sensor data stays in the DNG next to it, so nothing is lost for
quantitative work.

The slow part (demosaic + 74 MB TIFF write) runs on a worker thread so the live
preview resumes as soon as the capture itself is done.
"""

import datetime as _dt
import json
import logging
import os
import queue
import threading
import time

import numpy as np

log = logging.getLogger("photos")

_PATTERN_OFFSETS = {
    # pattern -> (row parity, col parity) of the red sample in the 2x2 tile
    "RGGB": (0, 0),
    "BGGR": (1, 1),
    "GRBG": (0, 1),
    "GBRG": (1, 0),
}


def pattern_from_format(fmt: str) -> str:
    """'SRGGB12_CSI2P' -> 'RGGB'."""
    if not fmt:
        return "RGGB"
    text = fmt.upper()
    if text.startswith("S"):
        text = text[1:]
    return text[:4] if text[:4] in _PATTERN_OFFSETS else "RGGB"


def unpack_mipi12(packed: np.ndarray, width: int, height: int, stride: int) -> np.ndarray:
    """Unpack MIPI CSI-2 packed 12-bit Bayer into a uint16 (height, width) array.

    Two pixels live in three bytes with the *first* pixel in the high bits:

        byte0        byte1                byte2
        p0[11:4]     p0[3:0] p1[11:8]     p1[7:0]

    so even pixels are ``(b0 << 4) | (b1 >> 4)`` and odd pixels are
    ``((b1 & 0x0F) << 8) | b2``. Getting this order wrong still yields values in
    the 0-4095 range, but it destroys the image: the tell-tale is that the two
    green sites of the Bayer tile stop matching each other.
    """
    rows = packed.reshape(height, stride)
    useful = rows[:, : (width * 3) // 2]
    triplets = useful.reshape(height, width // 2, 3).astype(np.uint16)
    b0 = triplets[:, :, 0]
    b1 = triplets[:, :, 1]
    b2 = triplets[:, :, 2]
    even = (b0 << 4) | (b1 >> 4)
    odd = ((b1 & 0x0F) << 8) | b2
    out = np.empty((height, width), dtype=np.uint16)
    out[:, 0::2] = even
    out[:, 1::2] = odd
    return out


def demosaic(raw: np.ndarray, pattern: str = "BGGR") -> np.ndarray:
    """Bilinear demosaic of a Bayer plane into a uint16 (h, w, 3) RGB image.

    Quantisation-free for 12-bit input (max 4095) and cheap enough (~1 s at
    12 MP) to run on the Pi.
    """
    rp, cp = _PATTERN_OFFSETS.get(pattern, (0, 0))
    h, w = raw.shape
    tiles = raw.reshape(h // 2, 2, w // 2, 2).transpose(0, 2, 1, 3)
    red = tiles[..., rp, cp].astype(np.uint32)
    blue = tiles[..., 1 - rp, 1 - cp].astype(np.uint32)
    green_r = tiles[..., rp, 1 - cp].astype(np.uint32)
    green_b = tiles[..., 1 - rp, cp].astype(np.uint32)

    def padded(plane):
        return np.pad(plane, 1, mode="edge")

    def upsample(plane):
        p = padded(plane)
        center = p[1:-1, 1:-1]
        left = p[1:-1, :-2]
        right = p[1:-1, 2:]
        up = p[:-2, 1:-1]
        down = p[2:, 1:-1]
        horiz = (center + (right if cp == 0 else left)) // 2
        vert = (center + (down if rp == 0 else up)) // 2
        col_n = right if cp == 0 else left
        row_n = down if rp == 0 else up
        if cp == 0 and rp == 0:
            corner = p[2:, 2:]
        elif cp == 0 and rp == 1:
            corner = p[2:, :-2]
        elif cp == 1 and rp == 0:
            corner = p[:-2, 2:]
        else:
            corner = p[:-2, :-2]
        diag = (center + col_n + row_n + corner) // 4
        full = np.empty((h, w), dtype=np.uint16)
        full[rp::2, cp::2] = center
        full[rp::2, (1 - cp)::2] = horiz
        full[(1 - rp)::2, cp::2] = vert
        full[(1 - rp)::2, (1 - cp)::2] = diag
        return full

    red_full = upsample(red)
    blue_full = upsample(blue)

    g1 = padded(green_r)
    g2 = padded(green_b)
    g1_left = (g1[1:-1, 1:-1] + g1[1:-1, :-2]) // 2
    g2_right = (g2[1:-1, 1:-1] + g2[1:-1, 2:]) // 2
    green_full = np.empty((h, w), dtype=np.uint16)
    green_full[rp::2, 1 - cp::2] = green_r
    green_full[1 - rp::2, cp::2] = green_b
    green_full[rp::2, cp::2] = g1_left
    green_full[1 - rp::2, 1 - cp::2] = g2_right

    rgb = np.empty((h, w, 3), dtype=np.uint16)
    rgb[..., 0] = red_full
    rgb[..., 1] = green_full
    rgb[..., 2] = blue_full
    return rgb


def srgb_encode(linear: np.ndarray) -> np.ndarray:
    """Apply the sRGB transfer function (linear light -> display values)."""
    out = np.empty_like(linear)
    low = linear <= 0.0031308
    np.multiply(linear, 12.92, out=out, where=low)
    high = ~low
    np.power(np.clip(linear, 1e-8, None), 1.0 / 2.4, out=out, where=high)
    np.multiply(out, 1.055, out=out, where=high)
    np.subtract(out, 0.055, out=out, where=high)
    return out


def render_still(raw, pattern, black=256.0, white=4095.0, gains=None, ccm=None,
                 gamma=True) -> np.ndarray:
    """Turn a 12-bit Bayer frame into a viewable 16-bit RGB image.

    Mirrors what a raw converter does, in the order the ISP does it:
    black-level subtract, white balance, colour matrix, demosaic, gamma.

    ``gains`` is picamera2's ``ColourGains`` (red, blue; green is 1.0) and
    ``ccm`` is its ``ColourCorrectionMatrix`` (3x3, camera RGB -> sRGB). Either
    may be None, in which case that step is skipped.
    """
    rgb = demosaic(raw, pattern)
    img = rgb.astype(np.float32)
    del rgb

    scale = 1.0 / float(white - black)
    img *= scale
    img -= black * scale
    np.clip(img, 0.0, None, out=img)

    if gains is not None:
        img[..., 0] *= float(gains[0])
        img[..., 2] *= float(gains[1])

    if ccm is not None:
        matrix = np.asarray(ccm, dtype=np.float32).reshape(3, 3)
        img = img @ matrix.T
        np.clip(img, 0.0, None, out=img)

    # Highlight clipping happens here, as it does in the ISP.
    np.clip(img, 0.0, 1.0, out=img)
    if gamma:
        img = srgb_encode(img)

    return np.clip(img * 65535.0 + 0.5, 0, 65535).astype(np.uint16)


def _rotate_array(rgb, rotation):
    """Rotate an (h, w, 3) array clockwise by 0/90/180/270 degrees."""
    rotation = int(rotation) % 360
    if rotation == 0:
        return rgb
    if rotation == 90:
        return np.ascontiguousarray(np.rot90(rgb, k=3))
    if rotation == 180:
        return np.ascontiguousarray(np.rot90(rgb, k=2))
    if rotation == 270:
        return np.ascontiguousarray(np.rot90(rgb, k=1))
    return rgb


def _pil_transpose(rotation):
    """PIL constant that rotates an image clockwise by *rotation* degrees.

    PIL's ROTATE_* constants turn counter-clockwise, hence the inversion.
    """
    from PIL import Image

    transpose = getattr(Image, "Transpose", Image)
    return {90: transpose.ROTATE_270, 180: transpose.ROTATE_180,
            270: transpose.ROTATE_90}[int(rotation) % 360]


def rotate_jpeg(path, rotation):
    """Rotate a JPEG on disk in place (used for the ISP-written still)."""
    rotation = int(rotation) % 360
    if rotation == 0:
        return False
    from PIL import Image

    with Image.open(path) as image:
        turned = image.transpose(_pil_transpose(rotation))
        turned.save(path, quality=95, subsampling=0)
    return True


def set_dng_orientation(path, rotation):
    """Write the DNG/TIFF Orientation tag (274) in place.

    Rotating the Bayer mosaic itself would change the sensor pattern, so the raw
    DNG keeps its pixels and records the rotation the way every raw converter
    expects: as metadata. Orientation 1/3/6/8 = 0/180/90/270 degrees.
    """
    value = {0: 1, 90: 6, 180: 3, 270: 8}.get(int(rotation) % 360)
    if value in (None, 1):
        return False
    try:
        with open(path, "r+b") as handle:
            header = handle.read(8)
            if header[:2] not in (b"II", b"MM"):
                return False
            order = "little" if header[:2] == b"II" else "big"
            ifd_offset = int.from_bytes(header[4:8], order)
            handle.seek(ifd_offset)
            count = int.from_bytes(handle.read(2), order)
            entries = handle.read(12 * count)
            for index in range(count):
                entry = entries[index * 12:(index + 1) * 12]
                tag = int.from_bytes(entry[0:2], order)
                if tag == 274:
                    handle.seek(ifd_offset + 2 + index * 12 + 8)
                    handle.write(value.to_bytes(2, order) + b"\x00\x00")
                    return True
    except OSError:
        log.exception("could not set DNG orientation on %s", path)
    return False


class TiffWorker(threading.Thread):
    """Serial background worker for the expensive demosaic + TIFF write."""

    def __init__(self, on_done=None):
        super().__init__(name="tiff-writer", daemon=True)
        self._jobs = queue.Queue()
        self._on_done = on_done

    def submit(self, job):
        self._jobs.put(job)

    def stop(self, timeout=90):
        """Finish any queued job, then end the thread."""
        if self.is_alive():
            self._jobs.put(None)
            self.join(timeout=timeout)

    def run(self):
        while True:
            job = self._jobs.get()
            if job is None:
                log.info("TIFF worker stopping")
                return
            try:
                self._do(job)
            except Exception:
                log.exception("TIFF job failed")
                if self._on_done:
                    self._on_done(job, None, "TIFF failed (see log)")

    def _do(self, job):
        import tifffile

        started = time.monotonic()
        raw = unpack_mipi12(job["packed"], job["width"], job["height"], job["stride"])
        rgb = render_still(
            raw,
            job["pattern"],
            black=job.get("black_level", 256.0),
            white=job.get("white_level", 4095.0),
            gains=job.get("gains") if job.get("apply_white_balance", True) else None,
            ccm=job.get("ccm") if job.get("apply_ccm", True) else None,
            gamma=job.get("gamma", True),
        )
        del raw
        rotation = int(job.get("rotation", 0) or 0)
        if rotation:
            rgb = _rotate_array(rgb, rotation)

        description = (
            "16-bit RGB from the 12-bit Bayer frame; "
            f"pattern={job['pattern']} black={job.get('black_level')} "
            f"white={job.get('white_level')} "
            f"white_balance={job.get('gains') if job.get('apply_white_balance', True) else None} "
            f"colour_matrix={'applied' if job.get('apply_ccm', True) and job.get('ccm') else 'none'} "
            f"transfer={'sRGB' if job.get('gamma', True) else 'linear'}; "
            "raw sensor data preserved in the DNG alongside."
        )
        tifffile.imwrite(
            job["path"], rgb, photometric="rgb",
            description=description, compression=None,
        )
        elapsed = time.monotonic() - started
        log.info("TIFF written in %.1fs: %s", elapsed, job["path"])
        if self._on_done:
            self._on_done(job, elapsed, None)


def capture_still(camera, storage, cfg, tiff_worker, on_tiff_done=None, rotation=0):
    """Capture one full-resolution still. Must run on the camera thread.

    *rotation* (clockwise degrees) is applied to the saved output: the JPEG is
    turned on disk, the TIFF is turned before it is written, and the DNG carries
    it as an Orientation tag (its Bayer data must stay sensor-native).
    """
    photo_cfg = cfg["photo"]
    rotation = int(rotation or 0) % 360
    session = storage.new_session("photo")
    raw_dir = storage.format_dir(session, "raw")
    serial = storage.next_serial(raw_dir, "still")
    stem = "still_%04d" % serial

    t_switch = time.monotonic()
    camera.switch_to_still()
    settle = max(0, int(photo_cfg.get("settle_frames", 3)))
    for _ in range(settle):
        stale = camera.capture_request()
        stale.release()
    switch_seconds = time.monotonic() - t_switch

    t_capture = time.monotonic()
    request = camera.capture_request()
    metadata = request.get_metadata()
    written = {}

    try:
        if photo_cfg.get("write_dng", True):
            path = raw_dir / (stem + ".dng")
            request.save_dng(str(path), "raw")
            if rotation:
                set_dng_orientation(str(path), rotation)
            written["dng"] = str(path)

        if photo_cfg.get("write_jpeg", True):
            jpeg_dir = storage.format_dir(session, "jpeg")
            path = jpeg_dir / (stem + ".jpg")
            request.save("main", str(path))
            if rotation:
                rotate_jpeg(str(path), rotation)
            written["jpeg"] = str(path)

        packed = None
        if photo_cfg.get("write_tiff", True):
            raw_cfg = camera.stream_config("raw") or {}
            packed = np.array(request.make_buffer("raw"), dtype=np.uint8)
            width = int(raw_cfg.get("size", (photo_cfg["width"], photo_cfg["height"]))[0])
            height = int(raw_cfg.get("size", (photo_cfg["width"], photo_cfg["height"]))[1])
            stride = int(raw_cfg.get("stride") or (width * 3 // 2))
            pattern = pattern_from_format(raw_cfg.get("format", "SRGGB12"))
            ccm = metadata.get("ColourCorrectionMatrix")
            tiff_dir = storage.format_dir(session, "tiff")
            job = {
                "packed": packed,
                "width": width,
                "height": height,
                "stride": stride,
                "pattern": pattern,
                "gains": metadata.get("ColourGains"),
                "apply_white_balance": photo_cfg.get("tiff_apply_white_balance", True),
                "ccm": list(ccm) if ccm is not None else None,
                "apply_ccm": photo_cfg.get("tiff_apply_colour_matrix", True),
                "gamma": photo_cfg.get("tiff_apply_gamma", True),
                "black_level": photo_cfg.get("black_level", 256),
                "white_level": photo_cfg.get("white_level", 4095),
                "rotation": rotation,
                "path": str(tiff_dir / (stem + ".tiff")),
                "stem": stem,
                "session": str(session),
            }
            tiff_worker.submit(job)
            written["tiff"] = job["path"]
    finally:
        request.release()
        camera.switch_to_video()

    summary = {
        "kind": "photo",
        "stem": stem,
        "session": str(session),
        "taken": _dt.datetime.now().astimezone().isoformat(timespec="milliseconds"),
        "files": {k: storage.relative(v) for k, v in written.items()},
        "mode": "4056x3040",
        "rotation": rotation,
        "capture_ms": round((time.monotonic() - t_capture) * 1000),
        "switch_ms": round(switch_seconds * 1000),
        "exposure_us": metadata.get("ExposureTime"),
        "gain": metadata.get("AnalogueGain"),
        "awb": list(metadata.get("ColourGains")) if metadata.get("ColourGains") else None,
        "colour_correction_matrix": list(metadata["ColourCorrectionMatrix"])
        if metadata.get("ColourCorrectionMatrix") is not None else None,
        "black_level": photo_cfg.get("black_level", 256),
        "white_level": photo_cfg.get("white_level", 4095),
        "sensor_temperature": metadata.get("SensorTemperature"),
        "notes": [
            "DNG holds the untouched 12-bit Bayer frame.",
            "TIFF is 16-bit sRGB, demosaiced from that Bayer frame with white balance",
            "and the colour-correction matrix applied.",
            "JPEG is the ISP-processed 8-bit image.",
            "rotation is the clockwise display/export rotation in degrees.",
        ],
    }
    _write_json(session / "photo.json", summary)
    return summary


def _write_json(path, payload):
    tmp = str(path) + ".tmp"
    try:
        with open(tmp, "w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2, default=str)
            handle.write("\n")
        os.replace(tmp, path)
    except OSError:
        log.exception("could not write %s", path)