"""Tests for StreamPipeline pause/resume.

A fake sensor stands in for the real one: the pause contract is about which
stages run, not about camera hardware, and a fake makes that observable without
a device. The sensor exposes the same surface the stages use - ``get_img``,
``compute_flow``, ``compute_hhd``, ``compute_contact``, ``reset``, ``read_info``
and ``_force_raw`` - so the pipeline runs its real loops unchanged.
"""

from __future__ import annotations

import threading
import time

import numpy as np
import pytest

import orisys
from orisys import PauseLevel, StreamPipeline


class FakeInfo:
    FNORMAL = "fnormal"
    FNORMAL_RAW = "fnormal_raw"
    FSHEARX = "fshearx"
    FSHEARY = "fsheary"
    VNORMAL = "vnormal"
    VSHEAR = "vshear"
    DEPTH = "depth"
    CENTROID = "centroid"
    CONTOUR = "contour"


class FakeSensor:
    """Minimal sensor surface, with counters for what each stage did."""

    def __init__(self, *, source=None) -> None:
        self.info = FakeInfo()
        self.force_unit = "raw"
        self.fnormal = 1.0
        self.fshearx = 2.0
        self.fsheary = 3.0
        self.ftor = 0.0
        self._force_raw = (10.0, 20.0, 30.0)
        self.img_size_raw = (8, 8)
        self.flow = None
        self.img = None
        self._force_filtered_raw = (10.0, 20.0, 30.0)
        self._force_raw_unfiltered = (10.0, 20.0, 30.0)
        self._force_unfiltered = (1.0, 2.0, 3.0)
        self._lock = threading.Lock()
        self.produced = 0
        self.flow_calls = 0
        self.hhd_calls = 0
        self.reset_calls = 0
        if source is not None:
            self._capture_source = source

    # ---- producer stage
    def get_img(self):
        with self._lock:
            self.produced += 1
        # A distinct frame each time, so the stages can be told apart.
        return np.full((4, 4), float(self.produced), dtype=np.float32)

    # ---- flow stage
    def compute_flow(self, frame=None):
        with self._lock:
            self.flow_calls += 1
        shape = (4, 4, 2)
        self.flow = np.zeros(shape, dtype=np.float32)

    # ---- hhd stage
    def compute_hhd(self, flow=None):
        with self._lock:
            self.hhd_calls += 1

    def compute_contact(self):
        return None

    def reset(self):
        with self._lock:
            self.reset_calls += 1

    def read_info(self, *which):
        values = []
        for key in which:
            if key in (FakeInfo.VNORMAL, FakeInfo.VSHEAR):
                values.append(np.zeros((4, 4, 2), dtype=np.float32))
            elif key == FakeInfo.DEPTH:
                values.append(np.zeros((4, 4), dtype=np.float32))
            elif key == FakeInfo.CENTROID:
                values.append(np.zeros(2, dtype=np.float32))
            elif key == FakeInfo.CONTOUR:
                values.append(np.zeros(0, dtype=np.float32))
            else:
                values.append(1.0)
        return tuple(values)


class FakeCaptureSource:
    """Records whether the camera was told to keep streaming."""

    def __init__(self) -> None:
        self.started = 0
        self.stopped = 0
        self.reads = 0
        self._running = True

    def is_open(self):
        return True

    def start_background(self):
        self.started += 1
        self._running = True

    def stop_background(self, timeout: float = 1.0):
        self.stopped += 1
        self._running = False

    def read_latest(self, *, copy: bool = True, wait_s: float = 0.0):
        self.reads += 1
        return True, np.zeros((4, 4), dtype=np.float32)


def _wait_for(predicate, timeout=2.0, interval=0.01):
    deadline = time.perf_counter() + timeout
    while time.perf_counter() < deadline:
        if predicate():
            return True
        time.sleep(interval)
    return predicate()


def test_pause_level_is_exported_from_the_package():
    assert orisys.PauseLevel is PauseLevel
    assert PauseLevel.CAPTURE_ONLY.value == "capture_only"
    assert PauseLevel.BOTH.value == "both"


def test_capture_only_keeps_reading_but_stops_computing():
    source = FakeCaptureSource()
    sensor = FakeSensor(source=source)
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 3)
        assert pipe.paused is False

        pipe.pause(PauseLevel.CAPTURE_ONLY)
        assert pipe.paused is True
        assert pipe.pause_level is PauseLevel.CAPTURE_ONLY

        flow_at_pause = sensor.flow_calls
        hhd_at_pause = sensor.hhd_calls
        produced_at_pause = sensor.produced
        time.sleep(0.4)

        # The read keeps going, so a live camera stays converged...
        assert sensor.produced > produced_at_pause, "capture-only must keep reading"
        assert source.stopped == 0, "the camera must not be stopped"
        # ...and the work stops: this is the part the hold is bought for.
        assert sensor.flow_calls == flow_at_pause, "flow must not run while held"
        assert sensor.hhd_calls == hhd_at_pause, "no HHD work may happen while held"

        published_at_pause = pipe.stats.published
        pipe.resume()
        assert pipe.paused is False
        assert _wait_for(lambda: pipe.stats.published > published_at_pause), (
            "resume must publish frames again"
        )
        assert sensor.hhd_calls > hhd_at_pause


