"""Raw frame recorder.

Frames come straight off the sensor's raw stream as MIPI-packed 12-bit Bayer
(``SBGGR12_CSI2P``) and are written to disk verbatim: no debayer, no
compression, no loss. One camera buffer per file.

Writes happen on a dedicated thread so the capture loop never blocks on the
SSD. The bounded queue plus the camera's own buffer count provide backpressure;
if the writer falls behind (the SSD sustains ~156 MB/s, i.e. ~29 fps at
2028x1520x12-bit) frames are dropped and counted, never silently reordered.
"""

import json
import logging
import os
import queue
import threading
import time

log = logging.getLogger("recorder")


def _write_raw_buffer(request, path):
    """Write the raw stream buffer to *path*; return (bytes_written, geometry)."""
    config = request.config.get("raw", {})
    try:
        # Zero-copy path: borrow the camera's mapped buffer and write it out.
        from picamera2.request import _MappedBuffer
    except ImportError:  # pragma: no cover - older picamera2
        _MappedBuffer = None

    if _MappedBuffer is not None:
        with _MappedBuffer(request, "raw", write=False) as buf:
            view = memoryview(buf)
            with open(path, "wb") as handle:
                handle.write(view)
            return view.nbytes, config

    data = request.make_buffer("raw")
    with open(path, "wb") as handle:
        handle.write(data)
    return int(data.nbytes), config


