from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ..capture import (
    CaptureConfig,
    CaptureSource,
    ReconnectPolicy,
    is_live_camera,
    normalize_source,
    open_capture,
)
from ..configs import load_config
from ..frame_processors import FrameUnwarper
from ..sensor import Sensor
from ..unroll import (
    build_or_load_domain,
    load_calib_json,
    resolve_uv_runtime_sizes,
    scale_camera_geometry,
)
from .exports import (
    DEFAULT_FINGER_EXPORT_DIR,
    DEFAULT_FINGER_MAP_CACHE,
    FingerExportPaths,
    load_finger_export_bundle,
)


def normalize_video_source(video):
    """Accept camera index, numeric string, or filesystem video path."""
    return normalize_source(video)


def open_source(video, config, *, auto_reconnect: bool = False, verbose: bool = False):
    """Open a capture source for finger pipelines (compat: returns ``.cap``-like handle).

    Historically returned a bare ``cv2.VideoCapture``. Callers that need the raw
    OpenCV handle can use ``source.cap``; prefer ``source.read()``.
    """
    capture_config = CaptureConfig.from_camera_dict(
        normalize_video_source(video),
        config.get("camera", {}),
        reconnect=ReconnectPolicy(enabled=auto_reconnect),
        verbose=verbose,
        background=False,
    )
    source = open_capture(capture_config, start_background=False)
    cap = source.cap
    cap._orisys_worker = source  # type: ignore[attr-defined]
    return cap


class TactileProcessor:
    """Narrow tactile-analysis wrapper around ``Sensor`` external-frame mode."""

    def __init__(self, sensor: Sensor, *, default_valid_mask=None):
        self._sensor = sensor
        self._default_valid_mask = default_valid_mask

    @property
    def raw_sensor(self) -> Sensor:
        """Escape hatch for legacy callers that still need the full ``Sensor``."""
        return self._sensor

    @property
    def info(self):
        return self._sensor.info

    @property
    def frame(self):
        return self._sensor.frame

    @property
    def img(self):
        return self._sensor.img

    @property
    def flow(self):
        return self._sensor.flow

    @property
    def vnormal(self):
        return self._sensor.vnormal

    @property
    def vshear(self):
        return self._sensor.vshear

    @property
    def depth_map(self):
        return self._sensor.depth_map

    @property
    def mask_stats(self):
        return self._sensor.mask_stats

    @property
    def valid_mask(self):
        return self._sensor.valid_mask

    @property
    def force_units(self) -> tuple[str, str, str]:
        return self._sensor.force_units

    def push_frame(self, frame, *, already_gray=False, valid_mask=None):
        """Feed one processed tactile frame into the analysis pipeline."""
        mask = self._default_valid_mask if valid_mask is None else valid_mask
        return self._sensor.get_img(
            frame=frame,
            already_processed=True,
            already_gray=already_gray,
            valid_mask=mask,
        )

    def warmup(self):
        return self._sensor.warmup_deformation()

    def pipeline(self, produce, **kwargs):
        """Start a :class:`~orisys.pipeline.StreamPipeline` over this processor's own Sensor.

        ``produce`` is the frame source: a zero-argument callable returning the next unrolled
        frame, or ``None`` at end of input. :class:`FingerRuntime` supplies it through
        :meth:`FingerRuntime.pipeline`, so call this directly only when you are producing frames
        yourself.

        Any keyword argument is forwarded to ``StreamPipeline``, so ``contact=True`` and
        ``include=(...)`` work exactly as they do there.
        """
        from orisys.pipeline import StreamPipeline

        return StreamPipeline(self._sensor, produce, **kwargs)

    def compute_deformation(self, *args, **kwargs):
        return self._sensor.compute_deformation(*args, **kwargs)

    def compute_contact(self, *args, **kwargs):
        return self._sensor.compute_contact(*args, **kwargs)

    def read_info(self, *args):
        return self._sensor.read_info(*args)

    def reset(self):
        return self._sensor.reset()

    def disconnect(self):
        return self._sensor.disconnect()


