src/ui/sessions.py

"""
Interactive selection session classes.

This module provides session classes for different types of region selection:
- ScorebugSelectionSession: For selecting the main scorebug region
- PlayClockSelectionSession: For selecting the play clock within scorebug
- TimeoutSelectionSession: For selecting timeout indicator regions
"""

from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np

from .models import BBox, SelectionState, SelectionViewConfig
from .selector import KeyHandler, RegionSelector


def _extract_sample_frames_for_selection(video_path: str, start_time: float, num_frames: int = 5, interval: float = 2.0) -> List[Tuple[float, np.ndarray[Any, Any]]]:
    """Extract sample frames for interactive selection.

    Note: Import is inside function to avoid circular imports.
    The video.frame_extractor module imports from ui, so importing
    it at module level would create a circular dependency.
    """
    # pylint: disable=import-outside-toplevel
    # Import must be inside function to avoid circular dependency:
    # ui.sessions -> video.frame_extractor -> ui (circular)
    from video.frame_extractor import extract_sample_frames

    return extract_sample_frames(video_path, start_time, num_frames, interval)


class InteractiveSelectionSession:
    """
    Base class for interactive region selection sessions.

    Handles the common event loop, key bindings, and display logic.
    Subclasses customize the rendering and available key handlers.
    """

    def __init__(
        self,
        window_name: str,
        frames: List[Tuple[float, np.ndarray[Any, Any]]],
        view_config: Optional[SelectionViewConfig] = None,
        video_path: Optional[str] = None,
    ):
        """Initialize the selection session."""
        self.window_name = window_name
        self.view_config = view_config or SelectionViewConfig()
        self.state = SelectionState(frames=list(frames), video_path=video_path)
        self.selector = RegionSelector(window_name, mode="two_click")

        # Key handlers map: key code -> handler function
        self._key_handlers: Dict[int, KeyHandler] = {}
        self._register_default_handlers()

    def _register_default_handlers(self) -> None:
        """Register the default key handlers."""
        self._key_handlers[ord("q")] = self._handle_quit
        self._key_handlers[ord("r")] = self._handle_reset

    def register_handler(self, key: int, handler: KeyHandler) -> None:
        """Register a custom key handler."""
        self._key_handlers[key] = handler

    # -------------------------------------------------------------------------
    # Key Handlers (small helpers with single-line docstrings)
    # -------------------------------------------------------------------------

    def _handle_quit(self, state: SelectionState, _selector: RegionSelector) -> None:
        """Handle 'q' key - quit selection without saving."""
        state.should_quit = True

    def _handle_reset(self, _state: SelectionState, selector: RegionSelector) -> None:
        """Handle 'r' key - reset current selection."""
        selector.reset()

    def _handle_next_frame(self, state: SelectionState, selector: RegionSelector) -> None:
        """Handle 'n' key - advance to next frame."""
        if state.frame_idx < len(state.frames) - 1:
            state.frame_idx += 1
            selector.reset()

    def _handle_prev_frame(self, state: SelectionState, selector: RegionSelector) -> None:
        """Handle 'p' key - go to previous frame."""
        if state.frame_idx > 0:
            state.frame_idx -= 1
            selector.reset()

    def _handle_skip_30s(self, state: SelectionState, selector: RegionSelector) -> None:
        """Handle 's' key - skip forward 30 seconds."""
        if not state.video_path:
            return
        last_ts = state.frames[-1][0]
        new_start = last_ts + 30
        print(f"  Skipping to {new_start:.1f}s...")
        new_frames = _extract_sample_frames_for_selection(state.video_path, new_start, num_frames=5, interval=2.0)
        if new_frames:
            state.frames = new_frames
            state.frame_idx = 0
            selector.reset()
            print(f"  Loaded {len(new_frames)} new frames starting at {new_start:.1f}s")
        else:
            print(f"  Could not load frames at {new_start:.1f}s (end of video?)")

    def _handle_skip_5min(self, state: SelectionState, selector: RegionSelector) -> None:
        """Handle 'S' key - skip forward 5 minutes."""
        if not state.video_path:
            return
        last_ts = state.frames[-1][0]
        new_start = last_ts + 300
        print(f"  Skipping to {new_start:.1f}s ({new_start / 60:.1f} min)...")
        new_frames = _extract_sample_frames_for_selection(state.video_path, new_start, num_frames=5, interval=2.0)
        if new_frames:
            state.frames = new_frames
            state.frame_idx = 0
            selector.reset()
            print(f"  Loaded {len(new_frames)} new frames starting at {new_start:.1f}s")
        else:
            print(f"  Could not load frames at {new_start:.1f}s (end of video?)")

    def _handle_confirm(self, state: SelectionState, selector: RegionSelector) -> None:
        """Handle Enter key - confirm selection."""
        if selector.selection_complete:
            state.should_confirm = True

    # -------------------------------------------------------------------------
    # Display Methods (to be overridden by subclasses)
    # -------------------------------------------------------------------------

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Render overlays on the display frame. Override in subclasses."""
        return display_frame

    def _get_display_frame(self) -> np.ndarray[Any, Any]:
        """Get the frame to display. Override in subclasses for custom views."""
        return self.state.frame.copy()

    def _draw_selection_overlay(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Draw the selection points and rectangles on the frame."""
        # Draw clicked points
        for i, point in enumerate(self.selector.points):
            color = (0, 255, 0) if i == 0 else (0, 0, 255)
            cv2.circle(display_frame, point, 5, color, -1)

        # Draw preview rectangle if we have first point and mouse position
        if len(self.selector.points) == 1 and self.selector.current_point:
            cv2.rectangle(display_frame, self.selector.points[0], self.selector.current_point, (0, 255, 255), 2)

        # Draw final rectangle if selection complete
        if self.selector.selection_complete:
            bbox = self.selector.get_bbox()
            if bbox:
                cv2.rectangle(display_frame, (bbox.x, bbox.y), (bbox.x2, bbox.y2), (0, 255, 0), 3)

        return display_frame

    # -------------------------------------------------------------------------
    # Main Event Loop
    # -------------------------------------------------------------------------

    def run(self) -> Optional[BBox]:
        """Run the interactive selection loop. Returns selected BBox or None."""
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(self.window_name, self.view_config.window_width, self.view_config.window_height)
        cv2.setMouseCallback(self.window_name, self.selector.mouse_callback)

        while True:
            # Get and render the display frame
            display_frame = self._get_display_frame()
            display_frame = self._draw_selection_overlay(display_frame)
            display_frame = self._render_frame(display_frame)

            cv2.imshow(self.window_name, display_frame)
            key = cv2.waitKey(30) & 0xFF

            # Handle key press
            if key in self._key_handlers:
                self._key_handlers[key](self.state, self.selector)

            # Check exit conditions
            if self.state.should_quit:
                cv2.destroyAllWindows()
                return None

            if self.state.should_confirm:
                bbox = self.selector.get_bbox()
                cv2.destroyAllWindows()
                return bbox

        return None


