"""A source element to watch directories for new frame files in real-time."""

# Copyright (C) 2022      Ron Tapia
# Copyright (C) 2024-2025 Becca Ewing, Yun-Jing Huang

from __future__ import annotations

import os
import queue
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, ClassVar

import gpstime
import gwframe
from sgnts.base import Offset, SeriesBuffer, TSResourceSource
from watchdog.events import PatternMatchingEventHandler
from watchdog.observers import Observer

if TYPE_CHECKING:
    import logging

    from gwframe import TimeSeries
    from sgn.base import SourcePad
    from sgn.subprocess import WorkerContext

from sgn_gwframe.base import from_T050017


class WorkerAction(Enum):
    """Actions that the worker can take in a single iteration."""

    DATA = "data"
    GAP = "gap"
    HEARTBEAT = "heartbeat"


class _GWFrameFileEventHandler(PatternMatchingEventHandler):
    """Custom watchdog event handler for tracking new frame files."""

    def __init__(self, queue, watch_suffix, current_watermark):
        self.queue = queue
        self.watch_suffix = watch_suffix
        self.current_watermark = current_watermark
        super().__init__(patterns=[f"*{self.watch_suffix}"])

    def on_closed(self, event):
        if event.is_directory:
            return
        self._handle_event(event.src_path)

    def on_moved(self, event):
        if event.is_directory:
            return
        self._handle_event(event.dest_path)

    def _handle_event(self, path):
        extension = os.path.splitext(path)[1]
        if extension != self.watch_suffix:
            return
        # skip any files past the current watermark
        _, _, start, _ = from_T050017(path)
        if start < self.current_watermark:
            return
        self.queue.put((path, start))


