"""Camera subscriber module for receiving video frames via ZMQ.

This module provides functionality to receive raw YUYV video frames from a ZMQ
publisher, convert them to BGR for display, and show them using OpenCV.

This is the main subscriber module. It has been refactored into smaller modules:
- subscriber_terminal.py: Terminal input handling
- subscriber_frame_utils.py: Frame processing utilities
- subscriber_recording.py: Recording functionality
"""

import hashlib
import json
import os
import time
from typing import Callable, Optional

import cv2
import numpy as np
import zmq
import zstandard

from .image_processing_transformations import ImageTransformer
from .image_processing_texts import Color
from .image_saving import save_current_frame
from .subscriber_terminal import TerminalInputHandler
from .subscriber_frame_utils import (
    convert_yuyv_to_bgr,
    scale_frame,
    crop_frame,
    NoSignalManager,
)
from .subscriber_recording import (
    RecordingManager,
    get_duration_step,
    get_frame_step,
    save_frame_as_format,
)

# Use X11 (xcb) platform plugin for Qt instead of offscreen
os.environ["QT_QPA_PLATFORM"] = "xcb"
# Use system fontconfig instead of bundled Qt fonts (which are missing)
os.environ["QT_QPA_FONTDIR"] = "/usr/share/fonts"


class CameraSubscriber:
    """Subscribes to camera frames via ZMQ.

    Attributes:
        connect_address: ZMQ connect address (e.g., "tcp://localhost:5555").
        frame_callback: Optional callback function to process received frames.
        display_width: Display width for the window (default: 640).
        display_height: Display height for the window (default: 480).
    """

    def __init__(
        self,
        connect_address: str = "tcp://localhost:5555",
        frame_callback: Optional[Callable[[np.ndarray], None]] = None,
        display_width: int = 640,
        display_height: int = 480,
        fov_degrees: float = 150.0,
        no_signal_timeout: float = 5.0,
        max_frames_in_fits: int = 10,
        record_duration: int = 10,
        burst_mode: str = "individual",
    ) -> None:
        """Initialize the camera subscriber.

        Args:
            connect_address: ZMQ connect address (default: "tcp://localhost:5555").
            frame_callback: Optional callback function for frame processing.
                          If None, frames will be displayed in a window.
            display_width: Display width for the window (default: 640).
            display_height: Display height for the window (default: 480).
            fov_degrees: Field of view in degrees for barrel correction (default: 150.0).
            no_signal_timeout: Timeout in seconds before showing 'no signal' image.
            max_frames_in_fits: Maximum number of frames to save in a FITS sequence.
            record_duration: Duration in seconds for time-based recording (default: 10).
            burst_mode: Burst mode, 'individual' or 'sequence' (default: 'individual').
        """
        self.connect_address = connect_address
        self.frame_callback = frame_callback
        self.display_width = display_width
        self.display_height = display_height
        self.no_signal_timeout = no_signal_timeout

        self._context: Optional[zmq.Context] = None
        self._socket: Optional[zmq.Socket] = None
        self._running = False
        self._input_handler: Optional[TerminalInputHandler] = None

        # Window state management
        self._window_visible = False
        self._window_name = "ZMQ Camera Stream"

        # Display mode: 'scale' (0) or crop positions
        self._display_mode = "scale"  # Default: scale full image to 640 width

        # Statistics
        self._last_seq = 0
        self._dropped_frames = 0

        # Decompressor for zstd compression
        self._decompressor: Optional[zstandard.ZstdDecompressor] = None

        # Image transformer for post-processing
        self._transformer = ImageTransformer(fov_degrees=fov_degrees)

        # 2x scaling toggle for display
        self._scale_2x = False

        # No signal detection (using refactored module)
        self._no_signal_manager = NoSignalManager(
            display_width=display_width,
            display_height=display_height,
            connect_address=connect_address,
            no_signal_timeout=no_signal_timeout,
        )

        # Frame storage for saving
        self._last_frame_data: Optional[bytes] = None
        self._last_frame_header: Optional[dict] = None

        # Recording manager (using refactored module)
        self._recording_manager = RecordingManager(
            max_frames_in_fits=max_frames_in_fits,
            record_duration=record_duration,
            burst_mode=burst_mode,
        )

    def _show_window(self) -> None:
        """Create and show the OpenCV window."""
        if not self._window_visible:
            # Use WINDOW_AUTOSIZE to maintain original frame dimensions
            cv2.namedWindow(self._window_name, cv2.WINDOW_AUTOSIZE)
            self._window_visible = True
            print("\n[Window shown - press 'h' to hide]")

    def _hide_window(self) -> None:
        """Destroy the OpenCV window."""
        if self._window_visible:
            cv2.destroyWindow(self._window_name)
            self._window_visible = False
            print("\n[Window hidden - press 'v' to show]")

    def _verify_checksum(self, data: bytes, expected_checksum: str) -> bool:
        """Verify MD5 checksum of data.

        Args:
            data: Raw bytes to verify.
            expected_checksum: Expected hexadecimal checksum.

        Returns:
            True if checksum matches, False otherwise.
        """
        computed = hashlib.md5(data).hexdigest()
        return computed == expected_checksum

    def _process_frame(self, frame: np.ndarray, header: dict) -> None:
        """Process a received frame.

        Args:
            frame: The decoded frame to process.
            header: Frame header dictionary.
        """
        # Apply image transformation (e.g., barrel correction)
        transformed_frame = self._transformer.transform(frame)

        # Apply display mode (scale or crop)
        if self._display_mode == "scale":
            display_frame = scale_frame(transformed_frame, self.display_width)
        else:  # crop modes
            display_frame = crop_frame(
                transformed_frame,
                position=self._display_mode,
                display_width=self.display_width,
                display_height=self.display_height,
            )

        # Apply 2x scaling if enabled (just before display)
        if self._scale_2x:
            display_frame = cv2.resize(
                display_frame,
                None,
                fx=2.0,
                fy=2.0,
                interpolation=cv2.INTER_LINEAR,
            )

        # Display frame if window is visible
        if self._window_visible:
            # Ensure window is created before displaying
            self._show_window()
            cv2.imshow(self._window_name, display_frame)
            # Process window events (needed for window to respond)
            cv2.waitKey(1)

        # Call user callback if provided (pass original transformed frame)
        if self.frame_callback is not None:
            self.frame_callback(transformed_frame)

    def _save_frame_as_format(self, format: str) -> None:
        """Save current frame as JPG or PNG.

        Args:
            format: 'jpg' or 'png'.
        """
        save_frame_as_format(
            self._last_frame_data,
            self._last_frame_header,
            format,
            convert_yuyv_to_bgr,
        )

    def _handle_keypress(self, key: str) -> bool:
        """Handle a keypress from the terminal.

        Args:
            key: The character that was pressed.

        Returns:
            True if the application should continue running, False to quit.
        """
        key = key.lower()

        if key == 'q':
            print("\n'q' pressed - quitting...")
            return False
        elif key == 's':
            # Save current frame to FITS file
            if self._last_frame_data is not None and self._last_frame_header is not None:
                filepath = save_current_frame(
                    self._last_frame_header,
                    self._last_frame_data,
                    output_dir=".",
                )
                if filepath:
                    print(f"\n[Frame saved: {filepath}]")
                else:
                    print("\n[Error: Failed to save frame]")
            else:
                print("\n[Error: No frame available to save]")
        elif key == 'j':
            # Save current frame as JPG
            self._save_frame_as_format("jpg")
        elif key == 'p':
            # Save current frame as PNG
            self._save_frame_as_format("png")
        elif key == 'v':
            self._show_window()
        elif key == 'h':
            self._hide_window()
        elif key == '5':
            self._display_mode = "center"
            print("\n[Display mode: crop center 640x480]")
        elif key == '0':
            self._display_mode = "scale"
            print("\n[Display mode: scale to 640 width]")
        elif key == '4':
            self._display_mode = "left"
            print("\n[Display mode: crop left-center 640x480]")
        elif key == '6':
            self._display_mode = "right"
            print("\n[Display mode: crop right-center 640x480]")
        elif key == '8':
            self._display_mode = "top"
            print("\n[Display mode: crop center-top 640x480]")
        elif key == '2':
            self._display_mode = "bottom"
            print("\n[Display mode: crop center-bottom 640x480]")
        elif key == '7':
            self._display_mode = "top-left"
            print("\n[Display mode: crop top-left 640x480]")
        elif key == '9':
            self._display_mode = "top-right"
            print("\n[Display mode: crop top-right 640x480]")
        elif key == '1':
            self._display_mode = "bottom-left"
            print("\n[Display mode: crop bottom-left 640x480]")
        elif key == '3':
            self._display_mode = "bottom-right"
            print("\n[Display mode: crop bottom-right 640x480]")
        elif key == 'b':
            enabled = self._transformer.toggle()
            status = "ON" if enabled else "OFF"
            print(f"\n[Image transformation: barrel correction {status}]")
        elif key == '+':
            new_k1 = self._transformer.adjust_k1(0.01)
            print(f"\n[Barrel correction k1: {new_k1:.2f} (+0.01)]")
        elif key == '-':
            new_k1 = self._transformer.adjust_k1(-0.01)
            print(f"\n[Barrel correction k1: {new_k1:.2f} (-0.01)]")
        elif key == 'x':
            self._scale_2x = not self._scale_2x
            status = "ON" if self._scale_2x else "OFF"
            print(f"\n[2x display scaling: {status}]")
        elif key == 'a':
            # Toggle burst recording (frame-limited)
            if self._recording_manager.is_recording:
                # Stop any active recording
                self._recording_manager.toggle_recording()
            else:
                # Start frame-limited recording
                self._recording_manager.toggle_recording(mode="frame")
        elif key == 't':
            # Toggle burst recording (time-limited)
            if self._recording_manager.is_recording:
                # Stop any active recording
                self._recording_manager.toggle_recording()
            else:
                # Start time-limited recording
                self._recording_manager.toggle_recording(mode="time")
        elif key == 'm':
            # Cycle burst mode
            self._recording_manager.burst_mode = "sequence" if self._recording_manager.burst_mode == "individual" else "individual"
            print(f"\n[Burst mode: {self._recording_manager.burst_mode}]")
        elif key == '[':
            # Decrease record duration with progressive steps
            step = get_duration_step(self._recording_manager.record_duration, decreasing=True)
            self._recording_manager.record_duration = max(1, self._recording_manager.record_duration - step)
            print(f"\n[Duration: {self._recording_manager.record_duration}s (step: -{step}s)]")
        elif key == ']':
            # Increase record duration with progressive steps
            step = get_duration_step(self._recording_manager.record_duration, decreasing=False)
            self._recording_manager.record_duration += step
            print(f"\n[Duration: {self._recording_manager.record_duration}s (step: +{step}s)]")
        elif key == ',':
            # Decrease burst count with progressive steps
            step = get_frame_step(self._recording_manager.max_frames_in_fits, decreasing=True)
            self._recording_manager.max_frames_in_fits = max(1, self._recording_manager.max_frames_in_fits - step)
            print(f"\n[Frame count: {self._recording_manager.max_frames_in_fits} (step: -{step})]")
        elif key == '.':
            # Increase burst count with progressive steps
            step = get_frame_step(self._recording_manager.max_frames_in_fits, decreasing=False)
            self._recording_manager.max_frames_in_fits += step
            print(f"\n[Frame count: {self._recording_manager.max_frames_in_fits} (step: +{step})]")
        elif key == '\x03':  # Ctrl+C
            print("\nCtrl+C pressed - quitting...")
            return False

        return True

    def start(self) -> None:
        """Start the camera subscriber.

        Connects to the ZMQ publisher and starts receiving frames.
        Press 'q' to stop, 'v' to show window, 'h' to hide window,
        's' to save as FITS, 'j' to save as JPG, 'p' to save as PNG,
        'a' for frame-limited recording, 't' for time-limited recording,
        'm' to cycle burst mode, '[' or ']' to change duration,
        ',' or '.' to change frame count.

        Raises:
            RuntimeError: If cannot connect to the ZMQ publisher.
        """
        # Initialize terminal input handler
        self._input_handler = TerminalInputHandler()
        self._input_handler.start()

        # Initialize ZMQ context and socket
        self._context = zmq.Context()
        self._socket = self._context.socket(zmq.SUB)
        self._socket.connect(self.connect_address)
        self._socket.setsockopt_string(zmq.SUBSCRIBE, "")

        # Set receive timeout to allow checking for quit
        self._socket.setsockopt(zmq.RCVTIMEO, 100)

        print(f"Camera subscriber connected to {self.connect_address}")
        print(
            f"Display resolution: {self.display_width}x{self.display_height}"
        )
        print("Commands: 'q' = quit, 'v' = show window, 'h' = hide window")
        print("          's' = save current frame as FITS")
        print("          'j' = save current frame as JPG")
        print("          'p' = save current frame as PNG")
        print("          '0' = scale mode (default)")
        print("          Crop modes (640x480):")
        print("            '7'='top-left'  '8'='top'      '9'='top-right'")
        print("            '4'='left'      '5'='center'   '6'='right'")
        print("            '1'='bottom-left' '2'='bottom' '3'='bottom-right'")
        print("          'b' = toggle barrel correction, '+'/'-' = adjust k1 (±0.01)")
        print(f"          'a' = frame-limited recording ({self._recording_manager.max_frames_in_fits} frames)")
        print(f"          't' = time-limited recording ({self._recording_manager.record_duration}s)")
        print(f"          'm' = cycle burst mode (individual/sequence)")
        print(f"          '[' / ']' = adjust duration")
        print(f"          ',' / '.' = adjust frame count")
        print("          'x' = toggle 2x display scaling")
        print("Press Ctrl+C to stop")

        self._running = True
        frame_count = 0

        # FPS and bandwidth calculation variables - rolling window of last 5 seconds
        stats_window_seconds = 5.0
        # List of (timestamp, bytes) tuples for rolling window calculation
        stats_window: list[tuple[float, int]] = []
        current_fps = 0.0
        current_mbps = 0.0

        try:
            while self._running:
                # Check for keyboard input
                if self._input_handler is not None:
                    key = self._input_handler.get_key()
                    if key is not None:
                        if not self._handle_keypress(key):
                            break

                try:
                    # Receive multipart message [header, image]
                    parts = self._socket.recv_multipart()
                    if len(parts) != 2:
                        print(f"Warning: Expected 2 parts, got {len(parts)}")
                        continue

                    header_bytes, image_bytes = parts

                    # Parse header
                    try:
                        header = json.loads(header_bytes.decode("utf-8"))
                    except (json.JSONDecodeError, UnicodeDecodeError) as e:
                        print(f"Warning: Failed to parse header: {e}")
                        continue

                    # Extract header fields
                    width = header.get("width", 0)
                    height = header.get("height", 0)
                    fmt = header.get("fmt", "UNKNOWN")
                    seq = header.get("seq", 0)
                    expected_checksum = header.get("checksum", "")
                    is_compressed = header.get("compressed", False)
                    compression = header.get("compression")
                    original_size = header.get("original_size", 0)

                    # Store compressed size for compression ratio calculation
                    compressed_size = len(image_bytes)

                    # Decompress if needed
                    if is_compressed and compression == "zstd":
                        if self._decompressor is None:
                            self._decompressor = zstandard.ZstdDecompressor()
                        try:
                            image_bytes = self._decompressor.decompress(image_bytes)
                        except zstandard.ZstdError as e:
                            print(f"Warning: Failed to decompress frame {seq}: {e}")
                            continue

                    # Calculate compression ratio (100% = no compression)
                    if is_compressed and original_size > 0:
                        compression_ratio = (compressed_size / original_size) * 100
                    else:
                        compression_ratio = 100.0  # No compression applied

                    # Verify checksum (on decompressed data)
                    if expected_checksum:
                        if not self._verify_checksum(image_bytes, expected_checksum):
                            print(f"Warning: Checksum mismatch for seq {seq}")
                            continue

                    # Check for dropped frames
                    if self._last_seq > 0 and seq > self._last_seq + 1:
                        dropped = seq - self._last_seq - 1
                        self._dropped_frames += dropped
                    self._last_seq = seq

                    # Convert YUYV to BGR
                    if fmt == "YUYV":
                        frame = convert_yuyv_to_bgr(image_bytes, width, height)
                    else:
                        # For other formats, try to decode as-is
                        frame = np.frombuffer(image_bytes, dtype=np.uint8)
                        # Assume it's already in a displayable format
                        continue

                    if frame is None:
                        continue

                    # Update last frame time for no-signal detection
                    self._no_signal_manager.update_frame_time()

                    # Store frame data and header for saving
                    self._last_frame_data = image_bytes
                    self._last_frame_header = header

                    # Record frame if burst recording is active
                    if self._recording_manager.is_recording:
                        self._recording_manager.process_frame(
                            header,
                            image_bytes,
                            save_current_frame,
                        )

                    frame_count += 1
                    frame_bytes = len(image_bytes)

                    # FPS and bandwidth calculation using rolling 5-second window
                    now = time.time()
                    # Add current frame to window
                    stats_window.append((now, frame_bytes))
                    # Remove entries older than 5 seconds
                    cutoff_time = now - stats_window_seconds
                    stats_window = [(t, b) for t, b in stats_window if t > cutoff_time]
                    # Calculate FPS and MB/s based on window
                    if len(stats_window) > 1:
                        window_duration = stats_window[-1][0] - stats_window[0][0]
                        if window_duration > 0:
                            current_fps = len(stats_window) / window_duration
                            total_bytes = sum(b for _, b in stats_window)
                            current_mbps = total_bytes / window_duration / 1_000_000

                    # Get ts_after from header for display
                    ts_after = header.get("ts_after", "N/A")

                    dropped_str = f" Dropped:{self._dropped_frames}" if self._dropped_frames > 0 else ""
                    compression_str = f" Comp:{compression_ratio:5.1f}%" if compression_ratio < 100.0 else ""
                    print(
                        f"Frame:{frame_count:7d} "
                        f"Seq:{seq:7d} "
                        f"Res:{width:4d}x{height:4d} "
                        f"Time:{ts_after} "
                        f"; {current_fps:5.1f} fps; "
                        f"{current_mbps:6.3f} MB/s"
                        f"{compression_str}"
                        f"{dropped_str};",
                        end="\r",
                    )

                    # Process frame (display if window visible, call callback)
                    self._process_frame(frame, header)

                except zmq.Again:
                    # Timeout - check for no-signal condition
                    if self._window_visible:
                        no_signal_frame = self._no_signal_manager.check_no_signal()
                        if no_signal_frame is not None:
                            # Ensure window is created before displaying
                            self._show_window()
                            # Display no signal frame
                            display_frame = no_signal_frame.copy()
                            if self._scale_2x:
                                display_frame = cv2.resize(
                                    display_frame,
                                    None,
                                    fx=2.0,
                                    fy=2.0,
                                    interpolation=cv2.INTER_LINEAR,
                                )
                            cv2.imshow(self._window_name, display_frame)
                            cv2.waitKey(1)
                    continue

        except KeyboardInterrupt:
            print("\nStopping subscriber...")
        finally:
            self.stop()

    def stop(self) -> None:
        """Stop the camera subscriber and release resources."""
        self._running = False

        # Stop input handler and restore terminal
        if self._input_handler is not None:
            self._input_handler.stop()
            self._input_handler = None

        # Hide window if visible
        if self._window_visible:
            cv2.destroyWindow(self._window_name)
            self._window_visible = False

        if self._socket is not None:
            self._socket.close()
            self._socket = None

        if self._context is not None:
            self._context.term()
            self._context = None

        print("Subscriber stopped")