class ScorebugSelectionSession(InteractiveSelectionSession):
    """Interactive session for selecting the scorebug region."""

    def __init__(self, frames: List[Tuple[float, np.ndarray[Any, Any]]], video_path: Optional[str] = None):
        """Initialize scorebug selection session."""
        super().__init__(
            window_name="Select Scorebug Region",
            frames=frames,
            view_config=SelectionViewConfig(window_width=1280, window_height=720),
            video_path=video_path,
        )
        # Register additional handlers for scorebug selection
        self._key_handlers[ord("n")] = self._handle_next_frame
        self._key_handlers[ord("p")] = self._handle_prev_frame
        self._key_handlers[ord("s")] = self._handle_skip_30s
        self._key_handlers[ord("S")] = self._handle_skip_5min
        self._key_handlers[13] = self._handle_confirm  # Enter key

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Render scorebug-specific overlays and instructions."""
        # Draw point labels
        for i, point in enumerate(self.selector.points):
            color = (0, 255, 0) if i == 0 else (0, 0, 255)
            label = "Top-Left" if i == 0 else "Bottom-Right"
            cv2.putText(display_frame, label, (point[0] + 10, point[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        # Frame info overlay
        frame_info = f"Frame {self.state.frame_idx + 1}/{len(self.state.frames)} @ {self.state.timestamp:.1f}s | n=next, p=prev, s=skip 30s, S=skip 5min"
        cv2.putText(display_frame, frame_info, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

        # Step instructions
        if len(self.selector.points) == 0:
            instruction = "Step 1: Click TOP-LEFT corner of scorebug"
            color = (0, 255, 255)
        elif len(self.selector.points) == 1:
            instruction = "Step 2: Click BOTTOM-RIGHT corner of scorebug"
            color = (0, 255, 255)
        else:
            instruction = "Press ENTER to confirm, 'r' to reset"
            color = (0, 255, 0)

        cv2.putText(display_frame, instruction, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        return display_frame


class PlayClockSelectionSession(InteractiveSelectionSession):
    """Interactive session for selecting the play clock region within a scorebug."""

    def __init__(self, frame: np.ndarray[Any, Any], scorebug_bbox: BBox, scale_factor: int = 3):
        """Initialize play clock selection session."""
        self.scorebug_bbox = scorebug_bbox
        self.scale_factor = scale_factor

        # Extract and scale the scorebug region
        sb = scorebug_bbox
        self.scorebug_region = frame[sb.y : sb.y2, sb.x : sb.x2].copy()
        self.scaled_region = cv2.resize(
            self.scorebug_region,
            (sb.width * scale_factor, sb.height * scale_factor),
            interpolation=cv2.INTER_LINEAR,
        )

        # Create a single-frame list for the base class
        super().__init__(
            window_name="Select Play Clock Region",
            frames=[(0.0, self.scaled_region)],
            view_config=SelectionViewConfig(
                window_width=sb.width * scale_factor,
                window_height=sb.height * scale_factor + 100,
                scale_factor=scale_factor,
            ),
        )
        self._key_handlers[13] = self._handle_confirm  # Enter key

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Render play clock selection overlays."""
        # Draw grid for reference
        grid_spacing = 50 * self.scale_factor
        for gx in range(0, display_frame.shape[1], grid_spacing):
            cv2.line(display_frame, (gx, 0), (gx, display_frame.shape[0]), (100, 100, 100), 1)
        for gy in range(0, display_frame.shape[0], grid_spacing):
            cv2.line(display_frame, (0, gy), (display_frame.shape[1], gy), (100, 100, 100), 1)

        # Instructions
        cv2.putText(display_frame, "Select Play Clock | r=reset, q=quit, ENTER=confirm", (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        return display_frame

    def get_unscaled_bbox(self) -> Optional[BBox]:
        """Get the selection converted back to original scorebug coordinates."""
        bbox = self.selector.get_bbox()
        if bbox is None:
            return None
        return bbox.unscaled(self.scale_factor)


class TimeoutSelectionSession(InteractiveSelectionSession):
    """Interactive session for selecting timeout indicator regions."""

    def __init__(self, frame: np.ndarray[Any, Any], scorebug_bbox: BBox, team: str, scale_factor: int = 3, padding: int = 50):
        """Initialize timeout selection session."""
        self.scorebug_bbox = scorebug_bbox
        self.scale_factor = scale_factor
        self.padding = padding
        self.team = team
        self.team_color = (0, 0, 255) if team == "home" else (255, 165, 0)

        # Extract padded region around scorebug
        sb = scorebug_bbox
        self.pad_x1 = max(0, sb.x - padding)
        self.pad_y1 = max(0, sb.y - padding)
        self.pad_x2 = min(frame.shape[1], sb.x2 + padding)
        self.pad_y2 = min(frame.shape[0], sb.y2 + padding)

        # Scorebug offset within padded region
        self.scorebug_offset = BBox(x=sb.x - self.pad_x1, y=sb.y - self.pad_y1, width=sb.width, height=sb.height)

        # Extract and scale
        padded_region = frame[self.pad_y1 : self.pad_y2, self.pad_x1 : self.pad_x2].copy()
        padded_w, padded_h = padded_region.shape[1], padded_region.shape[0]
        self.scaled_region = cv2.resize(padded_region, (padded_w * scale_factor, padded_h * scale_factor), interpolation=cv2.INTER_LINEAR)

        super().__init__(
            window_name=f"Select {team.upper()} Team Timeout Region",
            frames=[(0.0, self.scaled_region)],
            view_config=SelectionViewConfig(
                window_width=padded_w * scale_factor,
                window_height=padded_h * scale_factor + 50,
                scale_factor=scale_factor,
                padding=padding,
            ),
        )
        self._key_handlers[13] = self._handle_confirm  # Enter key

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Render timeout selection overlays."""
        # Draw scorebug boundary for reference
        sb_scaled = self.scorebug_offset.scaled(self.scale_factor)
        cv2.rectangle(display_frame, (sb_scaled.x, sb_scaled.y), (sb_scaled.x2, sb_scaled.y2), (128, 128, 128), 1)

        # Update point colors to team color
        if self.selector.selection_complete:
            bbox = self.selector.get_bbox()
            if bbox:
                cv2.rectangle(display_frame, (bbox.x, bbox.y), (bbox.x2, bbox.y2), self.team_color, 3)

        # Instructions
        cv2.putText(
            display_frame,
            f"Select {self.team.upper()} Timeout | r=reset, q=quit, ENTER=confirm",
            (10, 25),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 255, 255),
            1,
        )

        return display_frame

    def get_absolute_bbox(self) -> Optional[BBox]:
        """Get the selection converted to absolute frame coordinates."""
        bbox = self.selector.get_bbox()
        if bbox is None:
            return None
        # Unscale and add padding offset
        unscaled = bbox.unscaled(self.scale_factor)
        return unscaled.offset(self.pad_x1, self.pad_y1)


class FlagRegionSelectionSession(InteractiveSelectionSession):
    """Interactive session for selecting the FLAG indicator region (where '1st & 10' / 'FLAG' appears)."""

    def __init__(self, frame: np.ndarray[Any, Any], scorebug_bbox: BBox, scale_factor: int = 3, padding: int = 100):
        """
        Initialize FLAG region selection session.

        Args:
            frame: Full video frame (BGR)
            scorebug_bbox: Bounding box of the scorebug in the frame
            scale_factor: How much to scale up for easier selection
            padding: Pixels to pad around the scorebug
        """
        self.scorebug_bbox = scorebug_bbox
        self.scale_factor = scale_factor
        self.padding = padding

        # Extract padded region around scorebug (FLAG may be outside scorebug box)
        sb = scorebug_bbox
        self.pad_x1 = max(0, sb.x - padding)
        self.pad_y1 = max(0, sb.y - padding)
        self.pad_x2 = min(frame.shape[1], sb.x2 + padding)
        self.pad_y2 = min(frame.shape[0], sb.y2 + padding)

        # Scorebug offset within padded region
        self.scorebug_offset = BBox(x=sb.x - self.pad_x1, y=sb.y - self.pad_y1, width=sb.width, height=sb.height)

        # Extract and scale
        padded_region = frame[self.pad_y1 : self.pad_y2, self.pad_x1 : self.pad_x2].copy()
        padded_w, padded_h = padded_region.shape[1], padded_region.shape[0]
        self.scaled_region = cv2.resize(padded_region, (padded_w * scale_factor, padded_h * scale_factor), interpolation=cv2.INTER_LINEAR)

        super().__init__(
            window_name="Select FLAG Region (1st & 10 location)",
            frames=[(0.0, self.scaled_region)],
            view_config=SelectionViewConfig(
                window_width=padded_w * scale_factor,
                window_height=padded_h * scale_factor + 50,
                scale_factor=scale_factor,
                padding=padding,
            ),
        )
        self._key_handlers[13] = self._handle_confirm  # Enter key

    def _render_frame(self, display_frame: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        """Render FLAG region selection overlays."""
        # Draw scorebug boundary for reference
        sb_scaled = self.scorebug_offset.scaled(self.scale_factor)
        cv2.rectangle(display_frame, (sb_scaled.x, sb_scaled.y), (sb_scaled.x2, sb_scaled.y2), (128, 128, 128), 2)
        cv2.putText(display_frame, "Scorebug", (sb_scaled.x, sb_scaled.y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (128, 128, 128), 1)

        # Instructions
        cv2.putText(
            display_frame,
            "Select FLAG region (where '1st & 10' appears) | r=reset, q=quit, ENTER=confirm",
            (10, 25),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 255, 255),
            1,
        )

        return display_frame

    def get_bbox_relative_to_scorebug(self) -> Optional[BBox]:
        """Get the selection as offset from scorebug (like play clock region)."""
        bbox = self.selector.get_bbox()
        if bbox is None:
            return None

        # Unscale first
        unscaled = bbox.unscaled(self.scale_factor)

        # Convert to scorebug-relative coordinates
        # The selection is in padded region coords, scorebug_offset is where scorebug starts in padded region
        x_offset = unscaled.x - self.scorebug_offset.x
        y_offset = unscaled.y - self.scorebug_offset.y

        return BBox(x=x_offset, y=y_offset, width=unscaled.width, height=unscaled.height)