@dataclass(kw_only=True)
class GWFrameWatchSource(TSResourceSource):
    """Source element to watch a directory for new frame files in real-time.

    Monitors a directory for new frame files, automatically reading and sending
    data as files arrive. Handles discontinuities with gap detection, validates
    data integrity, and supports both live streaming (start=None) and replay
    from a specific GPS time.

    For multi-IFO coherent analysis, use multiple GWFrameWatchSource elements
    (one per IFO) and let the pipeline handle synchronization.

    Args:
        channels:
            list[str], channel names to read from frame files.
            Source pads will be automatically generated for each channel, with
            channel name as pad name.
        watch_dir:
            str, directory path to watch for new frame files.
        discont_wait_time:
            float, time to wait before sending a gap buffer when no data arrives,
            in seconds. Default: 60
        queue_timeout:
            float, time to wait for next file from the queue, in seconds. Default: 1
        watch_suffix:
            str, filename suffix to watch for. Default: ".gwf"

    Example:
        Basic usage::

            from sgn_gwframe.sources import GWFrameWatchSource
            from sgn.apps import Pipeline

            src = GWFrameWatchSource(
                name="L1_data",
                channels=["L1:GDS-CALIB_STRAIN"],
                watch_dir="/data/frames/L1",
                duration=1,  # 1 second buffers
            )

            pipeline = Pipeline()
            pipeline.insert(src, ...)
            pipeline.run()

    Note:
        Requires frame files following T050017 naming convention:
        {site}-{description}-{gpstime}-{duration}.gwf
    """

    allow_dynamic_source_pads: ClassVar[bool] = False

    channels: list[str]
    watch_dir: str
    discont_wait_time: float = 60
    queue_timeout: float = 1
    watch_suffix: str = ".gwf"

    @property
    def static_source_pads(self) -> list[str]:  # type: ignore[override]
        """Define source pads from channels."""
        return self.channels

    def worker_process(  # type: ignore[override]
        self,
        context: WorkerContext,
        channels: list[str],
        watch_dir: str,
        srcs: dict[str, SourcePad],
        watch_suffix: str,
        queue_timeout: float,
        discont_wait_time: float,
        start: int | None,
        logger: logging.Logger,
    ) -> None:
        """Worker process that watches directory and reads in frame data."""
        # Initialize on first call
        if "initialized" not in context.state:
            # Get starting time
            if start is None:
                start_time = int(gpstime.gpsnow())
                logger.debug(
                    "No start time specified, using current GPS time: %d", start_time
                )
            else:
                start_time = start
                logger.debug("Starting at GPS time: %d", start_time)

            # Setup watchdog observer and queue
            file_queue: queue.Queue[tuple[str, int]] = queue.Queue()
            event_handler = _GWFrameFileEventHandler(
                file_queue, watch_suffix, start_time
            )
            observer = Observer()
            observer.schedule(event_handler, path=watch_dir)
            observer.daemon = True
            observer.start()

            # Store state
            context.state["initialized"] = True
            context.state["file_queue"] = file_queue
            context.state["observer"] = observer
            context.state["next_buffer_time"] = start_time
            context.state["last_file_time"] = None
            context.state["rates"] = {}
            context.state["buffer_duration"] = None
            context.state["pending_frame"] = None
            context.state["pending_file_time"] = None

        # Check if we should stop
        if context.should_stop():
            observer = context.state.get("observer")  # type: ignore[assignment]
            if observer is not None:
                observer.stop()
                observer.join()
            return

        # Get state
        file_queue = context.state["file_queue"]
        next_buffer_time = context.state["next_buffer_time"]
        last_file_time = context.state["last_file_time"]
        rates = context.state["rates"]
        buffer_duration = context.state["buffer_duration"]
        pending_frame = context.state["pending_frame"]
        pending_file_time = context.state["pending_file_time"]

        # Determine what action to take based on current state
        action: WorkerAction | None = None
        frame_data = None
        gap_start_time = None
        gap_end_time = None
        file_time = None

        # Case 1: Handle pending discontinuity
        if pending_frame is not None:
            gap_start_time = next_buffer_time
            gap_end_time = gap_start_time + buffer_duration
            action = WorkerAction.GAP

            # Check if gap is filled - then also send the pending data
            if gap_end_time >= pending_file_time:
                frame_data = pending_frame

        # Case 2: Try to get a file from the queue
        else:
            try:
                filepath, file_time = _drain_and_get_file(
                    file_queue, queue_timeout, next_buffer_time, last_file_time, logger
                )
            except queue.Empty:
                # Timeout case - need rates to send heartbeat or gap
                if not rates or buffer_duration is None:
                    # Try to initialize rates from a sample file in the directory
                    result = _initialize_rates_from_sample(
                        watch_dir, watch_suffix, channels, logger
                    )
                    if result is not None:
                        rates, buffer_duration = result
                        context.state["rates"] = rates
                        context.state["buffer_duration"] = buffer_duration
                    else:
                        # Can't send anything without knowing rates
                        return

                # Determine if we should send gap or heartbeat
                time_waiting = gpstime.gpsnow() - next_buffer_time
                if time_waiting >= discont_wait_time:
                    logger.warning(
                        "No data received for %.1f seconds (>= %.1f second timeout), "
                        "sending gap buffer",
                        time_waiting,
                        discont_wait_time,
                    )
                    gap_start_time = next_buffer_time
                    gap_end_time = gap_start_time + buffer_duration
                    action = WorkerAction.GAP
                else:
                    action = WorkerAction.HEARTBEAT
            else:
                # Successfully got a file from queue - try to read it
                try:
                    frame_data = _read_and_validate_frame(
                        filepath,
                        channels,
                        rates,
                        buffer_duration,
                        file_time,
                    )
                except (FileNotFoundError, OSError, RuntimeError, ValueError):
                    logger.exception("Failed to read or validate file %s", filepath)
                    # Send gap if we know duration, otherwise heartbeat
                    if buffer_duration is not None:
                        gap_start_time = next_buffer_time
                        gap_end_time = gap_start_time + buffer_duration
                        action = WorkerAction.GAP
                    else:
                        action = WorkerAction.HEARTBEAT
                else:
                    # Successfully read frame - process it
                    # Initialize rates if this is the first file
                    if not rates:
                        rates, buffer_duration = _initialize_rates_and_duration(
                            frame_data, channels
                        )
                        context.state["rates"] = rates
                        context.state["buffer_duration"] = buffer_duration

                    # Check for discontinuity
                    if last_file_time is not None and file_time > next_buffer_time:
                        logger.warning(
                            "Discontinuity detected: expected time %d, got %d "
                            "(gap of %.1f seconds)",
                            next_buffer_time,
                            file_time,
                            file_time - next_buffer_time,
                        )
                        # Store frame and send gap next iteration
                        context.state["pending_frame"] = frame_data
                        context.state["pending_file_time"] = file_time
                        gap_start_time = next_buffer_time
                        gap_end_time = gap_start_time + buffer_duration
                        action = WorkerAction.GAP
                        frame_data = None  # Don't send data yet
                    else:
                        # Continuous data
                        action = WorkerAction.DATA

        # Execute action based on determined state
        # Initialize rates if needed for GAP or HEARTBEAT actions
        if action in (WorkerAction.GAP, WorkerAction.HEARTBEAT) and (
            not rates or buffer_duration is None
        ):
            result = _initialize_rates_from_sample(
                watch_dir, watch_suffix, channels, logger
            )
            if result is not None:
                rates, buffer_duration = result
                context.state["rates"] = rates
                context.state["buffer_duration"] = buffer_duration
            else:
                msg = (
                    f"Cannot initialize GWFrameWatchSource: sample rates for "
                    f"requested channels unknown and no existing files exist "
                    f"in {watch_dir}"
                )
                raise RuntimeError(msg)

        if action == WorkerAction.DATA:
            assert frame_data is not None
            assert file_time is not None
            _send_data_buffers(context, frame_data, srcs)
            logger.debug(
                "Sent frame at time %d, delay=%.3f seconds",
                file_time,
                gpstime.gpsnow() - file_time,
            )
            context.state["next_buffer_time"] = int(file_time + buffer_duration)
            context.state["last_file_time"] = file_time

        elif action == WorkerAction.GAP:
            assert gap_start_time is not None
            assert gap_end_time is not None
            _send_gap_buffers(
                context, channels, srcs, rates, gap_start_time, gap_end_time
            )
            context.state["next_buffer_time"] = gap_end_time

            # If we also have pending data that's now ready, send it
            if frame_data is not None:
                _send_data_buffers(context, frame_data, srcs)
                context.state["next_buffer_time"] = int(
                    pending_file_time + buffer_duration
                )
                context.state["last_file_time"] = pending_file_time
                context.state["pending_frame"] = None
                context.state["pending_file_time"] = None

        elif action == WorkerAction.HEARTBEAT:
            _send_heartbeat_buffers(context, channels, srcs, rates, next_buffer_time)


