#!/usr/bin/env python3
"""Lightweight HTTP worker that keeps YOLOv8 segmentation warm and serves occluder masks."""

import base64
import io
import json
import os
import sys
import threading
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from cgi import FieldStorage

import numpy as np
from PIL import Image, ImageFilter

try:
    from ultralytics import YOLO
except Exception as exc:  # pragma: no cover
    print(f"[worker] failed to import ultralytics: {exc}", file=sys.stderr)
    raise

MODEL_PATH = os.environ.get("OCCLUDER_SEG_MODEL", "models/yolov8x-seg.pt")
CONF_THRESHOLD = float(os.environ.get("OCCLUDER_SEG_CONF", "0.3"))
IOU_THRESHOLD = float(os.environ.get("OCCLUDER_SEG_IOU", "0.5"))
MEDIAN = int(os.environ.get("OCCLUDER_SEG_MEDIAN", "3"))
DILATE = int(os.environ.get("OCCLUDER_SEG_DILATE", "3"))
MAX_MASKS = int(os.environ.get("OCCLUDER_SEG_MAX", "40"))
EXCLUDED_CLASSES = {
    name.strip().lower()
    for name in os.environ.get("OCCLUDER_SEG_EXCLUDE", "").split(",")
    if name.strip()
}

_WORKER_LOCK = threading.Lock()
_MODEL = YOLO(MODEL_PATH)

def _clamp_filter_size(value: int) -> int:
    value = max(0, int(value))
    if value <= 1:
        return 0
    return value if value % 2 == 1 else value + 1

def _infer_mask(image_bytes: bytes):
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    np_image = np.array(image)
    results = _MODEL.predict(
        source=np_image,
        conf=CONF_THRESHOLD,
        iou=IOU_THRESHOLD,
        verbose=False,
    )
    if not results:
        raise RuntimeError("YOLO returned no predictions")

    result = results[0]
    if result.masks is None or result.masks.data is None:
        raise RuntimeError("YOLO segmentation masks missing")

    masks = result.masks.data
    try:
        masks_np = masks.cpu().numpy()
    except AttributeError:
        masks_np = np.asarray(masks)

    boxes = result.boxes
    if boxes is None or boxes.cls is None:
        raise RuntimeError("YOLO bounding boxes missing")

    try:
        cls_ids = boxes.cls.cpu().numpy().astype(int)
        confidences = boxes.conf.cpu().numpy().astype(float)
    except AttributeError:
        cls_ids = np.asarray(boxes.cls).astype(int)
        confidences = np.asarray(boxes.conf).astype(float)

    names_map = getattr(result, "names", {}) or {}
    names = {int(k): str(v) for k, v in names_map.items()}

    height, width = masks_np.shape[1], masks_np.shape[2]
    threshold = float(os.environ.get("OCCLUDER_SEG_THRESHOLD", "0.2"))
    min_area = float(os.environ.get("OCCLUDER_SEG_MIN_AREA", "1000"))

    selected = []
    selected_meta = []

    for idx, class_id in enumerate(cls_ids):
        name = names.get(int(class_id), str(int(class_id))).lower()
        if name in EXCLUDED_CLASSES:
            continue
        score = float(confidences[idx])
        mask = masks_np[idx]
        solid = float((mask > threshold).sum())
        if solid < min_area:
            continue
        selected.append(idx)
        selected_meta.append(
            {
                "index": int(idx),
                "class": names.get(int(class_id), str(int(class_id))),
                "class_id": int(class_id),
                "score": score,
                "area_px": int(solid),
            }
        )

    if not selected:
        raise RuntimeError("No foreground instances survived filtering")

    selected = selected[: max(1, MAX_MASKS)]
    union = np.zeros((height, width), dtype=np.uint8)
    for idx in selected:
        union |= (masks_np[idx] > threshold).astype(np.uint8)

    mask_img = Image.fromarray(union * 255, mode="L")
    median_size = _clamp_filter_size(MEDIAN)
    if median_size:
        mask_img = mask_img.filter(ImageFilter.MedianFilter(size=median_size))
    dilate_size = _clamp_filter_size(DILATE)
    if dilate_size:
        mask_img = mask_img.filter(ImageFilter.MaxFilter(size=dilate_size))

    mask_arr = np.array(mask_img)
    coverage = float(mask_arr.sum() / 255) / float(mask_arr.size)
    mask_bytes = io.BytesIO()
    mask_img.save(mask_bytes, format="PNG")
    mask_b64 = base64.b64encode(mask_bytes.getvalue()).decode("ascii")

    meta = {
        "ok": True,
        "model": os.path.basename(MODEL_PATH),
        "conf": CONF_THRESHOLD,
        "iou": IOU_THRESHOLD,
        "instances": selected_meta,
    }

    return {
        "mask": mask_b64,
        "width": int(width),
        "height": int(height),
        "coverage": coverage,
        "source": "yolo-worker",
        "meta": meta,
    }


class OccluderHandler(BaseHTTPRequestHandler):
    server_version = "OccluderWorker/1.0"

    def _json_response(self, status: HTTPStatus, payload: dict) -> None:
        data = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def do_GET(self):  # noqa: N802
        if self.path == "/health":
            self._json_response(HTTPStatus.OK, {"ok": True})
        else:
            self._json_response(HTTPStatus.NOT_FOUND, {"ok": False, "error": "not_found"})

    def do_POST(self):  # noqa: N802
        if self.path != "/mask":
            self._json_response(HTTPStatus.NOT_FOUND, {"ok": False, "error": "not_found"})
            return
        content_length = int(self.headers.get("Content-Length", "0"))
        if content_length <= 0:
            self._json_response(HTTPStatus.BAD_REQUEST, {"ok": False, "error": "missing body"})
            return
        raw = self.rfile.read(content_length)
        environ = {
            "REQUEST_METHOD": "POST",
            "CONTENT_TYPE": self.headers.get("Content-Type", ""),
            "CONTENT_LENGTH": str(content_length),
        }
        fs = FieldStorage(fp=io.BytesIO(raw), headers=self.headers, environ=environ)
        if "file" not in fs:
            self._json_response(HTTPStatus.BAD_REQUEST, {"ok": False, "error": "missing file"})
            return
        file_item = fs["file"]
        image_bytes = file_item.file.read()
        try:
            with _WORKER_LOCK:
                result = _infer_mask(image_bytes)
            self._json_response(HTTPStatus.OK, {"ok": True, **result})
        except Exception as exc:  # pragma: no cover
            self._json_response(
                HTTPStatus.INTERNAL_SERVER_ERROR,
                {"ok": False, "error": str(exc)},
            )

    def log_message(self, format, *args):  # noqa: A003
        if os.environ.get("OCCLUDER_SEG_VERBOSE") == "1":
            super().log_message(format, *args)


def main():
    host = os.environ.get("OCCLUDER_WORKER_HOST", "127.0.0.1")
    port = int(os.environ.get("OCCLUDER_WORKER_PORT", "5055"))
    server = ThreadingHTTPServer((host, port), OccluderHandler)
    print(f"[worker] listening on http://{host}:{port}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()