timeline.py

"""
PySide6-based Timeline widget.
Migrated from timeline.py (CustomTkinter-based) without Scaling class logic.
"""

from __future__ import annotations
from typing import assert_never

import numpy as np
from PySide6.QtWidgets import QWidget, QApplication
from PySide6.QtCore import Qt, Signal, QRectF, QPointF, QEvent
from PySide6.QtGui import QPainter, QColor, QPen, QFont, QPolygonF, QBrush, QCursor

from app_types import DeleteSelectionAction, SeekAction, SelectionBeginAction, SelectionCompleteAction, TimelineMouseAction
from constants import *
from sound import Sound
from segments import Segment, Segments
from util import *


class Timeline(QWidget):
    """
    Timeline widget displaying waveform and edit info.
    Does custom painting via paintEvent.
    
    Drawn elements:
    - Waveform visualization
    - Playhead position
    - Segment markers
    - Start/end position markers
    - Start-of-segment temporary marker
    """
    
    HEIGHT = 200
    BARS_PER_SECOND = 200
    BAR_WIDTH_BASE = 1.0
    ZOOM_MIN = 0.20
    ZOOM_MAX = 1.0

    # Signal emitted when the canvas is clicked
    canvas_clicked = Signal(object) # TimelineMouseAction

    def __init__(self, parent=None):
        super().__init__(parent)
        
        self.setFixedHeight(self.HEIGHT)
        self.setMinimumWidth(400)
        self._font = QFont("Arial", 10)
        
        # Initialize properties
        self._zoom = 1.0
        self._bar_width = self.BAR_WIDTH_BASE * self._zoom
        self._pixels_per_second = self.BARS_PER_SECOND * self._bar_width
        self._bars: np.ndarray = np.array([])  # numpy array for fast slicing
        self._samples_per_bar: float = 0.0  # Float to avoid cumulative drift
        self._duration = 0.0
        self._min_position = 0.0
        self._max_position = 0.0
        self._scroll_position = 0.0  # horizontal scroll offset in seconds
        self._playhead_position = 0.0
        self._segments: Segments | None = None
        self._selection_start = -1.0
        self._subtitle_text = ""
                
        # Install event filter to catch Control key events globally
        # Defer installation to ensure QApplication instance exists
        self._event_filter_installed = False
        
        # Tracks control and shift key state as source of truth
        self._control_pressed = False
        self._shift_pressed = False

        self.setMouseTracking(True)
        self._mouse_prompt = ""
        
        # Click data storage (for emitting with signal)
        self._click_data = {}

    def init(self, sound: Sound | None, segments_reference: Segments | None) -> None:
        """
        Sets the waveform's sound and segments reference.
        Computes the bars array for waveform visualization.
        """
        self._duration = sound.duration if sound else 0.0
        self._segments = segments_reference
        self._selection_start = -1.0
        self._playhead_position = 0.0
        self._scroll_position = self._calculate_scroll_for_playhead(self._scroll_position)
        
        if sound:
            self._min_position = 0.0
            self._max_position = sound.duration
            
            # Compute bars array for waveform visualization
            # Store samples_per_bar as float to avoid cumulative drift from integer truncation
            num_bars = int(sound.duration * self.BARS_PER_SECOND)
            n_samples = sound.data.shape[-1]  # Works for both mono and stereo
            self._samples_per_bar = n_samples / num_bars  # Float, e.g., 176.4
            
            # Convert to mono if stereo
            if sound.data.ndim == 1:
                mono_data = sound.data
            else:
                # Stereo: average channels for visualization
                mono_data = np.mean(sound.data, axis=0)
            
            # Compute bar heights using vectorized reduceat for performance
            # Each bar i represents time [i/BARS_PER_SECOND, (i+1)/BARS_PER_SECOND]
            abs_data = np.abs(mono_data)
            
            # Generate bar boundaries using float samples_per_bar
            bar_indices = (np.arange(num_bars + 1) * self._samples_per_bar).astype(int)
            
            # Use reduceat for vectorized sum computation per segment
            sums = np.add.reduceat(abs_data, bar_indices[:-1])
            
            # Compute sample counts per bar (handle variable lengths due to float rounding)
            counts = np.diff(bar_indices)
            counts = np.maximum(counts, 1)  # Avoid division by zero
            
            self._bars = sums / counts
            self._bars = np.clip(self._bars, 0.0, 1.0)
            
            # Logarithmic scaling for better visual representation
            FACTOR = 33
            self._bars = np.log1p(self._bars * FACTOR) / np.log1p(100)
        else:
            self._bars = np.array([])
            self._min_position = 0.0
            self._max_position = 0.0
        
        self.update()

    # ---------------
    # Getters/setters

    @property
    def zoom(self) -> float:
        """Controls zoom level; more bars per second = more zoom"""
        return self._zoom

    @zoom.setter
    def zoom(self, value: float) -> None:
        value = max(self.ZOOM_MIN, min(self.ZOOM_MAX, value))
        self._zoom = value
        self._bar_width = self.BAR_WIDTH_BASE * self._zoom
        self._pixels_per_second = self.BARS_PER_SECOND * self._bar_width
        self.update()

    @property
    def min_position(self) -> float:
        return self._min_position

    @min_position.setter
    def min_position(self, value: float) -> None:
        self._min_position = value
        self.update()

    @property
    def max_position(self) -> float:
        return self._max_position

    @max_position.setter
    def max_position(self, value: float) -> None:
        self._max_position = value
        self.update()

    @property
    def scroll_position(self) -> float:
        return self._scroll_position  # read-only

    @property
    def playhead_position(self) -> float:
        return self._playhead_position

    @playhead_position.setter
    def playhead_position(self, position: float) -> None:
        if position != self._playhead_position:
            self._playhead_position = position
            self._scroll_position = self._calculate_scroll_for_playhead(position)
            self.update()

    @property
    def segment_start(self) -> float:
        """ Segment start marker """
        return self._selection_start

    @segment_start.setter
    def segment_start(self, value: float) -> None:
        self._selection_start = value
        self.update()

    @property
    def mouse_prompt(self) -> str:
        return self._mouse_prompt # read-only
    
    def update_mouse_prompt(self) -> None:
        mouse_action = self._get_mouse_action()
        if not mouse_action:
            s = ""
        else:
            s = ""
            match mouse_action:
                case SeekAction(position):
                    s = f"Click to seek {make_position_string(position)}"
                case SelectionBeginAction(position):
                    s = f"Click to start new selection"
                case SelectionCompleteAction(end, start):
                    s = f"Click to complete selection (Press ESC to cancel)"
                case DeleteSelectionAction(index):
                    s = f"Click to delete selection"
                case _:
                    s = ""

        self._mouse_prompt = s

    def _calculate_scroll_for_playhead(self, position: float) -> float:
        """Calculate the appropriate scroll position to keep playhead visible.
        
        When position is at 0 (beginning), allows 10px extra space on the left
        for mouse selection. Similarly, when position is at the sound's full
        duration, allows 10px extra space on the right.
        """
        
        # Guard clauses
        if not self._duration:
            return 0.0
        canvas_width = self.width()
        if canvas_width <= 0:
            return self._scroll_position  # Return current if not sized yet
        
        view_width_seconds = canvas_width / self._pixels_per_second
        
        # Extra padding in seconds (10px worth) for selection space at edges
        padding_seconds = 10.0 / self._pixels_per_second
        
        # Calculate scroll to center the playhead
        centered_scroll = position - view_width_seconds / 2
        
        # Clamp to allow 10px padding on both ends
        min_scroll = -padding_seconds  # 10px space on left when position = 0
        max_scroll = self._duration - view_width_seconds + padding_seconds  # 10px space on right
        
        return max(min_scroll, min(max_scroll, centered_scroll))

    def set_segments(self, segments: Segments) -> None:
        """Set the segments reference"""
        self._segments = segments
        self._selection_start = -1.0
        self.update()

    def set_subtitle_text(self, text: str) -> None:
        """Set the subtitle text to display"""
        if text != self._subtitle_text:
            self._subtitle_text = text
            self.update()

    def local_x_to_position(self, local_x: int) -> float:
        """ Convert local x to time position (unclamped) """
        scroll_pixels = self._scroll_position * self._pixels_per_second
        position = (local_x + scroll_pixels) / self._pixels_per_second
        return position

    def _mouse_x_to_position(self) -> float:
        """ Returns mouse-x's time position, or -1 if mouse not over Timeline """
        mouse_coords = self._get_mouse_coords()
        return self.local_x_to_position(mouse_coords[0]) if mouse_coords else -1

    def _get_mouse_coords(self) -> tuple[int, int] | None:
        """ Gets current mouse position in local coordinates, or None if not over widget. """
        if not self.underMouse():
            return None
        pos = self.mapFromGlobal(QCursor.pos())
        mouse_x, mouse_y = pos.x(), pos.y()
        # Check bounds (underMouse can be true at edges)
        if 0 <= mouse_x < self.width() and 0 <= mouse_y < self.height():
            return (mouse_x, mouse_y)
        return None

    def get_mouse_over_segment_index(self) -> int:
        """ Returns the index of the segment which the mouse is hovering over, or -1 """
        
        if not self._segments:
            return -1

        mouse_pos = self._get_mouse_coords()
        if not mouse_pos:
            return -1
        
        mouse_x, _ = mouse_pos
        position = self.local_x_to_position(mouse_x)
                        
        for i, segment in enumerate(self._segments.segments):
            if segment.start <= position <= segment.end:
                return i
        return -1

    def _get_mouse_action(self) -> TimelineMouseAction | None:
        """
        Answers the question:
        What should happen when the mouse is clicked at it current position
        
        New behavior:
        - First click: sets start point (SelectionBeginAction)
        - Second click: completes selection (SelectionCompleteAction)
        - End point can be before or after start point (normalized later)
        """
        if not self._duration:
            return None
        
        coords = self._get_mouse_coords()
        if not coords:
            # Mouse is not over timeline
            self._mouse_status = None
            return None
        
        x, y = coords
        position = self.local_x_to_position(x)
        position = clamp(position, 0.0, self._duration)

        if self._control_pressed:
            return SeekAction(position)

        no_selection_marker = (self._selection_start == -1)
        
        segment_index = self.get_mouse_over_segment_index()
        is_delete = segment_index > -1 and no_selection_marker
        if is_delete:
            return DeleteSelectionAction(segment_index)
        
        if no_selection_marker:
            # First click: set start point
            return SelectionBeginAction(begin_position=position, is_silence_detect=False)
        else:
            # Second click: complete selection
            return SelectionCompleteAction(
                complete_position=position,
                begin_position=self._selection_start,
                is_silence_detect=False
            )
        
    @property
    def mouse_position(self) -> tuple[int, int] | None:
        """
        Get local mouse position or None if mouse not over widget
        """
        return self._get_mouse_coords()

    # ----------
    # Draw logic

    def _draw_segments(self, painter: QPainter, mouse_action: TimelineMouseAction | None) -> None:
        """Draws light-colored rects for segments (selections) """
        
        if not self._segments:
            return

        canvas_width = self.width()
        canvas_height = self.height()
        if canvas_width <= 0 or canvas_height <= 0:
            return

        if isinstance(mouse_action, DeleteSelectionAction):
            mouse_over_index = mouse_action.index
        else:
            mouse_over_index = -1

        for i, segment in enumerate(self._segments.segments):
            is_delete_style = (i == mouse_over_index)
            self._draw_segment(painter, segment, is_delete_style=is_delete_style)

        # Draw "X" at mouse position as deletion indicator
        if mouse_over_index > -1:
            if False:
                mouse_coords = self._get_mouse_coords()
                assert mouse_coords
                mouse_x = mouse_coords[0]
                
                # Get the segment being hovered
                segment = self._segments.segments[mouse_over_index]
                scroll_pixels = self._scroll_position * self._pixels_per_second
                x1 = segment.start * self._pixels_per_second - scroll_pixels
                x2 = segment.end * self._pixels_per_second - scroll_pixels
                
                # Clamp x position within segment bounds with 20px margin
                clamped_x = max(x1 + 20, min(mouse_x, x2 - 20))
                
                # Use larger font and center the "X"
                large_font = QFont("Arial", 20, QFont.Weight.Bold)
                painter.setPen(QColor(255, 255, 255))  # White
                painter.setFont(large_font)
                
                # Draw at clamped position, centered
                painter.drawText(int(clamped_x) - 50, 0, 100, 40, Qt.AlignmentFlag.AlignCenter, "X")

        if isinstance(mouse_action, SelectionCompleteAction):
            a, b = mouse_action.begin_position, mouse_action.complete_position
            if a != b:
                if a > b:
                    a, b = b, a
                segment = Segment(a, b)
                self._draw_segment(painter, segment, is_delete_style=True)

    def _draw_segment(self, painter: QPainter, segment: Segment, is_delete_style: bool=False) -> None:

            scroll_pixels = self._scroll_position * self._pixels_per_second
            x1 = segment.start * self._pixels_per_second - scroll_pixels
            x2 = segment.end * self._pixels_per_second - scroll_pixels

            # Skip segments entirely outside the visible area
            if x2 < 0 or x1 > self.width():
                return

            # Clamp to visible area
            x1 = max(0, x1)
            x2 = min(self.width(), x2)

            painter.setPen(Qt.PenStyle.NoPen)

            if is_delete_style:
                # Checkered fill style signifies that clicking on selection will delete it
                painter.setBrush(QBrush(QCOLOR_LIGHT_GRAY, Qt.BrushStyle.DiagCrossPattern))
            else:
                painter.setBrush(QColor(208, 208, 208))  # #d0d0d0
            painter.drawRect(round(x1), 0, round(x2 - x1), self.height())

    def _draw_wave_plus(self, painter: QPainter, mouse_action: TimelineMouseAction | None) -> None:
        """ Draws waveform, vertical bars at second boundaries, and markers. """

        if not self._duration:
            return

        canvas_width = self.width()
        canvas_height = self.height()
        if canvas_width <= 0 or canvas_height <= 0 or len(self._bars) == 0:
            return

        center_y = canvas_height / 2
        scroll_pixels = self._scroll_position * self._pixels_per_second

        # Calculate visible bar range (culling)
        first_bar_index = max(0, int(scroll_pixels / self._bar_width))
        last_bar_index = min(len(self._bars), int((scroll_pixels + canvas_width) / self._bar_width) + 1)
        
        # Draw vertical line at each second interval
        max_seconds = int(self._duration)
        first_visible_second = max(0, int(self._scroll_position))
        last_visible_second = min(max_seconds, int(self._scroll_position + canvas_width / self._pixels_per_second) + 1)
        
        # Create dashed pen for second markers (no Scaling.canvas_scale)
        dash_pen = QPen(QColor(187, 187, 187), 0.5, Qt.PenStyle.DashLine)
        painter.setPen(dash_pen)
        
        for second in range(first_visible_second, last_visible_second + 1):
            x = round(second * self._pixels_per_second - scroll_pixels)
            painter.drawLine(x, 0, x, canvas_height)
            
            # Draw timestamp to the right of every 10-second line
            if second % 10 == 0:
                painter.setPen(QCOLOR_LIGHT_GRAY)
                painter.setFont(self._font)
                painter.drawText(int(x + 5), 15, make_position_string(second, whole_seconds=True))
                painter.setPen(dash_pen)  # Restore dashed pen

        # Draw sound wave bars
        if first_bar_index < last_bar_index:
            visible_bars = self._bars[first_bar_index:last_bar_index]
            
            # Use single pen for all bars
            painter.setPen(QPen(QColor(0, 0, 0), max(1, int(self._bar_width))))
            
            for i, bar_height in enumerate(visible_bars):
                x = round((first_bar_index + i) * self._bar_width - scroll_pixels)
                half_height = bar_height * self.HEIGHT / 2
                painter.drawLine(x, int(center_y - half_height), x, int(center_y + half_height))
        
        # Start marker (green)
        start_x = self._min_position * self._pixels_per_second - scroll_pixels
        if 0 <= start_x <= canvas_width:
            painter.setPen(QPen(QColor(0, 128, 0), 2))
            painter.drawLine(round(start_x), 0, round(start_x), canvas_height)
            painter.setFont(self._font)
            painter.drawText(int(start_x + 5), 15, f"{make_position_string(self._min_position)}")

        # End marker (red)
        end_x = self._max_position * self._pixels_per_second - scroll_pixels
        if 0 <= end_x <= canvas_width:
            painter.setPen(QPen(QColor(255, 0, 0), 2))
            painter.drawLine(round(end_x), 0, round(end_x), canvas_height)
            painter.drawText(QRectF(end_x - 100 - 5, 0, 100, 20), Qt.AlignmentFlag.AlignRight, f"{make_position_string(self._max_position)}")

        mouse_time_position = self._mouse_x_to_position()

        # Selection start marker
        if self._selection_start >= 0:
            self._draw_selection_marker(painter, "start", mouse_time_position)

        # Mouse hover graphic
        if mouse_action:
            match mouse_action:
                case SeekAction(mouse_time_position):
                    # Draw hover-styled playhead
                    self._draw_playhead(painter, position=mouse_time_position, is_hover=True)
                case SelectionBeginAction(begin_position, _):
                    self._draw_selection_marker(painter, "hover_ibar", begin_position)
                case SelectionCompleteAction(complete_position, _, __):
                    self._draw_selection_marker(painter, "hover_end", mouse_time_position)
                case DeleteSelectionAction(index):
                    # Do nothing, delete hover indicator handled elsewhere
                    ...
                case _:
                    assert_never(mouse_action)

        # Playhead (magenta)
        self._draw_playhead(painter, self._playhead_position)

    def _draw_playhead(self, painter: QPainter, position: float, is_hover: bool=False) -> None:

        canvas_width = self.width()
        canvas_height = self.height()
        scroll_pixels = self._scroll_position * self._pixels_per_second

        x = round( position * self._pixels_per_second - scroll_pixels )
        if not (0 <= x <= canvas_width):
            return
        
        color = QColor(255, 128, 255) if is_hover else QColor(255, 0, 255)

        # Draw playhead line
        if is_hover:
            painter.setPen(QPen(color, 2, Qt.PenStyle.DashLine))
        else:
            painter.setPen(QPen(color, 2))
        painter.drawLine(x, 0, x, canvas_height)

        # Draw triangle at base of playhead
        painter.setPen(Qt.PenStyle.NoPen)
        painter.setBrush(QBrush(color))
        triangle = QPolygonF([
            QPointF(x - 10, canvas_height),
            QPointF(x + 10, canvas_height),
            QPointF(x, canvas_height - 10)
        ])
        painter.drawPolygon(triangle)
        
        # Draw position label
        painter.setPen(QPen(color, 1))
        painter.setFont(self._font)
        label_string = make_position_string(position)
        label_y = canvas_height - 25
        if x >= canvas_width - 100:
            label_rect = QRectF(x - 100 - 10, label_y, 100, 25)
            label_align = Qt.AlignmentFlag.AlignRight
        else:
            label_rect = QRectF(x + 10, label_y, 100, 25)
            label_align = Qt.AlignmentFlag.AlignLeft
        painter.drawText(label_rect, label_align, label_string)


    def _draw_overlay(self, painter: QPainter) -> None:
        """Draw overlay elements like subtitle text."""
        
        if self._playhead_position == -1:
            return
        
        canvas_width = self.width()
        canvas_height = self.height()
        if canvas_width <= 0 or canvas_height <= 0:
            return
        
        if self._subtitle_text:
            painter.setPen(QColor(255, 0, 0))
            painter.setFont(self._font)
            # Draw text anchored at bottom-right
            painter.drawText(canvas_width - 5, canvas_height - 5, self._subtitle_text)

    def _draw_selection_marker(
            self, painter: QPainter, style: str, mouse_time_position: float
    ) -> None:
        """
        Draws a segment start marker graphic in one of three styles ("start", "hover_end", "hover_ibar")
        """
        STROKE_WIDTH = 1.5
        CAP_LENGTH = 8

        if style == "start":
            position = self._selection_start
            points_left = (mouse_time_position < self._selection_start)
            points_right = not points_left
            color = QCOLOR_BLACK
        elif style == "hover_end":
            position = mouse_time_position
            points_left = (mouse_time_position > self._selection_start)
            points_right = not points_left
            color = QCOLOR_BLACK
        else: # "hover_ibar":
            position = mouse_time_position
            points_left = True
            points_right = True
            color = QCOLOR_LIGHT_GRAY 

        scroll_pixels = self._scroll_position * self._pixels_per_second
        x = round(position * self._pixels_per_second - scroll_pixels)

        VERTICAL_MARGIN = 3
        y1 = VERTICAL_MARGIN
        y2 = self.height() - VERTICAL_MARGIN

        painter.setPen(QPen(color, STROKE_WIDTH, Qt.PenStyle.DashLine))

        # Vertical line
        painter.drawLine(x, y1, x, y2) 

        # End caps
        if points_left:
            painter.drawLine(x, y1, x - CAP_LENGTH, y1)
            painter.drawLine(x, y2, x - CAP_LENGTH, y2)
        if points_right:
            painter.drawLine(x, y1, x + CAP_LENGTH, y1)
            painter.drawLine(x, y2, x + CAP_LENGTH, y2)

    # ----------------------
    # QWidget event handlers

    def paintEvent(self, event):
        """Paint the timeline with segments, waveform, markers, and overlay."""
        painter = QPainter(self)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)        
        painter.fillRect(self.rect(), QCOLOR_BG)

        # Draw layers in order
        mouse_action = self._get_mouse_action()
        self._draw_segments(painter, mouse_action)
        self._draw_wave_plus(painter, mouse_action)
        self._draw_overlay(painter)

    def mousePressEvent(self, event):
        """Handle mouse click - emit canvas_clicked signal."""
        
        if event.button() != Qt.MouseButton.LeftButton:
            return        

        data = self._get_mouse_action()
        if not data:
            return

        self._click_data = data
        self.canvas_clicked.emit(self._click_data)

    def resizeEvent(self, event) -> None:
        """Recalculate scroll position when widget is resized"""
        super().resizeEvent(event)
        self._scroll_position = self._calculate_scroll_for_playhead(self._playhead_position)
        self.update()

    def showEvent(self, event) -> None:
        """ Install event filter when widget is shown """
        super().showEvent(event)
        if not self._event_filter_installed:
            app = QApplication.instance()
            if app:
                app.installEventFilter(self)
                self._event_filter_installed = True

    def eventFilter(self, obj, event) -> bool:
        """Event filter to catch Control and Shift key events globally"""
        if event.type() == QEvent.Type.KeyPress:
            if event.key() == Qt.Key.Key_Control and not self._control_pressed:
                self._control_pressed = True
                self.update()
            elif event.key() == Qt.Key.Key_Shift and not self._shift_pressed:
                self._shift_pressed = True
                self.update()
        elif event.type() == QEvent.Type.KeyRelease:
            if event.key() == Qt.Key.Key_Control and self._control_pressed:
                self._control_pressed = False
                self.update()
            elif event.key() == Qt.Key.Key_Shift and self._shift_pressed:
                self._shift_pressed = False
                self.update()
        return super().eventFilter(obj, event)

QCOLOR_BG = QColor(240, 240, 240)