"""Resident reference DDRNet adapter; run only inside the pinned module image.

Raw pinhole RGB is the only input. No device mask, scene labels or actor truth.
The container publishes this port on Windows loopback only and owns no weights.
"""

import hashlib
import importlib.util
import json
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

import numpy as np
from PIL import Image


def checked(path, expected):
    value = Path(path)
    if hashlib.sha256(value.read_bytes()).hexdigest() != expected:
        raise RuntimeError("Pinned DDRNet asset changed: " + value.name)
    return value


def main():
    checkpoint = checked(
        "/assets/ddrnet-checkpoint.pth",
        "b99c2838051bcd7b092fd3970aa62a77d5c0bbb809c9b9afb2ff4b0ebdaa4ee6",
    )
    runner = checked(
        "/assets/ddrnet-goose-runner.py",
        "b18ad60f277eea69a240a28f290611b94627fb9707faf1bb3e6e22102dad67c1",
    )
    spec = importlib.util.spec_from_file_location("polygon_pinned_goose", runner)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    model, _, _ = module.load_model("ddrnet", checkpoint)
    tensor, _ = module.preprocess(Image.fromarray(np.zeros((600, 800, 3), np.uint8)))
    module.infer(model, tensor)

    class Handler(BaseHTTPRequestHandler):
        def setup(self):
            super().setup()
            self.connection.settimeout(10)

        def reply(self, status, body, kind="application/octet-stream"):
            self.send_response(status)
            self.send_header("Content-Type", kind)
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path != "/ready":
                self.reply(404, b"")
                return
            self.reply(200, json.dumps({"model": "ddrnet-goose-pytorch-reference"}).encode())

        def do_POST(self):
            if self.path != "/infer" or self.headers.get("Content-Length") != "1440000":
                self.reply(400, b"Expected 800x600 raw RGB uint8")
                return
            try:
                raw = self.rfile.read(1440000)
                if len(raw) != 1440000:
                    raise ValueError("Incomplete camera frame")
                rgb = np.frombuffer(raw, np.uint8).reshape(600, 800, 3)
                tensor, _ = module.preprocess(Image.fromarray(rgb))
                mask, _ = module.infer(model, tensor)
                if mask.shape != (512, 512) or np.any(mask < 0) or np.any(mask >= 64):
                    raise ValueError("DDRNet output contract changed")
                self.reply(200, mask.astype(np.uint8).tobytes())
            except (TimeoutError, ValueError, RuntimeError):
                self.reply(500, b"DDRNet inference failed")

        def log_message(self, *_):
            pass

    # External exposure is restricted by the Compose loopback publication.
    HTTPServer(("0.0.0.0", 8010), Handler).serve_forever()


if __name__ == "__main__":
    main()