"""Monitoring collectors: one per event stream, each with start/stop/dump."""

import os
import threading
from abc import ABC, abstractmethod
from collections import deque
from datetime import datetime
from typing import TYPE_CHECKING, Optional

from feathersdk.utils.feathertypes import NamedTuple
from feathersdk.utils.logger import warning

from ..motors.robstride import RobstrideMotor, RobstrideMotorError

if TYPE_CHECKING:
    from ..motors.motors_manager import MotorsManager


def _parquet_metadata(reference_wall_time: float, reference_monotonic: float) -> dict:
    return {
        b"reference_wall_time": str(reference_wall_time).encode(),
        b"reference_monotonic": str(reference_monotonic).encode(),
    }


class MonitoringCollector(ABC):
    """Protocol for monitoring event stream collectors."""

    @abstractmethod
    def start(self) -> None:
        """Register listeners, reset state. Called when monitoring platform starts."""
        pass

    @abstractmethod
    def stop(self) -> None:
        """Unregister listeners. Platform will call dump() after stop to flush."""
        pass

    @abstractmethod
    def dump(
        self,
        session_dir: str,
        reference_wall_time: float,
        reference_monotonic: float,
    ) -> None:
        """Write buffered data to Parquet in session_dir. Use given metadata."""
        pass


class MotorFeedbackRecord(NamedTuple):
    timestamp: float
    motor_name: str
    angle: float
    velocity: float
    torque: float
    temp: float
    mode: int
    errors: int


class MotorFeedbackCollector(MonitoringCollector):
    """Collects motor feedback at full rate; dumps to motor_feedback_*.parquet."""

    def __init__(self, motors_manager: "MotorsManager", config: dict):
        self.motors_manager = motors_manager
        self._max_buffer_size = config.get("max_buffer_size", 500_000)
        self._buffer: deque = deque(maxlen=self._max_buffer_size)
        self._buffer_lock = threading.Lock()
        self._enabled = False

    def _get_record(self, motor: RobstrideMotor, timestamp: float) -> MotorFeedbackRecord:
        return MotorFeedbackRecord(
            timestamp=timestamp,
            motor_name=motor.motor_name,
            angle=motor.angle.value,
            velocity=motor.velocity.value,
            torque=motor.torque.value,
            temp=motor.temp.value,
            mode=motor.mode.value.value,
            errors=RobstrideMotorError.errors_to_int(motor.errors.value),
        )

    def _on_motor_feedback(self, motor: RobstrideMotor, timestamp: float) -> None:
        try:
            record = self._get_record(motor, timestamp)
            with self._buffer_lock:
                self._buffer.append(record)
        except Exception as e:
            warning("MotorFeedbackCollector: Failed to record feedback for motor {0}: {1}", motor.motor_name, e)

    def start(self) -> None:
        RobstrideMotor.add_feedback_listener(self._on_motor_feedback)
        self._enabled = True

    def stop(self) -> None:
        if self._enabled:
            RobstrideMotor.remove_feedback_listener(self._on_motor_feedback)
            self._enabled = False

    def dump(
        self,
        session_dir: str,
        reference_wall_time: float,
        reference_monotonic: float,
    ) -> None:
        with self._buffer_lock:
            if not self._buffer:
                return
            old_buffer = self._buffer
            self._buffer = deque(maxlen=self._max_buffer_size)

        collector_dir = os.path.join(session_dir, "motor_feedback")
        os.makedirs(collector_dir, exist_ok=True)

        try:
            import pandas as pd
            import pyarrow as pa
            import pyarrow.parquet as pq

            records = list(old_buffer)
            columns = [
                "timestamp", "motor_name", "angle", "velocity",
                "torque", "temp", "mode", "errors",
            ]
            df = pd.DataFrame(records, columns=columns)
            table = pa.Table.from_pandas(df)
            merged = {**(table.schema.metadata or {}), **_parquet_metadata(reference_wall_time, reference_monotonic)}
            table = table.replace_schema_metadata(merged)
            path = os.path.join(collector_dir, f"motor_feedback_{datetime.now().strftime('%Y%m%d_%H%M%S')}.parquet")
            pq.write_table(table, path)
        except Exception as e:
            warning("MotorFeedbackCollector: Failed to write parquet file: {0}", e)
            raise


class MotorCommandRecord(NamedTuple):
    """One record per loc_ref or spd_ref command; use float('nan') for N/A."""
    timestamp: float
    motor_name: str
    target_position: float
    target_velocity: float