@dataclass(frozen=True)
class FingerGeometry:
    """Resolved one-time geometry/bootstrap state for a finger UV runtime."""

    config: dict
    camera_config: dict
    export_paths: FingerExportPaths
    data: dict
    source_K: np.ndarray
    dist_coeffs: np.ndarray | None
    native_width: int
    native_height: int
    source_width: int
    source_height: int
    out_width: int
    out_height: int
    source_scale: float
    map_cache: str
    edge_expand: float


def resolve_finger_geometry(
    *,
    export_dir,
    config="finger",
    map_cache=DEFAULT_FINGER_MAP_CACHE,
    source_scale=None,
    calib=None,
) -> FingerGeometry:
    """Resolve config/export/camera geometry for a UV finger runtime."""
    config = load_config(config)
    camera_config = config.get("camera", {})
    out_width, out_height, source_scale = resolve_uv_runtime_sizes(
        camera_config,
        source_scale=source_scale,
    )
    bundle = load_finger_export_bundle(
        export_dir,
        map_cache=map_cache,
        source_scale=source_scale,
    )

    if calib:
        source_K, dist_coeffs, img_size = load_calib_json(calib)
        native_width, native_height = img_size
    else:
        source_K = bundle.data["K"]
        dist_coeffs = None
        native_width = int(bundle.data["width"])
        native_height = int(bundle.data["height"])

    source_K, source_width, source_height = scale_camera_geometry(
        source_K,
        native_width,
        native_height,
        source_scale,
    )

    return FingerGeometry(
        config=config,
        camera_config=camera_config,
        export_paths=bundle.export_paths,
        data=bundle.data,
        source_K=source_K,
        dist_coeffs=dist_coeffs,
        native_width=native_width,
        native_height=native_height,
        source_width=source_width,
        source_height=source_height,
        out_width=out_width,
        out_height=out_height,
        source_scale=float(source_scale),
        map_cache=bundle.map_cache,
        edge_expand=bundle.edge_expand,
    )


