"""Active Learning storage writer (Contextra).

Specced in configs/system_config.yaml :: active_learning:
  raw_frames    -> JPEG of every decoded frame (q85, retention 30d)
  crop_objects  -> per-object BBox crops (retention 30d)
  metadata_json -> JSON backup of the merged-frame metadata
  annotations   -> YOLO-format label files from detections
  labeled       -> operator-tagged results (written by the frontend, NOT here)

Design
------
All file I/O runs in ONE background worker thread so the GStreamer callback
and the metadata persist loop stay fast.

* Pixel work (raw frame + crops) is fed from the det pipeline's consume(),
  the only place the decoded frame is reachable (buffer.extract(0) -> GPU
  tensor -> HWC numpy). The producer converts the GPU tensor to numpy and
  enqueues (fn, hwc_array, dets, site, stream).
* metadata_json is fed from StorageWriter.persist(), where the fully merged
  frame (det+pose+seg) is available.

Storage layout (filenames embed site/stream to stay collision-free):
  <raw_frames>/<site>__<stream>__<fn:08d>.jpg
  <crop_objects>/<site>__<stream>__<fn:08d>__<idx:02d>.jpg
  <annotations>/<site>__<stream>__<fn:08d>.txt
  <metadata_json>/<pkey>.json

Retention: files older than retention_days are pruned once at startup.
Every failure is non-fatal (logged, never raises into the pipeline).
"""
import json
import os
import queue
import threading
import time

import cv2
import numpy as np

CONFIG_PATHS = (
    os.environ.get("SYSTEM_CONFIG_PATH", "/etc/contextra/system_config.yaml"),
)

DEFAULTS = {
    "enabled": False,
    "raw_frames": {
        "enabled": False, "quality": 85, "retention_days": 30,
        "storage_path": "/data/active_learning/raw_frames/",
    },
    "crop_objects": {
        "enabled": False, "retention_days": 30,
        "storage_path": "/data/active_learning/crop_objects/",
    },
    "metadata_json": {
        "enabled": False,
        "storage_path": "/data/active_learning/metadata_json/",
    },
    "annotations": {
        "enabled": False, "label_format": "yolo",
        "storage_path": "/data/active_learning/annotations/",
    },
}

_QUEUE_MAX = 64  # bounded; drop-with-count on backpressure


def _deep_merge(base: dict, over: dict) -> dict:
    out = dict(base)
    for k, v in (over or {}).items():
        if isinstance(v, dict) and isinstance(out.get(k), dict):
            out[k] = _deep_merge(out[k], v)
        else:
            out[k] = v
    return out


def load_config():
    """Return the active_learning config block (yaml over DEFAULTS)."""
    cfg = DEFAULTS
    for p in CONFIG_PATHS:
        if os.path.isfile(p):
            try:
                import yaml
                with open(p, "r", encoding="utf-8") as fh:
                    data = yaml.safe_load(fh) or {}
                al = data.get("active_learning") or {}
                cfg = _deep_merge(DEFAULTS, al)
            except Exception as e:  # noqa: BLE001
                print(f"[AL] config load failed ({p}): {e}; using defaults", flush=True)
            break
    return cfg