class RawRecorder:
    """Owns the writer thread and the session index for one recording."""

    def __init__(self, storage, cfg, on_state_change=None):
        self.storage = storage
        self.cfg = cfg
        self.on_state_change = on_state_change
        self.sensor_fps = float(cfg.get("sensor.sensor_fps", 40))
        self.frame_period_ns = int(1e9 / self.sensor_fps)
        self.target_fps = int(cfg.get("record.target_fps", 30))
        self.min_free_mb = int(cfg.get("archive_min_free_mb", 500))

        self._queue = queue.Queue(maxsize=int(cfg.get("record.write_queue", 8)))
        self._thread = None
        self._active = False
        self._stop_requested = threading.Event()
        self._lock = threading.Lock()

        self.session_dir = None
        self.rotation = 0
        self.stop_reason = None
        self._raw_geometry = None
        self._index_handle = None
        self._start_monotonic = None
        self._start_iso = None
        self._stats = {}

    # ------------------------------------------------------------------ state
    @property
    def active(self) -> bool:
        return self._active

    def stats(self):
        with self._lock:
            data = dict(self._stats)
        if data and self._start_monotonic:
            elapsed = max(1e-6, time.monotonic() - self._start_monotonic)
            data["elapsed"] = elapsed
            data["actual_fps"] = data.get("frames", 0) / elapsed
            data["mb_per_s"] = data.get("bytes", 0) / elapsed / 1e6
        return data

    # ------------------------------------------------------------------ start
    def start(self, session_dir, rotation=0):
        if self._active:
            return None
        session_dir = os.fspath(session_dir)
        os.makedirs(os.path.join(session_dir, "raw"), exist_ok=True)
        self.session_dir = session_dir
        # Clockwise display rotation in degrees. Raw Bayer frames are stored
        # sensor-native, so the rotation is recorded here for whatever reads
        # them later (tools/unpack_raw.py applies it and also bakes it in when
        # writing .npy/.tiff/.png output).
        self.rotation = int(rotation or 0) % 360
        self.stop_reason = None
        self._stop_requested.clear()
        self._start_monotonic = time.monotonic()
        self._start_iso = _iso_now()
        self._stats = {
            "frames": 0, "bytes": 0, "dropped": 0, "missed": 0,
            "first_sensor_ts": None, "last_sensor_ts": None, "started": self._start_iso,
        }
        self._index_handle = open(os.path.join(session_dir, "index.jsonl"), "w", encoding="utf-8")
        self._write_session_json(final=False)
        self._active = True
        self._thread = threading.Thread(target=self._writer_loop, name="raw-writer", daemon=True)
        self._thread.start()
        log.info("recording started: %s (target %d fps)", session_dir, self.target_fps)
        if self.on_state_change:
            self.on_state_change()
        return session_dir

    # ----------------------------------------------------------------- submit
    def submit(self, request):
        """Hand a completed request to the writer. Takes ownership of *request*."""
        if not self._active:
            request.release()
            return
        try:
            self._queue.put_nowait(request)
        except queue.Full:
            with self._lock:
                self._stats["dropped"] += 1
            request.release()

    # ------------------------------------------------------------------- stop
    def stop(self, reason=None):
        if not self._active:
            return self.stats()
        self._active = False
        self.stop_reason = reason or "stopped"
        self._stop_requested.set()
        thread, self._thread = self._thread, None
        if thread is not None:
            thread.join(timeout=30)
        # Anything still queued was never written; count it honestly.
        leftovers = 0
        while True:
            try:
                request = self._queue.get_nowait()
            except queue.Empty:
                break
            leftovers += 1
            request.release()
        if leftovers:
            with self._lock:
                self._stats["dropped"] += leftovers
        if self._index_handle is not None:
            self._index_handle.close()
            self._index_handle = None
        self._write_session_json(final=True)
        summary = self.stats()
        log.info("recording stopped (%s): %d frames, %.1f MB, %d dropped, %.1f fps",
                 self.stop_reason, summary.get("frames", 0), summary.get("bytes", 0) / 1e6,
                 summary.get("dropped", 0), summary.get("actual_fps", 0.0))
        if self.on_state_change:
            self.on_state_change()
        return summary

    # ------------------------------------------------------------ writer side
    def _writer_loop(self):
        free_check_countdown = 15  # ~0.5s at 30fps
        while True:
            try:
                request = self._queue.get(timeout=0.5)
            except queue.Empty:
                if self._stop_requested.is_set():
                    break
                continue
            try:
                self._write_one(request)
            except Exception:  # a bad frame must not kill the recording
                log.exception("failed writing frame")
                with self._lock:
                    self._stats["dropped"] += 1
            finally:
                request.release()

            free_check_countdown -= 1
            if free_check_countdown <= 0:
                free_check_countdown = 15
                reason = self._check_environment()
                if reason:
                    self.stop_reason = reason
                    self._active = False
                    self._stop_requested.set()
                    break
            if self._stop_requested.is_set() and self._queue.empty():
                break

    def _check_environment(self):
        if not self.storage.is_mounted():
            log.warning("SSD disappeared while recording")
            return "ssd removed"
        free_mb = self.storage.free_mb()
        if free_mb and free_mb < self.min_free_mb:
            log.warning("SSD below minimum free space (%d MB)", free_mb)
            return "disk full"
        return None

    def _write_one(self, request):
        metadata = request.get_metadata()
        with self._lock:
            stats = self._stats
            frame_no = stats["frames"] + 1
            missed = 0
            sensor_ts = metadata.get("SensorTimestamp")
            # Judge gaps against the duration the sensor is actually running at,
            # which auto exposure can change; assuming the mode's maximum rate
            # would report a "missed" frame on every single frame in dim light.
            period_ns = self.frame_period_ns
            frame_duration_us = metadata.get("FrameDuration")
            if frame_duration_us:
                period_ns = max(1, int(float(frame_duration_us) * 1000))
            if sensor_ts is not None:
                last = stats["last_sensor_ts"]
                if last is not None and sensor_ts > last:
                    gap = sensor_ts - last
                    missed = max(0, int(round(gap / period_ns)) - 1)
                stats["last_sensor_ts"] = sensor_ts
                if stats["first_sensor_ts"] is None:
                    stats["first_sensor_ts"] = sensor_ts
            stats["missed"] += missed

        filename = "frame_%06d.raw" % frame_no
        path = os.path.join(self.session_dir, "raw", filename)
        written, raw_cfg = _write_raw_buffer(request, path)

        elapsed = time.monotonic() - self._start_monotonic
        record = {
            "frame": frame_no,
            "file": "raw/" + filename,
            "t": round(elapsed, 6),
            "sensor_ts": sensor_ts,
            "exposure_us": metadata.get("ExposureTime"),
            "gain": _round(metadata.get("AnalogueGain")),
            "awb": _round_pair(metadata.get("ColourGains")),
            "lux": _round(metadata.get("Lux")),
            "bytes": written,
            "missed_before": missed,
        }
        if frame_no % int(self.cfg.get("record.index_stride", 1)) == 0:
            self._index_handle.write(json.dumps(record, separators=(",", ":")) + "\n")

        with self._lock:
            self._stats["frames"] = frame_no
            self._stats["bytes"] += written
        return raw_cfg

    # --------------------------------------------------------------- metadata
    def set_raw_geometry(self, geometry):
        """Remember the raw stream geometry so session.json can describe it."""
        self._raw_geometry = geometry

    def _write_session_json(self, final: bool):
        summary = self.stats()
        payload = {
            "kind": "record",
            "started": self._start_iso,
            "ended": _iso_now() if final else None,
            "stop_reason": self.stop_reason if final else None,
            "sensor": self.cfg.get("sensor", {}),
            "raw": getattr(self, "_raw_geometry", None),
            "target_fps": self.target_fps,
            "rotation": getattr(self, "rotation", 0),
            "frames": summary.get("frames", 0),
            "bytes": summary.get("bytes", 0),
            "dropped": summary.get("dropped", 0),
            "missed": summary.get("missed", 0),
            "actual_fps": round(summary.get("actual_fps", 0.0), 3),
            "mb_per_s": round(summary.get("mb_per_s", 0.0), 2),
            "elapsed_s": round(summary.get("elapsed", 0.0), 3),
            "layout": "date/HH-MM-SS/<format>/",
            "frame_orientation": "sensor-native; apply 'rotation' clockwise for display",
            "notes": [
                "Frames are MIPI-packed 12-bit Bayer, one file each, no processing.",
                "'missed' counts sensor frames that never made it to disk.",
                "Use tools/unpack_raw.py on a desktop to unpack to 16-bit .npy/.tiff.",
            ],
        }
        path = os.path.join(self.session_dir, "session.json")
        tmp = path + ".tmp"
        try:
            with open(tmp, "w", encoding="utf-8") as handle:
                json.dump(payload, handle, indent=2)
                handle.write("\n")
            os.replace(tmp, path)
        except OSError:
            log.exception("could not write %s", path)


def _iso_now():
    import datetime
    return datetime.datetime.now().astimezone().isoformat(timespec="milliseconds")


def _round(value, digits=4):
    try:
        return round(float(value), digits)
    except (TypeError, ValueError):
        return None


def _round_pair(value):
    if not value:
        return None
    try:
        return [round(float(value[0]), 4), round(float(value[1]), 4)]
    except (TypeError, ValueError, IndexError):
        return None