class FingerRuntime:
    """Reusable UV finger runtime composed from capture, UV remap, and tactile analysis."""

    def __init__(
        self,
        *,
        geometry: FingerGeometry,
        domain,
        tactile: TactileProcessor,
        frame_processor: FrameUnwarper,
        capture,
        capture_source: CaptureSource | None = None,
    ):
        self.geometry = geometry
        self.domain = domain
        self.tactile = tactile
        self.sensor = tactile
        self.frame_processor = frame_processor
        self.capture = capture
        self.capture_source = capture_source
        # Deprecated alias.
        self.capture_worker = capture_source

    @classmethod
    def open(
        cls,
        *,
        export_dir=DEFAULT_FINGER_EXPORT_DIR,
        config="finger",
        map_cache=DEFAULT_FINGER_MAP_CACHE,
        source_scale=None,
        calib=None,
        cal_path=None,
        calibration_dir=None,
        device_id=None,
        video=0,
        camera_id=None,
        flip_source_vertical=True,
        verbose=False,
        mirror=False,
        auto_reconnect: bool = False,
        capture_mode: str = "latest",
    ) -> FingerRuntime:
        """Open a live finger runtime from export bundle, config, and video source."""
        del capture_mode  # deprecated; latest-frame capture only
        if camera_id is not None:
            video = camera_id
        video = normalize_video_source(video)
        geometry = resolve_finger_geometry(
            export_dir=export_dir,
            config=config,
            map_cache=map_cache,
            source_scale=source_scale,
            calib=calib,
        )

        mesh_path, intrinsics_path, pose_path = geometry.export_paths.as_tuple()
        domain = build_or_load_domain(
            mesh_path,
            intrinsics_path,
            pose_path,
            geometry.map_cache,
            geometry.source_K,
            geometry.source_width,
            geometry.source_height,
            geometry.out_width,
            geometry.out_height,
            edge_expand=geometry.edge_expand,
            include_surface=True,
            verbose=verbose,
            source_scale=geometry.source_scale,
        )
        domain.remap.ensure_fast_maps()

        sensor = Sensor(
            "external", config_name=config, cal_path=cal_path,
            calibration_dir=calibration_dir, device_id=device_id,
            verbose=verbose, mirror=mirror,
        )

        capture_config = CaptureConfig.from_camera_dict(
            video,
            geometry.config.get("camera", {}),
            reconnect=ReconnectPolicy(enabled=auto_reconnect),
            verbose=verbose,
            background=False,
        )
        capture_source = open_capture(capture_config, start_background=False)
        if not capture_source.is_open():
            raise RuntimeError("Cannot open camera/video source")
        if is_live_camera(video):
            capture_source.start_background()
        capture = capture_source.cap

        if not is_live_camera(video):
            # File playback: wall-clock cadence is not signal time, so pin the
            # force-filter design rate to the configured camera fps.
            try:
                fps = float(geometry.config.get("camera", {}).get("fps", 0.0) or 0.0)
            except (TypeError, ValueError):
                fps = 0.0
            if fps > 0:
                sensor.configure(lowpass_sample_rate_hz=fps)
        frame_processor = FrameUnwarper(
            domain.remap,
            camera_matrix=geometry.source_K if geometry.dist_coeffs is not None else None,
            dist_coeffs=geometry.dist_coeffs,
            source_width=geometry.source_width,
            source_height=geometry.source_height,
            flip_source_vertical=flip_source_vertical,
        )
        tactile = TactileProcessor(sensor, default_valid_mask=frame_processor.valid_mask)
        return cls(
            geometry=geometry,
            domain=domain,
            tactile=tactile,
            frame_processor=frame_processor,
            capture=capture,
            capture_source=capture_source,
        )

    @property
    def valid_mask(self):
        return self.frame_processor.valid_mask

    def read_source(self):
        if self.capture_source is not None and self.capture_source.running:
            return self.capture_source.read(copy=True, wait_s=1.0)
        if self.capture_source is not None:
            return self.capture_source.read_raw()
        return self.capture.read()

    def process_frame(self, frame, *, color=False):
        """Prepare, unroll, and feed one source frame into the tactile processor."""
        if color:
            prepared = self.frame_processor.prepare_source(frame)
            unrolled = self.frame_processor.flatten_prepared(prepared)
            self.tactile.push_frame(unrolled)
            return unrolled

        prepared_gray = self.frame_processor.prepare_source_gray(frame)
        unrolled = self.frame_processor.flatten_prepared_gray(prepared_gray)
        self.tactile.push_frame(unrolled, already_gray=True)
        return unrolled

    def grab_and_process(self, *, color=False):
        ok, frame = self.read_source()
        if not ok:
            return False, None
        return True, self.process_frame(frame, color=color)

    def get_img(self, *, color=False):
        """Sensor-style alias: capture, unroll, and return the tactile frame or None."""
        ok, unrolled = self.grab_and_process(color=color)
        return unrolled if ok else None

    def warmup(self, *, color=False):
        ok, unrolled = self.grab_and_process(color=color)
        if ok:
            self.tactile.warmup()
        return ok, unrolled

    def pipeline(self, **kwargs):
        """Return a :class:`~orisys.pipeline.StreamPipeline` for this runtime.

        The recommended way to run a finger sensor. Each finished frame arrives as a single
        :class:`~orisys.pipeline.FrameSnapshot`, so there are no stage calls to sequence and no
        tuple to unpack::

            with runtime.pipeline(contact=True, include=("image", "depth")) as pipe:
                while True:
                    snap = pipe.wait()
                    if snap is None:
                        break
                    print(snap.normal_force, snap.shear_x, snap.shear_y)

        Any keyword argument is forwarded to ``StreamPipeline``. Use ``contact=True`` to run
        contact analysis, and ``include=(...)`` to select which arrays each snapshot carries.
        """
        from orisys.pipeline import StreamPipeline

        return StreamPipeline.for_finger_runtime(self, **kwargs)

    def close(self):
        if self.capture_source is not None:
            self.capture_source.close()
            self.capture_source = None
            self.capture_worker = None
            self.capture = None
        elif self.capture is not None:
            self.capture.release()
            self.capture = None
        self.tactile.disconnect()