class ActiveLearningWriter:
    """Queue + worker thread for active-learning artifact writes."""

    def __init__(self, config=None, site_id="SITE01", stream_id="CAM_01", enabled=None):
        self.cfg = config or load_config()
        self.site_id = site_id
        self.stream_id = stream_id
        self.enabled = bool(self.cfg.get("enabled", False)) if enabled is None else bool(enabled)
        self._queue = queue.Queue(maxsize=_QUEUE_MAX)
        self._stop = threading.Event()
        self._thread = None
        self._stats = {"frames": 0, "crops": 0, "annots": 0, "meta": 0, "dropped": 0, "errors": 0}
        self._dirs = {}
        if not self.enabled:
            return
        for key in ("raw_frames", "crop_objects", "metadata_json", "annotations"):
            sub = self.cfg.get(key, {}) or {}
            if not sub.get("enabled", False):
                continue
            path = sub.get("storage_path")
            if not path:
                continue
            try:
                os.makedirs(path, exist_ok=True)
                if not os.access(path, os.W_OK):
                    raise OSError("not writable")
                self._dirs[key] = os.path.normpath(path)
            except Exception as e:  # noqa: BLE001
                print(f"[AL] {key} dir unavailable ({path}): {e}", flush=True)
                sub["enabled"] = False
        if not self._dirs:
            print("[AL] no writable artifact dirs; active_learning DISABLED", flush=True)
            self.enabled = False
            return
        self._retain()
        self._thread = threading.Thread(target=self._worker, name="active-learning", daemon=True)
        self._thread.start()
        print(f"[AL] active_learning ON dirs={list(self._dirs)}", flush=True)

    # -- retention ---------------------------------------------------------
    def _retain(self):
        now = time.time()
        for key, path in self._dirs.items():
            days = int((self.cfg.get(key, {}) or {}).get("retention_days", 30) or 30)
            cutoff = now - days * 86400
            try:
                n = 0
                for name in os.listdir(path):
                    fp = os.path.join(path, name)
                    try:
                        if os.path.isfile(fp) and os.path.getmtime(fp) < cutoff:
                            os.remove(fp)
                            n += 1
                    except Exception:  # noqa: BLE001
                        pass
                if n:
                    print(f"[AL] retention {key}: removed {n} stale files", flush=True)
            except Exception as e:  # noqa: BLE001
                print(f"[AL] retention {key} failed: {e}", flush=True)

    # -- producers ---------------------------------------------------------
    def submit_frame(self, fn, hwc, dets, site_id=None, stream_id=None):
        """Enqueue a decoded frame (HWC uint8) + its det boxes for pixel work."""
        if not self.enabled:
            return
        try:
            self._queue.put_nowait({
                "type": "frame", "fn": int(fn), "hwc": hwc, "dets": dets or [],
                "site": site_id or self.site_id, "stream": stream_id or self.stream_id,
            })
        except queue.Full:
            self._stats["dropped"] += 1

    def submit_meta(self, pkey, payload):
        """Enqueue a metadata-JSON backup (built from the merged frame in persist)."""
        if not self.enabled or "metadata_json" not in self._dirs:
            return
        try:
            self._queue.put_nowait({"type": "meta", "pkey": pkey, "payload": payload})
        except queue.Full:
            self._stats["dropped"] += 1

    # -- worker ------------------------------------------------------------
    def _worker(self):
        q = self._queue
        while not self._stop.is_set():
            try:
                job = q.get(timeout=0.5)
            except queue.Empty:
                continue
            try:
                if job["type"] == "frame":
                    self._do_frame(job)
                elif job["type"] == "meta":
                    self._do_meta(job)
            except Exception as e:  # noqa: BLE001
                self._stats["errors"] += 1
                print(f"[AL] worker job failed (non-fatal): {e}", flush=True)
            finally:
                q.task_done()

    def _do_frame(self, job):
        fn = job["fn"]
        hwc = job["hwc"]
        dets = job["dets"]
        tag = f"{job['site']}__{job['stream']}__{fn:08d}"
        # raw frame
        if "raw_frames" in self._dirs:
            rf = self.cfg["raw_frames"]
            q = int(rf.get("quality", 85))
            fp = os.path.join(self._dirs["raw_frames"], tag + ".jpg")
            if cv2.imwrite(fp, hwc, [cv2.IMWRITE_JPEG_QUALITY, q]):
                self._stats["frames"] += 1
        # crops + yolo annotations (need det boxes)
        if dets and hwc is not None:
            H, W = hwc.shape[:2]
            yolo_lines = []
            if "crop_objects" in self._dirs:
                cd = self._dirs["crop_objects"]
                for i, d in enumerate(dets):
                    cx, cy, bw, bh = (float(v) for v in d["box"])
                    x0 = int(np.clip((cx - bw / 2) * W, 0, W - 1))
                    y0 = int(np.clip((cy - bh / 2) * H, 0, H - 1))
                    x1 = int(np.clip((cx + bw / 2) * W, x0 + 1, W))
                    y1 = int(np.clip((cy + bh / 2) * H, y0 + 1, H))
                    if x1 - x0 < 2 or y1 - y0 < 2:
                        continue
                    fp = os.path.join(cd, f"{tag}__{i:02d}.jpg")
                    if cv2.imwrite(fp, hwc[y0:y1, x0:x1]):
                        self._stats["crops"] += 1
            if "annotations" in self._dirs and dets:
                fmt = self.cfg["annotations"].get("label_format", "yolo")
                if fmt == "coco":
                    self._write_coco(tag, fn, hwc, dets)
                elif fmt == "yolo":
                    yolo_lines = []
                    for d in dets:
                        cx, cy, bw, bh = (float(v) for v in d["box"])
                        cid = int(d["class_id"])
                        ycls = cid - 1 if cid >= 1 else 0
                        yolo_lines.append(f"{ycls} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
                    if yolo_lines:
                        fp = os.path.join(self._dirs["annotations"], tag + ".txt")
                        with open(fp, "w", encoding="utf-8") as fh:
                            fh.write("\n".join(yolo_lines) + "\n")
                        self._stats["annots"] += 1

    def _write_coco(self, tag, fn, hwc, dets):
        """Emit one COCO-format JSON per frame (images/annotations/categories).

        Single image entry = the full raw frame (<tag>.jpg in raw_frames);
        N annotations reference it. bbox = pixel xywh (COCO), category_id = canonical 0-based.
        """
        from app.multi_pipeline import _coco_name
        H, W = hwc.shape[:2]
        if not dets:
            return
        annotations = []
        for i, d in enumerate(dets):
            cx, cy, bw, bh = (float(v) for v in d["box"])
            x0 = max((cx - bw / 2) * W, 0.0)
            y0 = max((cy - bh / 2) * H, 0.0)
            wpx = min(bw * W, W - x0)
            hpx = min(bh * H, H - y0)
            if wpx <= 0 or hpx <= 0:
                continue
            annotations.append({
                "id": i,
                "image_id": 0,
                "category_id": int(d["class_id"]),
                "bbox": [round(x0, 2), round(y0, 2), round(wpx, 2), round(hpx, 2)],
                "area": round(wpx * hpx, 2),
                "iscrowd": 0,
            })
        if not annotations:
            return
        cats = sorted({int(d["class_id"]) for d in dets})
        doc = {
            "info": {"description": f"contextra active-learning fn={fn} site={tag}"},
            "images": [{"id": 0, "file_name": tag + ".jpg", "width": int(W), "height": int(H)}],
            "annotations": annotations,
            "categories": [{"id": c, "name": _coco_name(c), "supercategory": ""} for c in cats],
        }
        fp = os.path.join(self._dirs["annotations"], f"{tag}.json")
        with open(fp, "w", encoding="utf-8") as fh:
            json.dump(doc, fh, ensure_ascii=False)
        self._stats["annots"] += 1

    def _do_meta(self, job):
        if "metadata_json" not in self._dirs:
            return
        fp = os.path.join(self._dirs["metadata_json"], f"{job['pkey']}.json")
        with open(fp, "w", encoding="utf-8") as fh:
            json.dump(job["payload"], fh, ensure_ascii=False, default=float)
        self._stats["meta"] += 1

    # -- lifecycle ---------------------------------------------------------
    def close(self):
        if not self.enabled or self._thread is None:
            return
        try:
            self._queue.join()
        except Exception:  # noqa: BLE001
            pass
        self._stop.set()
        self._thread.join(timeout=10)
        print(f"[AL] closed stats={self._stats}", flush=True)