def run_subscriber(
    connect_address: str = "tcp://localhost:5555",
    show_window: bool = False,
    display_width: int = 640,
    display_height: int = 480,
    fov_degrees: float = 150.0,
    no_signal_timeout: float = 5.0,
    max_frames_in_fits: int = 10,
    record_duration: int = 10,
    burst_mode: str = "individual",
) -> None:
    """Run the camera subscriber with the specified parameters.

    Args:
        connect_address: ZMQ connect address.
        show_window: Whether to initially display frames in a window.
                     Note: User can toggle with 's' and 'h' keys.
        display_width: Display width for the window (default: 640).
        display_height: Display height for the window (default: 480).
        fov_degrees: Field of view in degrees for barrel correction (default: 150.0).
        no_signal_timeout: Timeout in seconds before showing 'no signal' image.
        max_frames_in_fits: Maximum number of frames to save in a FITS sequence.
        record_duration: Duration in seconds for time-based recording (default: 10).
        burst_mode: Burst mode, 'individual' or 'sequence' (default: 'individual').
    """
    subscriber = CameraSubscriber(
        connect_address=connect_address,
        frame_callback=None,
        display_width=display_width,
        display_height=display_height,
        fov_degrees=fov_degrees,
        no_signal_timeout=no_signal_timeout,
        max_frames_in_fits=max_frames_in_fits,
        record_duration=record_duration,
        burst_mode=burst_mode,
    )

    # Show window initially if requested
    if show_window:
        subscriber._window_visible = True
        # Window will be created on first frame

    subscriber.start()
