# SPDX-License-Identifier: Apache-2.0
"""Input handling shared by the Python API, the HTTP server and the CPU reference: up to eight cameras + calibration +
ego speed + stream -> one :class:`~.preprocess.MeteorFrame`.

Accepted ``images`` (Python API): a list of ``ttaw.io.CameraImage`` (what the server decodes ``/predict`` into), a list
of dicts ``{"camera": name, "image": path | bytes | array | PIL, "intrinsics": ..., "T_ref_from_camera": ...}``, or a
mapping ``{camera name: image}``. Calibration not given inline comes from ``calibration["cameras"][name]``
(``intrinsics`` of the image AS SENT, ``T_ref_from_camera`` = camera optical frame -> base_link (METEOR's ego frame:
x forward, y left, z up, origin on the ground), any ``ttaw.io.parse_transform`` spelling) or from a preset
(``{"preset": name}`` -> ``calib/<name>.json``). The six wide / corner cameras are required; a missing
``CAM_FRONT_NARROW`` / ``CAM_BACK_NARROW`` is an absent camera (zero image + its donor's K and pose, as trained:
``config.DONOR``). ``ego_speed`` (m/s) is the graph's ``v0``. ``stream`` = ``{"id", "reset", "timestamp_s",
"T_world_from_ego"}``: the pose drives the host temporal post-processing (BEV seg fusion, yaw tracks).

:func:`load_sample` reads the bundle's sample manifests (``samples/<name>.json``: image paths relative to the file, a
calibration preset, the ego speed and an optional stream) into the keyword arguments of ``model(...)``. A preset is
looked up in the ``calib/`` directory beside the manifest's ``samples/`` directory first (a staged sample carries its
own preset), then in the package's ``calib/``; it is returned inline.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple

import numpy as np

from ..reference import config as C
from ..ttaw import io as tio
from .preprocess import MeteorFrame, assemble_frame, pose_from_transform

__all__ = ["load_cameras", "prepare_request", "load_sample", "SAMPLES_DIR", "CALIB_DIR", "stream_pose"]

PKG_DIR = Path(__file__).resolve().parents[1]
SAMPLES_DIR = PKG_DIR / "samples"
CALIB_DIR = PKG_DIR / "calib"
REQUIRED = tuple(c for c in C.CAMERAS if c not in C.DONOR)


def _camera_from(entry: Any, name: Optional[str], cams_cal: Mapping[str, Any]) -> "tio.CameraImage":
    if isinstance(entry, tio.CameraImage):
        cam = entry
    elif isinstance(entry, Mapping):
        if "camera" not in entry and name is None:
            raise tio.InputError("every image needs a 'camera' name")
        name = str(entry.get("camera", name))
        if "data" in entry:                                           # the /predict spelling (base64)
            return tio.decode_cameras([dict(entry, camera=name)], {"cameras": dict(cams_cal)}, order=None,
                                      require_calibration=True)[0]
        src = entry.get("image", entry.get("path"))
        if src is None:
            raise tio.InputError(f"camera {name!r}: no 'image' / 'path' / 'data'")
        cam = tio.CameraImage(name=name, image=tio.load_image(src), intrinsics=entry.get("intrinsics"),
                              T_ref_from_camera=entry.get("T_ref_from_camera", entry.get("extrinsics")),
                              timestamp_s=entry.get("timestamp_s"))
    else:
        if name is None:
            raise tio.InputError("an image without a camera name: pass dicts with 'camera' or a {name: image} map")
        cam = tio.CameraImage(name=str(name), image=tio.load_image(entry))
    cal = dict(cams_cal.get(cam.name) or {})
    k = cam.intrinsics if cam.intrinsics is not None else cal.get("intrinsics")
    t = cam.T_ref_from_camera if cam.T_ref_from_camera is not None else cal.get(
        "T_ref_from_camera", cal.get("extrinsics"))
    if k is None or t is None:
        raise tio.InputError(f"camera {cam.name!r} needs 'intrinsics' and 'T_ref_from_camera' (inline or in "
                             "calibration.cameras)")
    return tio.CameraImage(name=cam.name, image=cam.image,
                           intrinsics=tio.parse_intrinsics(k, field=f"{cam.name}.intrinsics"),
                           T_ref_from_camera=tio.parse_transform(t, field=f"{cam.name}.T_ref_from_camera"),
                           distortion=cam.distortion, timestamp_s=cam.timestamp_s)


def resolve_calibration(calibration: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
    calibration = dict(calibration or {})
    if "preset" in calibration:
        calibration = tio.resolve_calibration(calibration, CALIB_DIR)
    return calibration


def load_cameras(images: Any, calibration: Optional[Mapping[str, Any]]) -> Dict[str, "tio.CameraImage"]:
    """``images`` (+ ``calibration``) -> {camera name: CameraImage with K and T_ref_from_camera} (the present ones)."""
    if images is None:
        raise tio.InputError(f"this model needs 'images': the cameras {list(C.CAMERAS)}")
    cams_cal = dict(resolve_calibration(calibration).get("cameras") or {})
    if isinstance(images, Mapping):
        cams = [_camera_from(v, k, cams_cal) for k, v in images.items()]
    else:
        cams = [_camera_from(e, None, cams_cal) for e in images]
    by_name: Dict[str, Any] = {}
    for c in cams:
        if c.name in by_name:
            raise tio.InputError(f"camera {c.name!r} given twice")
        by_name[c.name] = c
    missing = [n for n in REQUIRED if n not in by_name]
    extra = [n for n in by_name if n not in C.CAMERAS]
    if missing or extra:
        raise tio.InputError(f"cameras must be {list(C.CAMERAS)} (the narrow ones may be absent); missing {missing}, "
                             f"unexpected {extra}")
    for c in by_name.values():
        a = np.asarray(c.image)
        if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8:
            raise tio.InputError(f"camera {c.name!r}: expected an (H, W, 3) uint8 image, got {a.shape} {a.dtype}")
        if a.shape[0] < C.IMG_H or a.shape[1] < C.IMG_W:
            raise tio.InputError(f"camera {c.name!r}: {a.shape[1]}x{a.shape[0]} is smaller than 768x432")
    return by_name


def stream_pose(stream: Optional[Mapping[str, Any]]) -> Optional[Tuple[float, float, float]]:
    """(x, y, yaw) from ``stream["T_world_from_ego"]`` (or ``stream["pose"]`` = [x, y, yaw]); None without one."""
    if not stream:
        return None
    if stream.get("pose") is not None:
        p = [float(v) for v in stream["pose"]]
        if len(p) != 3:
            raise tio.InputError("stream.pose must be [x, y, yaw]")
        return p[0], p[1], p[2]
    t = stream.get("T_world_from_ego")
    if t is None:
        return None
    return pose_from_transform(tio.parse_transform(t, field="stream.T_world_from_ego"))


def prepare_request(images: Any, calibration: Optional[Mapping[str, Any]], ego_speed: Any,
                    stream: Optional[Mapping[str, Any]] = None, *, inverse: str = "float32",
                    layout: str = "rows") -> MeteorFrame:
    """The API inputs -> the network input of one frame (``MeteorFrame``), with the ego pose of ``stream``."""
    if ego_speed is None:
        raise tio.InputError("this model needs 'ego_speed' (m/s, the graph's v0)")
    try:
        v0 = float(ego_speed)
    except (TypeError, ValueError):
        raise tio.InputError(f"ego_speed must be a number (m/s), got {ego_speed!r}") from None
    if not np.isfinite(v0):
        raise tio.InputError("ego_speed must be finite")
    cams = load_cameras(images, calibration)
    try:
        frame = assemble_frame({n: c.image for n, c in cams.items()}, {n: c.intrinsics for n, c in cams.items()},
                               {n: c.T_ref_from_camera for n, c in cams.items()}, v0, inverse=inverse,
                               pose=stream_pose(stream), layout=layout)
    except ValueError as e:
        raise tio.InputError(str(e)) from None
    frame.meta["cameras"] = sorted(cams)
    return frame


def load_sample(path: Any, *, with_stream: bool = True) -> Dict[str, Any]:
    """A sample manifest ``{"images": {camera: relative path}, "calibration": {...} | {"preset": name},
    "ego_speed": m/s, "stream": {...}}`` -> the keyword arguments of ``model(...)``."""
    p = Path(path)
    spec = json.loads(p.read_text())
    calib = dict(spec.get("calibration") or {})
    beside = p.parent.parent / "calib"
    if "preset" in calib and beside.resolve() != CALIB_DIR.resolve() and (beside / f"{calib['preset']}.json").is_file():
        calib = tio.resolve_calibration(calib, beside)
    calib = resolve_calibration(calib)
    images = [{"camera": name, "image": str((p.parent / rel).resolve())} for name, rel in spec["images"].items()]
    cams = [_camera_from(e, None, dict(calib.get("cameras") or {})) for e in images]
    out: Dict[str, Any] = {"images": cams, "calibration": calib, "ego_speed": float(spec["ego_speed"])}
    if with_stream and spec.get("stream"):
        out["stream"] = dict(spec["stream"])
    return out