def test_pause_both_stops_capture_as_well_as_compute():
    source = FakeCaptureSource()
    sensor = FakeSensor(source=source)
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 3)

        pipe.pause(PauseLevel.BOTH)
        assert _wait_for(lambda: source.stopped >= 1), "capture must be stopped"
        produced_at_pause = sensor.produced
        hhd_at_pause = sensor.hhd_calls
        published_at_pause = pipe.stats.published
        time.sleep(0.4)

        assert sensor.produced == produced_at_pause, "no capture while fully held"
        assert sensor.hhd_calls == hhd_at_pause
        assert pipe.stats.published == published_at_pause

        pipe.resume()
        assert _wait_for(lambda: source.started >= 1), "capture must restart"
        assert _wait_for(lambda: pipe.stats.published > published_at_pause)


def test_resume_rebases_the_tracker_by_default():
    sensor = FakeSensor()
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 2)
        resets_before = sensor.reset_calls
        pipe.pause(PauseLevel.CAPTURE_ONLY)
        time.sleep(0.1)
        pipe.resume()
        assert _wait_for(lambda: sensor.reset_calls > resets_before), (
            "a held reference frame would otherwise read as deformation"
        )


def test_resume_can_keep_the_reference_frame():
    sensor = FakeSensor()
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 2)
        resets_before = sensor.reset_calls
        pipe.pause(PauseLevel.CAPTURE_ONLY)
        time.sleep(0.1)
        pipe.resume(reset=False)
        assert _wait_for(lambda: pipe.stats.published > 0)
        time.sleep(0.2)
        assert sensor.reset_calls == resets_before


def test_wait_returns_none_while_paused_and_frames_after_resume():
    sensor = FakeSensor()
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert pipe.wait(timeout=2.0) is not None
        pipe.pause(PauseLevel.CAPTURE_ONLY)
        assert pipe.wait(timeout=0.1) is None, "a hold publishes nothing"
        # A paused pipeline is not a finished pipeline: None here must not be
        # read as end of input.
        assert pipe.running is True
        pipe.resume()
        assert pipe.wait(timeout=2.0) is not None


def test_pause_is_idempotent_and_level_changes_apply():
    source = FakeCaptureSource()
    sensor = FakeSensor(source=source)
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 2)
        pipe.pause(PauseLevel.CAPTURE_ONLY)
        pipe.pause(PauseLevel.CAPTURE_ONLY)
        assert pipe.paused is True
        pipe.pause(PauseLevel.BOTH)
        assert pipe.pause_level is PauseLevel.BOTH
        assert _wait_for(lambda: source.stopped >= 1)
        pipe.resume()
        pipe.resume()  # second resume is a no-op
        assert pipe.paused is False


def test_pause_accepts_a_plain_string_level():
    sensor = FakeSensor()
    with StreamPipeline(sensor, sensor.get_img, include=("image",)) as pipe:
        assert _wait_for(lambda: pipe.stats.published >= 1)
        pipe.pause("capture_only")
        assert pipe.pause_level is PauseLevel.CAPTURE_ONLY
        pipe.resume()


def test_stop_while_paused_terminates_promptly():
    sensor = FakeSensor()
    pipe = StreamPipeline(sensor, sensor.get_img, include=("image",))
    pipe.start()
    try:
        assert _wait_for(lambda: pipe.stats.published >= 2)
        pipe.pause(PauseLevel.BOTH)
        time.sleep(0.1)
        started = time.perf_counter()
        pipe.stop(timeout=3.0)
        assert time.perf_counter() - started < 2.0, "stop must not wait out the pause"
        assert pipe.running is False
    finally:
        if pipe.running:
            pipe.stop()


def test_end_of_input_is_still_reported_as_none():
    """The pause path must not swallow the end-of-stream sentinel.

    Counted through ``on_frame``, which delivers every frame exactly once: with
    a fast producer and a 5-frame source, ``wait`` is documented latest-wins and
    legitimately skips frames, so it cannot be used to count them.
    """
    frames = {"produced": 0}
    delivered: list[int] = []

    def produce():
        frames["produced"] += 1
        if frames["produced"] > 5:
            return None
        return np.zeros((4, 4), dtype=np.float32)

    sensor = FakeSensor()
    with StreamPipeline(
        sensor, produce, include=("image",), on_frame=lambda snap: delivered.append(snap.seq)
    ) as pipe:
        assert _wait_for(lambda: not pipe.running, timeout=3.0), "producer never finished"
        assert len(delivered) == 5, delivered
        # Once the input has ended, wait reports the end rather than hanging.
        assert pipe.wait(timeout=0.5) is None
        assert pipe.running is False