"""ImagePipeline (T05): independent image block, delegation, read-only results.

The image block is verified three ways: the same prepared frames through the
published Sensor path and through a standalone ImagePipeline must produce
bit-identical flow/HHD/force/field results; measured arrays must be read-only
copies that later frames cannot overwrite; and the optical-flow reference
generation must be written when the image thread computes a frame (a reset
requested while a frame is in flight applies to the next frame only).
"""

from __future__ import annotations

import json
import threading
import time

import numpy as np
import pytest

from orisys import Sensor, configs
from orisys.image_pipeline import (
    ImageMeasurement,
    ImagePipeline,
    ImageSettings,
    sensor_clock,
)
from orisys.sensors import build_sensor_context
from orisys.stream_contract import FrameIdentity, MeasurementFrame


def _frames(count: int = 4, size: int = 365) -> list[np.ndarray]:
    """Deterministic grey frames whose pattern translates a few pixels."""
    height = width = size
    yy, xx = np.mgrid[0:height, 0:width]
    frames = []
    for index in range(count):
        cy = 130.0 + 2.0 * index
        cx = 210.0 + 3.0 * index
        blob = 120.0 * np.exp(-(((xx - cx) / 25.0) ** 2 + ((yy - cy) / 21.0) ** 2))
        gradient = 40.0 + 30.0 * (xx / width) + 20.0 * (yy / height)
        frames.append(np.clip(gradient + blob, 0.0, 255.0).astype(np.uint8))
    return frames


def _config(name: str = "box") -> dict:
    return configs.load_config(name)


def _context():
    cfg = _config()
    return build_sensor_context(
        config_name="box",
        config=cfg,
        device_id="BOX-001",
        bundle_candidates=["box"],
    )


def _identity(context, frame_id: int = 0) -> FrameIdentity:
    return context.frame_identity(
        session_id="run-1", frame_id=frame_id, source_frame_id=frame_id
    )


# --------------------------------------------------------------- same-run parity
def test_standalone_pipeline_matches_the_published_sensor_path() -> None:
    """Old (Sensor delegation) and new (standalone) paths, same frames, same run."""
    sensor = Sensor(
        "external", config_name="box", verbose=False, backend="cpu", cuda=False, mirror=False
    )
    try:
        settings = ImageSettings.from_config(
            _config(),
            cuda=False,
            backend="cpu",
            flow_ema_alpha=sensor._force_readout.flow_ema_alpha,
        )
        assert settings.as_dict() == sensor.image_settings.as_dict()
        pipeline = ImagePipeline(settings)
        context = _context()

        for index, raw in enumerate(_frames()):
            prepared = sensor.get_img(
                frame=raw, already_processed=True, already_gray=True
            )
            assert np.array_equal(prepared, raw), "same-size resize must be identity"

            sensor.compute_flow()
            sensor.compute_hhd()
            sensor_flow = np.array(sensor.flow, copy=True)
            sensor_vnormal = np.array(sensor.vnormal, copy=True)
            sensor_vshear = np.array(sensor.vshear, copy=True)
            sensor_field = sensor.current_field()
            sensor_force = tuple(
                sensor.read_info(
                    sensor.info.FNORMAL_RAW,
                    sensor.info.FSHEARX_RAW,
                    sensor.info.FSHEARY_RAW,
                )
            )

            measurement = pipeline.measure(
                raw,
                identity=_identity(context, index),
                clock=sensor_clock(1_000_000 + index, clock_id="host"),
                include=("flow", "vnormal", "vshear", "field", "valid_mask"),
            )
            assert np.array_equal(measurement.flow, sensor_flow), index
            assert np.array_equal(measurement.vnormal, sensor_vnormal), index
            assert np.array_equal(measurement.vshear, sensor_vshear), index
            assert measurement.force_raw == pytest.approx(sensor_force), index
            for key, value in sensor_field.items():
                if isinstance(value, np.ndarray):
                    assert np.array_equal(measurement.field[key], value), (index, key)
                else:
                    assert measurement.field[key] == value, (index, key)
    finally:
        sensor.disconnect()


def test_standalone_pipeline_runs_without_constructing_a_sensor() -> None:
    context = _context()
    pipeline = ImagePipeline(
        ImageSettings.from_config(_config(), cuda=False, backend="cpu"), context=context
    )
    measurement = pipeline.measure(
        _frames()[0],
        identity=_identity(context),
        clock=sensor_clock(123, clock_id="host"),
        source_dt_s=None,
        include=("flow",),
    )
    assert isinstance(measurement, ImageMeasurement)
    assert isinstance(measurement.frame, MeasurementFrame)
    assert measurement.frame.stamp.sensor_id == "BOX-001"
    assert measurement.frame.stamp.session_id == "run-1"
    assert measurement.frame.reference_generation == 0
    assert measurement.frame.source_dt_s is None
    # The T03 header is the serialised contract: it round-trips unchanged.
    assert (
        MeasurementFrame.from_dict(json.loads(json.dumps(measurement.frame.to_dict())))
        == measurement.frame
    )
    assert pipeline.context is context


