"""Resident, pinned ADE20K surface provider. RGB only; no scene truth input."""

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

import numpy as np
import torch
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor

# These are surface candidates, not permission to drive. The metric terrain
# planner still checks step height, slope and the full rover footprint.
SURFACES = (3, 6, 9, 11, 13, 29, 34, 46, 52, 91)
FILES = {
    "config.json": "ee7400840fdb1e5045f0b2eba78bf053df8e33a309c4acec31705a48c8cc5c00",
    "preprocessor_config.json": "8039d1d210abaa7117ad78e58cdfd6141a2ec72c03dae891b3cd76737e422c6c",
    "pytorch_model.bin": "187ca07bea003a5717c63d04ea90b07f33cd033c0ebf44b4b89fce5070d6c8f3",
}


def main():
    for name, expected in FILES.items():
        if hashlib.sha256((Path("/assets") / name).read_bytes()).hexdigest() != expected:
            raise RuntimeError("SegFormer asset identity changed: " + name)
    processor = SegformerImageProcessor.from_pretrained("/assets", local_files_only=True)
    model = (
        SegformerForSemanticSegmentation.from_pretrained("/assets", local_files_only=True)
        .cuda()
        .eval()
    )
    torch.set_num_threads(2)

    def infer(rgb):
        image = Image.fromarray(rgb).crop((100, 0, 700, 600))
        tensor = processor(images=image, return_tensors="pt")["pixel_values"].cuda()
        with torch.inference_mode():
            logits = model(tensor).logits
            logits = torch.nn.functional.interpolate(
                logits, size=(512, 512), mode="bilinear", align_corners=False
            )
            confidence, labels = logits.softmax(1).max(1)
        labels = labels[0].cpu().numpy().astype(np.uint8)
        candidate = np.isin(labels, SURFACES) & (confidence[0].cpu().numpy() >= 0.55)
        return labels.tobytes() + candidate.astype(np.uint8).tobytes()

    infer(np.zeros((600, 800, 3), dtype=np.uint8))

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

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

        def do_GET(self):
            self.reply(
                200 if self.path == "/ready" else 404,
                json.dumps({"model": "segformer-b2-ade150"}).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")
                self.reply(200, infer(np.frombuffer(raw, np.uint8).reshape(600, 800, 3)))
            except (TimeoutError, ValueError, RuntimeError):
                self.reply(500, b"Surface inference failed")

        def log_message(self, *_):
            pass

    HTTPServer(("0.0.0.0", 8010), Handler).serve_forever()


if __name__ == "__main__":
    main()