def _initialize_rates_and_duration(
    frame: dict[str, TimeSeries], channels: list[str]
) -> tuple[dict[str, int], float]:
    """Extract sample rates and buffer duration from first frame.

    Returns:
        tuple of (rates dict, buffer_duration in seconds)
    """
    rates = {channel: int(series.sample_rate) for channel, series in frame.items()}
    buffer_duration = frame[channels[0]].duration
    return rates, buffer_duration


def _send_gap_buffers(
    context: WorkerContext,
    channels: list[str],
    srcs: dict[str, SourcePad],
    rates: dict[str, int],
    gap_start: int,
    gap_end: int,
) -> None:
    """Send gap buffers for all channels between gap_start and gap_end."""
    for channel in channels:
        pad = srcs[channel]
        gap_samples = int((gap_end - gap_start) * rates[channel])
        gap_buf = SeriesBuffer(
            offset=Offset.fromsec(gap_start),
            sample_rate=rates[channel],
            data=None,
            shape=(gap_samples,),
        )
        context.output_queue.put((pad, gap_buf))


def _send_data_buffers(
    context: WorkerContext,
    frame: dict[str, TimeSeries],
    srcs: dict[str, SourcePad],
) -> None:
    """Send data buffers for all channels from the frame."""
    for channel, series in frame.items():
        pad = srcs[channel]
        buf = SeriesBuffer(
            offset=Offset.fromsec(series.t0),
            sample_rate=int(series.sample_rate),
            data=series.array,
        )
        context.output_queue.put((pad, buf))