class MotorCommandCollector(MonitoringCollector):
    """Collects motor target position/velocity at command-sent time."""

    def __init__(self, motors_manager: "MotorsManager", config: dict):
        self.motors_manager = motors_manager
        self._max_buffer_size = config.get("max_buffer_size", 500_000)
        self._buffer: deque = deque(maxlen=self._max_buffer_size)
        self._buffer_lock = threading.Lock()
        self._enabled = False

    def _on_motor_command(
        self, motor: RobstrideMotor, param_name: str, value: float, timestamp: float
    ) -> None:
        if param_name == "loc_ref":
            record = MotorCommandRecord(
                timestamp=timestamp,
                motor_name=motor.motor_name,
                target_position=float(value),
                target_velocity=float("nan"),
            )
        else:
            record = MotorCommandRecord(
                timestamp=timestamp,
                motor_name=motor.motor_name,
                target_position=float("nan"),
                target_velocity=float(value),
            )
        with self._buffer_lock:
            self._buffer.append(record)

    def start(self) -> None:
        RobstrideMotor.add_command_listener(self._on_motor_command)
        self._enabled = True

    def stop(self) -> None:
        if self._enabled:
            RobstrideMotor.remove_command_listener(self._on_motor_command)
            self._enabled = False

    def dump(
        self,
        session_dir: str,
        reference_wall_time: float,
        reference_monotonic: float,
    ) -> None:
        with self._buffer_lock:
            if not self._buffer:
                return
            old_buffer = self._buffer
            self._buffer = deque(maxlen=self._max_buffer_size)

        collector_dir = os.path.join(session_dir, "motor_commands")
        os.makedirs(collector_dir, exist_ok=True)

        import pandas as pd
        import pyarrow as pa
        import pyarrow.parquet as pq

        records = list(old_buffer)
        columns = ["timestamp", "motor_name", "target_position", "target_velocity"]
        df = pd.DataFrame(records, columns=columns)
        table = pa.Table.from_pandas(df)
        merged = {**(table.schema.metadata or {}), **_parquet_metadata(reference_wall_time, reference_monotonic)}
        table = table.replace_schema_metadata(merged)
        path = os.path.join(collector_dir, f"motor_commands_{datetime.now().strftime('%Y%m%d_%H%M%S')}.parquet")
        pq.write_table(table, path)


class StepTimingCollector(MonitoringCollector):
    """Collects root and per-platform step durations; single Parquet with system_name column."""

    ROOT_SYSTEM_NAME = "root"

    def __init__(self, motors_manager: "MotorsManager", config: dict):
        self._max_buffer_size = config.get("step_timing_max_buffer_size", 60_000)
        self._buffer: deque = deque(maxlen=self._max_buffer_size)
        self._buffer_lock = threading.Lock()

    def record_step_timing(
        self,
        step_duration: float,
        timestamp: float,
        step_start_delay: Optional[float] = None,
    ) -> None:
        """Record overall root loop step duration (downsampled by caller)."""
        with self._buffer_lock:
            self._buffer.append(
                (timestamp, step_duration, self.ROOT_SYSTEM_NAME, step_start_delay)
            )

    def record_system_step_timing(
        self,
        system_name: str,
        step_duration: float,
        timestamp: float,
        step_start_delay: Optional[float] = None,
    ) -> None:
        """Record one platform's step() duration."""
        with self._buffer_lock:
            self._buffer.append(
                (timestamp, step_duration, system_name, step_start_delay)
            )

    def start(self) -> None:
        pass

    def stop(self) -> None:
        pass

    def dump(
        self,
        session_dir: str,
        reference_wall_time: float,
        reference_monotonic: float,
    ) -> None:
        with self._buffer_lock:
            if not self._buffer:
                return
            old_buffer = self._buffer
            self._buffer = deque(maxlen=self._max_buffer_size)

        collector_dir = os.path.join(session_dir, "step_timing")
        os.makedirs(collector_dir, exist_ok=True)

        import pandas as pd
        import pyarrow as pa
        import pyarrow.parquet as pq

        records = list(old_buffer)
        columns = ["timestamp", "step_duration", "system_name", "step_start_delay"]
        df = pd.DataFrame(records, columns=columns)
        table = pa.Table.from_pandas(df)
        merged = {
            **(table.schema.metadata or {}),
            **_parquet_metadata(reference_wall_time, reference_monotonic),
        }
        table = table.replace_schema_metadata(merged)
        path = os.path.join(collector_dir, f"step_timing_{datetime.now().strftime('%Y%m%d_%H%M%S')}.parquet")
        pq.write_table(table, path)