# ---------------------------------------------------------------- read-only data
def test_measured_arrays_are_readonly_copies_that_later_frames_do_not_rewrite() -> None:
    context = _context()
    pipeline = ImagePipeline(
        ImageSettings.from_config(_config(), cuda=False, backend="cpu"), context=context
    )
    frames = _frames()
    first = pipeline.measure(
        frames[0],
        identity=_identity(context, 0),
        clock=sensor_clock(10, clock_id="host"),
        include=("flow", "vnormal", "vshear", "field", "valid_mask"),
    )
    seen = {
        "flow": np.array(first.flow, copy=True),
        "vnormal": np.array(first.vnormal, copy=True),
        "field": np.array(first.field["d"], copy=True),
    }
    for name, array in (
        ("flow", first.flow),
        ("vnormal", first.vnormal),
        ("field", first.field["d"]),
    ):
        assert array.flags.writeable is False, name
        with pytest.raises(ValueError):
            array[0, 0] = 1.0

    # Computing the next frame must not overwrite what was already published.
    second = pipeline.measure(
        frames[1],
        identity=_identity(context, 1),
        clock=sensor_clock(20, clock_id="host"),
        include=("flow", "field"),
    )
    assert np.array_equal(first.flow, seen["flow"])
    assert np.array_equal(first.vnormal, seen["vnormal"])
    assert np.array_equal(first.field["d"], seen["field"])
    assert not np.array_equal(second.flow, first.flow)


def test_include_selects_only_requested_arrays() -> None:
    context = _context()
    pipeline = ImagePipeline(
        ImageSettings.from_config(_config(), cuda=False, backend="cpu"), context=context
    )
    measurement = pipeline.measure(
        _frames()[0],
        identity=_identity(context),
        clock=sensor_clock(1, clock_id="host"),
        include=("flow",),
    )
    assert measurement.flow is not None
    assert measurement.vnormal is None and measurement.vshear is None
    assert measurement.field is None and measurement.valid_mask is None
    with pytest.raises(ValueError, match="unknown include"):
        pipeline.measure(
            _frames()[0],
            identity=_identity(context),
            clock=sensor_clock(1, clock_id="host"),
            include=("not-a-field",),
        )


# ------------------------------------------------------ reference-generation seam
class _GatedPipeline(ImagePipeline):
    """Barrier-style gate: hold a frame after its reference version was captured."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.entered = threading.Event()
        self.release = threading.Event()
        self.gate_once = True

    def _compute_flow_state(self, image, *, valid_mask=None):
        if self.gate_once:
            self.gate_once = False
            self.entered.set()
            assert self.release.wait(timeout=5.0)
        return super()._compute_flow_state(image, valid_mask=valid_mask)


def test_inflight_frame_keeps_its_own_reference_version() -> None:
    """The stream can adopt a new version while an older frame is still in flight."""
    context = _context()
    pipeline = _GatedPipeline(
        ImageSettings.from_config(_config(), cuda=False, backend="cpu"), context=context
    )
    frames = _frames()
    results: list[ImageMeasurement] = []

    def compute_first() -> None:
        results.append(
            pipeline.measure(
                frames[0],
                identity=_identity(context, 0),
                clock=sensor_clock(100, clock_id="host"),
                reference_generation=0,
            )
        )

    worker = threading.Thread(target=compute_first)
    worker.start()
    try:
        assert pipeline.entered.wait(timeout=5.0), "first frame never reached the gate"
        # The stream adopts version 1 while frame 0 is provably in flight.
        pipeline.set_reference_generation(1)
        assert pipeline.reference_generation == 1
    finally:
        pipeline.release.set()
        worker.join(timeout=5.0)
    assert not worker.is_alive()
    in_flight = results[0]
    assert in_flight.frame.reference_generation == 0, "old frame was relabelled"

    next_measurement = pipeline.measure(
        frames[1],
        identity=_identity(context, 1),
        clock=sensor_clock(200, clock_id="host"),
        reference_generation=1,
    )
    assert next_measurement.frame.reference_generation == 1
    assert in_flight.frame.reference_generation == 0


def test_reference_version_is_adopted_not_counted_by_the_image_block() -> None:
    """:class:`ImagePipeline` stamps the stream's version; a reset does not bump it."""
    context = _context()
    pipeline = ImagePipeline(
        ImageSettings.from_config(_config(), cuda=False, backend="cpu"), context=context
    )
    frames = _frames()
    first = pipeline.measure(
        frames[0],
        identity=_identity(context, 0),
        clock=sensor_clock(1, clock_id="host"),
        reference_generation=0,
    )
    assert first.frame.reference_generation == 0
    pipeline.reset_image_state()
    zeroed = pipeline.measure(
        frames[1],
        identity=_identity(context, 1),
        clock=sensor_clock(2, clock_id="host"),
        reference_generation=0,
    )
    assert zeroed.frame.reference_generation == 0, "reset must not invent a version"
    pipeline.set_reference_generation(3)
    stamped = pipeline.measure(
        frames[2],
        identity=_identity(context, 2),
        clock=sensor_clock(3, clock_id="host"),
        reference_generation=3,
    )
    assert stamped.frame.reference_generation == 3
    assert pipeline.reference_generation == 3