def _send_heartbeat_buffers(
    context: WorkerContext,
    channels: list[str],
    srcs: dict[str, SourcePad],
    rates: dict[str, int],
    offset_time: int,
) -> None:
    """Send zero-duration heartbeat buffers to keep pipeline alive."""
    for channel in channels:
        pad = srcs[channel]
        heartbeat_buf = SeriesBuffer(
            offset=Offset.fromsec(offset_time),
            sample_rate=rates[channel],
            data=None,
            shape=(0,),
        )
        context.output_queue.put((pad, heartbeat_buf))


def _drain_and_get_file(
    file_queue: queue.Queue[tuple[str, int]],
    queue_timeout: float,
    next_buffer_time: int,
    last_file_time: int | None,
    logger: logging.Logger,
) -> tuple[str, int]:
    """Drain old files from queue and return current/future file.

    Raises:
        queue.Empty: If no file available within timeout
    """
    while True:
        filepath, file_time = file_queue.get(timeout=queue_timeout)

        # Discard files that are in the past
        if last_file_time is not None and file_time < next_buffer_time:
            logger.debug(
                "Discarding old file at time %d (expected %d)",
                file_time,
                next_buffer_time,
            )
            continue  # Get next file

        # Found current/future file
        return filepath, file_time


def _read_and_validate_frame(
    filepath: str,
    channels: list[str],
    rates: dict[str, int],
    buffer_duration: float | None,
    file_time: int,
) -> dict[str, TimeSeries]:
    """Read and validate frame file.

    Returns:
        Frame data dict

    Raises:
        FileNotFoundError: If file doesn't exist
        OSError, RuntimeError, ValueError: If read or validation fails
    """
    # Check file exists
    if not os.path.exists(filepath):
        msg = f"File no longer exists: {filepath}"
        raise FileNotFoundError(msg)

    # Try reading with error recovery
    frame = gwframe.read(filepath, channel=channels)

    # Validate if we have expected rates/duration
    if rates and buffer_duration is not None:
        for channel, series in frame.items():
            # Validate sample rate
            actual_rate = int(series.sample_rate)
            expected_rate = rates[channel]
            if actual_rate != expected_rate:
                msg = (
                    f"Sample rate mismatch for {channel}: "
                    f"expected {expected_rate}, got {actual_rate}"
                )
                raise ValueError(msg)

            # Validate duration
            if series.duration != buffer_duration:
                msg = (
                    f"Duration mismatch for {channel}: "
                    f"expected {buffer_duration}, got {series.duration}"
                )
                raise ValueError(msg)

            # Validate t0 matches filename
            if series.t0 != file_time:
                msg = (
                    f"Time mismatch for {channel}: "
                    f"filename suggests {file_time}, data has {series.t0}"
                )
                raise ValueError(msg)

    return frame


def _initialize_rates_from_sample(
    watch_dir: str,
    watch_suffix: str,
    channels: list[str],
    logger: logging.Logger,
) -> tuple[dict[str, int], float] | None:
    """Initialize rates and duration by reading a sample file from watch_dir.

    Returns:
        (rates dict, buffer_duration) if successful, None otherwise
    """
    files = sorted(
        [f for f in os.listdir(watch_dir) if f.endswith(watch_suffix)],
        reverse=True,  # Get most recent file
    )
    if not files:
        return None

    sample_file = os.path.join(watch_dir, files[0])
    try:
        sample_frame = gwframe.read(sample_file, channel=channels)
        rates, buffer_duration = _initialize_rates_and_duration(sample_frame, channels)
        logger.debug(
            "Initialized from sample file %s: rates=%s, duration=%s",
            sample_file,
            rates,
            buffer_duration,
        )
    except (OSError, RuntimeError, ValueError, KeyError, IndexError) as e:
        logger.warning("Could not read sample file %s: %s", sample_file, e)
        return None
    else:
        return rates, buffer_duration