def _drive_pipeline(sensor: Sensor, frames: list[np.ndarray], *, request_reset: bool):
    """Run one StreamPipeline; record per-flow-frame adopted versions and snapshots."""
    from orisys.pipeline import StreamPipeline

    flow_versions: list[int] = []
    snapshots: list = []
    original = sensor.compute_flow

    def recording_flow(*args, **kwargs):
        flow_versions.append(int(sensor.image_pipeline.reference_generation))
        return original(*args, **kwargs)

    sensor.compute_flow = recording_flow
    queue = iter(frames)

    def produce():
        try:
            raw = next(queue)
        except StopIteration:
            return None
        return sensor.get_img(frame=raw, already_processed=True, already_gray=True)

    pipeline = StreamPipeline(
        sensor, produce, include=(), on_frame=snapshots.append, verbose=False
    )
    if request_reset:
        pipeline.request_reset()
    pipeline.start()
    try:
        deadline = time.perf_counter() + 10.0
        while pipeline.running and time.perf_counter() < deadline:
            time.sleep(0.01)
        assert pipeline.error is None
    finally:
        pipeline.stop()
        sensor.compute_flow = original
    snapshot_versions = [int(snapshot.tracker_reset_seq) for snapshot in snapshots]
    return flow_versions, snapshot_versions


def test_pipeline_driven_reference_versions_match_every_snapshot_on_a_reused_sensor() -> None:
    """One version source: StreamPipeline's claimed version reaches the image block."""
    sensor = Sensor(
        "external", config_name="box", verbose=False, backend="cpu", cuda=False, mirror=False
    )
    frames = _frames()
    try:
        sensor.reset()
        flow_versions, snapshot_versions = _drive_pipeline(sensor, frames, request_reset=False)
        assert flow_versions == snapshot_versions == [0, 0, 0, 0]

        flow_versions, snapshot_versions = _drive_pipeline(sensor, frames, request_reset=True)
        assert flow_versions == snapshot_versions == [1, 1, 1, 1]

        # A direct Sensor.reset() between two runs must not leave the image block
        # on an old private counter: the next stream's version 0 wins.
        sensor.reset()
        flow_versions, snapshot_versions = _drive_pipeline(sensor, frames, request_reset=False)
        assert flow_versions == snapshot_versions == [0, 0, 0, 0]
        assert sensor.image_pipeline.reference_generation == 0
    finally:
        sensor.disconnect()


# ------------------------------------------------------ published flow-EMA binding
def test_runtime_flow_ema_configuration_keeps_head_behaviour() -> None:
    """``Sensor.configure(flow_ema_alpha=...)`` must affect the next frame (HEAD)."""
    from orisys.force_readout import ForceReadout

    context = _context()
    sensor = Sensor(
        "external", config_name="box", verbose=False, backend="cpu", cuda=False, mirror=False
    )
    # HEAD-equivalent reference: the published Tracker/ForceReadout composition.
    reference_readout = ForceReadout({"flow_ema_alpha": 1.0})
    reference = ImagePipeline(
        ImageSettings.from_config(
            _config(), cuda=False, backend="cpu", flow_ema_alpha=1.0
        ),
        flow_ema=reference_readout,
    )
    try:
        for index, raw in enumerate(_frames()):
            if index == 2:
                sensor.configure(flow_ema_alpha=0.5)
                reference_readout.configure(flow_ema_alpha=0.5)
            prepared = sensor.get_img(
                frame=raw, already_processed=True, already_gray=True
            )
            sensor.compute_flow()
            sensor.compute_hhd()
            measurement = reference.measure(
                prepared,
                identity=_identity(context, index),
                clock=sensor_clock(1000 + index, clock_id="host"),
                include=("flow", "vnormal"),
                reference_generation=0,
            )
            assert sensor.image_pipeline.flow_ema_alpha == 0.5 if index >= 2 else True
            assert np.array_equal(np.array(sensor.flow, copy=True), measurement.flow), index
            assert np.array_equal(
                np.array(sensor.vnormal, copy=True), measurement.vnormal
            ), index
            sensor_force = tuple(
                sensor.read_info(
                    sensor.info.FNORMAL_RAW,
                    sensor.info.FSHEARX_RAW,
                    sensor.info.FSHEARY_RAW,
                )
            )
            assert sensor_force == measurement.force_raw, index
    finally:
        sensor.disconnect()


def test_static_flow_ema_config_is_bound_once(tmp_path) -> None:
    """A config-declared alpha<1 keeps the same flow through Sensor and a standalone pipeline."""
    config = _config()
    config["filters"] = dict(config.get("filters") or {}, flow_ema_alpha=0.5)
    config_path = tmp_path / "box-ema05.json"
    config_path.write_text(json.dumps(config), encoding="utf-8")

    context = _context()
    pipeline = ImagePipeline(
        ImageSettings.from_config(config, cuda=False, backend="cpu")
    )
    sensor = Sensor(
        "external",
        config_name=str(config_path),
        verbose=False,
        backend="cpu",
        cuda=False,
        mirror=False,
    )
    try:
        assert pipeline.settings.flow_ema_alpha == 0.5
        assert sensor.image_pipeline.flow_ema_alpha == 0.5
        for index, raw in enumerate(_frames()):
            prepared = sensor.get_img(
                frame=raw, already_processed=True, already_gray=True
            )
            sensor.compute_flow()
            sensor.compute_hhd()
            measurement = pipeline.measure(
                prepared,
                identity=_identity(context, index),
                clock=sensor_clock(index, clock_id="host"),
                include=("flow",),
                reference_generation=0,
            )
            assert np.array_equal(np.array(sensor.flow, copy=True), measurement.flow), index
    finally:
        sensor.disconnect()


def test_flow_ema_matches_the_force_readout_recurrence() -> None:
    from orisys.force_readout import ForceReadout

    settings = ImageSettings.from_config(
        _config(), cuda=False, backend="cpu", flow_ema_alpha=0.5
    )
    pipeline = ImagePipeline(settings)
    readout = ForceReadout({"flow_ema_alpha": 0.5})
    rng = np.random.default_rng(7)
    for _ in range(4):
        flow = rng.normal(0.0, 1.0, (6, 5, 2)).astype(np.float32)
        pipeline_flow = pipeline._apply_flow_ema(flow.copy())
        readout_flow = readout.apply_flow_ema(flow.copy())
        assert np.array_equal(pipeline_flow, readout_flow)
    pipeline.reset_image_state()
    readout.reset_flow_ema()
    flow = rng.normal(0.0, 1.0, (6, 5, 2)).astype(np.float32)
    assert np.array_equal(
        pipeline._apply_flow_ema(flow.copy()), readout.apply_flow_ema(flow.copy())
    )


# -------------------------------------------------------------- delegation paths
def test_sensor_delegates_its_image_block_to_one_pipeline() -> None:
    sensor = Sensor(
        "external", config_name="box", verbose=False, backend="cpu", cuda=False, mirror=False
    )
    try:
        assert isinstance(sensor.image_pipeline, ImagePipeline)
        assert (
            sensor.image_pipeline.settings.as_dict() == sensor.image_settings.as_dict()
        )
        sensor.get_img(frame=_frames()[0], already_processed=True, already_gray=True)
        sensor.compute_flow()
        assert sensor.flow is sensor.image_pipeline.flow
        sensor.compute_hhd()
        assert sensor.vnormal is sensor.image_pipeline.vnormal
        assert sensor.vshear is sensor.image_pipeline.vshear
        before = sensor.image_pipeline.reference_generation
        sensor.reset_image_state()
        # The reference version belongs to the stream; a local reset must not
        # invent one. Only an adopted stream version changes it.
        assert sensor.image_pipeline.reference_generation == before
        sensor.set_stream_reference_generation(7)
        assert sensor.image_pipeline.reference_generation == 7
        # The published Newton/readout accessors keep working.
        assert sensor.force_unit in ("N", "raw")
    finally:
        sensor.disconnect()


def test_finger_runtime_exposes_the_same_image_pipeline() -> None:
    from orisys.finger.runtime import FingerRuntime

    pipeline = object()
    runtime = FingerRuntime.__new__(FingerRuntime)
    runtime.tactile = type(
        "Tactile", (), {"raw_sensor": type("Raw", (), {"image_pipeline": pipeline})()}
    )()
    assert runtime.image_pipeline is pipeline

    runtime.tactile = type("Tactile", (), {"raw_sensor": object()})()
    assert runtime.image_pipeline is None