segments.py

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any
import re
from typing import Any

from util import *

class Segment:

    def __init__(self, start: float, end: float, data: Any=None):
        
        if start >= end:
            raise ValueError(f"Segment start ({start}) must be less than end ({self._end})")
        self._start = start
        self._end = end
        
        self.data = data

    @property
    def start(self) -> float:
        return self._start

    @start.setter
    def start(self, value: float) -> None:
        if value >= self._end:
            raise ValueError(f"Segment start ({value}) must be less than end ({self._end})")
        self._start = value

    @property
    def end(self) -> float:
        return self._end

    @end.setter
    def end(self, value: float) -> None:
        if value <= self._start:
            raise ValueError(f"Segment end ({value}) must be greater than start ({self._start})")
        self._end = value

    def __str__(self) -> str:
        return make_position_range_string(self._start, self._end)

# ---

class Segments:
    """
    Manages a list of Segment instances with start and end positions (plus an optional data field).
    Items are monotonically increasing, non-overlapping.
    """

    _segments: list[Segment]

    def __init__(self, file_path: str, max_value: float) -> None:
        self._file_path = file_path
        self._dirty = False
        self._max_value = max_value

        if file_path:
            self._segments = Segments.make_list_from_json(self.json_path, max_value) or []
        else:
            self._segments = []

    @property
    def segments(self) -> list[Segment]:
        return self._segments
    
    @property
    def dirty(self) -> bool:
        return self._dirty
    
    @property
    def max_value(self) -> float:
        return self._max_value
    
    @property
    def json_path(self) -> str:
        return Segments.make_json_path_using(self._file_path)

    def insert(self, start: float, end: float, can_merge: bool=False, data: Any=None) -> int:
        """
        Inserts segment using `start` and `end`.
        If can_merge=True, merges with overlapping or adjacent segments.
        Returns index of inserted/merged segment, or -1 if not legal
        """
        
        end = min(end, self._max_value)
        if start >= end:
            return -1
        
        if can_merge:
            return self._insert_with_merge(start, end, data)
        else:
            insertion_point = self.can_insert(start, end)
            if insertion_point == -1:
                return -1
            self._segments.insert(insertion_point, Segment(start, end, data))
            self._dirty = True
            return insertion_point
    
    def _insert_with_merge(self, start: float, end: float, data: Any=None) -> int:
        """
        Insert segment, merging with any overlapping or adjacent segments.
        Returns index of merged segment.
        """
        # Find all segments that overlap or touch the new segment
        merge_indices = []
        new_start = start
        new_end = end
        
        for i, seg in enumerate(self._segments):
            # Check if segments overlap or touch
            if not (seg.end < start or seg.start > end):
                # Segments overlap or touch
                merge_indices.append(i)
                new_start = min(new_start, seg.start)
                new_end = max(new_end, seg.end)
        
        # Remove merged segments (in reverse order to preserve indices)
        for i in reversed(merge_indices):
            del self._segments[i]
        
        # Find insertion point for the merged segment
        insertion_point = self.can_insert(new_start, new_end)
        if insertion_point == -1:
            # This shouldn't happen since we removed overlapping segments
            insertion_point = 0
        
        # Insert the merged segment
        self._segments.insert(insertion_point, Segment(new_start, new_end, data))
        self._dirty = True
        return insertion_point

    def delete_at(self, index: int) -> bool:
        """
        Deletes segment at given index
        Returns False if index is invalid
        """
        if index < 0 or index >= len(self._segments):
            return False
        del self._segments[index]
        self._dirty = True
        return True

    def index_of(self, position: float) -> int:
        """
        Returns index of segment which contains position (inclusive), or -1
        """
        return Segments.segment_index_of(position, self.segments)
    
    def can_insert(self, start: float, end: float) -> int:
        """ Returns insertion point or -1"""

        if start >= end:
            return -1

        low, high = 0, len(self._segments) # binary search
        while low < high:
            mid = (low + high) // 2
            if self._segments[mid].start < start:
                low = mid + 1
            else:
                high = mid

        # Check for overlap with previous segment
        if low > 0 and self._segments[low - 1].end > start:
            return -1

        # Check for overlap with next segment
        if low < len(self._segments) and self._segments[low].start < end:
            return -1
        
        return low

    @staticmethod
    def segment_index_of(position: float, segments: list[Segment]) -> int:
        """
        Returns index of segment which contains position (inclusive), or -1
        """
        low, high = 0, len(segments) # Binary search
        while low < high:
            mid = (low + high) // 2
            seg = segments[mid]
            if position < seg.start:
                high = mid
            elif position > seg.end:
                low = mid + 1
            else:
                return mid # hit
        return -1
    
    @staticmethod
    def where(position: float, segments: list[Segment]) -> tuple[str, int | None]:
        """
        Returns where `position` is in relation to `segments`, described by a tuple:
            "inside", index - position is "inside" segment of the given index
            "before", index - position is "before" segment of the given index
            "before_first", None - position is before the first segment
            "after_last", None - position is after the last segment
            "none", None  - there are no segments
        """
        if not segments:
            return "none", None # meh
        
        if position < segments[0].start:
            return "before_first", None
        
        if position > segments[-1].end:
            return "after_last", None
        
        low, high = 0, len(segments)
        while low < high:
            mid = (low + high) // 2
            seg = segments[mid]
            if position < seg.start:
                high = mid
            elif position > seg.end:
                low = mid + 1
            else:
                return "inside", mid # position is inside segment (inclusive)
        
        # position is between two segments, return the higher indexed segment
        return "before", low

    def save_json(self) -> bool:
        """
        Serializes segments list to JSON file. 
        
        Disregards any data element from the items in the list.
        
        Format: [[start1, end1], [start2, end2], ...]
        
        If no segments, deletes pre-existing file if any.
        
        Returns True for success
        """        
        if not self._dirty: 
            return True # nothing to do

        if not self.segments:
            did_delete = self.delete_json()
            if did_delete:
                self._dirty = False
            return did_delete
        
        try:
            data = [[seg.start, seg.end] for seg in self._segments]
            with open(self.json_path, 'w') as f:
                json.dump(data, f)
        except Exception as e:
            print(f"* {e}")
            return False

        self._dirty = False
        return True
    
    def delete_json(self) -> bool:
        """ 
        Returns False if error deleting.
        """
        return Segments.delete_json_at_path(self.json_path)

    def invert(self) -> None:
        """ Inverts segments list """
        self._segments = Segments.make_inverse(self._segments, self._max_value)
    
    def make_tuples_list(self) -> list[tuple[float, float]]:
        return Segments.make_tuple_twos(self._segments)

    # ---

    @staticmethod
    def make_tuple_twos(segment_list: list[Segment]) -> list[tuple[float, float]]:
        """ Returns list of (start, end) tuples from Segment list """
        return [(seg.start, seg.end) for seg in segment_list]

    @staticmethod
    def make_inverse(segment_list: list[Segment], max_value: float) -> list[Segment]:
        """
        Returns a new list of Segment instances representing the "inverse" (gaps) of the input segment_list.
        The inverse consists of the gaps between segments, from 0 to the first segment,
        between consecutive segments, and from the last segment to duration.

        Ignores (loses) 'data' component.
        """
        if not segment_list:
            return [Segment(0, max_value)]
        if len(segment_list) == 1 and segment_list[0].start == 0.0 and segment_list[0].end == max_value:
            return []
        
        result: list[Segment] = []
        
        # Gap from start (0) to first segment
        first_start = segment_list[0].start
        if first_start > 0:
            result.append(Segment(0, first_start))
        
        # Gaps between consecutive segments
        for i in range(len(segment_list) - 1):
            current_end = segment_list[i].end
            next_start = segment_list[i + 1].start
            if current_end < next_start:
                result.append(Segment(current_end, next_start))
        
        # Gap from last segment to duration
        last_end = segment_list[-1].end
        if last_end < max_value:
            result.append(Segment(last_end, max_value))
        
        return result

    @staticmethod
    def make_json_path_using(path: str) -> str:
        if not path:
            return ""
        p = Path(path)
        p = p.with_stem(p.stem + "_segments")
        p = p.with_suffix(".json")
        return str(p)

    @staticmethod
    def delete_json_at_path(path: str) -> bool:
        """ Returns False if error deleting """
        if os.path.exists(path):
            try:
                os.unlink(path)
            except Exception as e:
                print(f"* {e}")
                return False
        return True

    # ---

    @staticmethod
    def make_list_from_json(file_path: str, max_value: float) -> list[Segment] | None:
        """
        Loads segment list from JSON file whose format looks like: [[start1, end1], [start2, end2], ...]
        Returns None on error.
        """
        if not os.path.exists(file_path):
            return None
        
        try:
            with open(file_path, 'r') as f:
                data = json.load(f)        
        except Exception as e:
            print(f"* {e} {file_path}")
            return None

        temp_segments = Segments("", max_value)
        for item in data:
            if isinstance(item, (list, tuple)) and len(item) == 2:
                a = float(item[0])
                b = float(item[1])
                if a == b:
                    continue # not legal
                if a > b:
                    b, a = a, b
                if temp_segments._segments and a < temp_segments._segments[-1].end:
                    continue
                temp_segments.insert(a, b, can_merge=False)
        return temp_segments.segments

    @staticmethod
    def load_from_srt_file(path: str) -> Segments | None:
        """
        Makes Segments instance from srt (subtitles) file.
        Returns None on error.
        """
        try:
            with open(path, 'r') as f:
                content = f.read()
        except Exception as e:
            print(f"* {e}")
            return None
        
        temp_segments = Segments("", float("infinity"))
        
        # SRT timing format: HH:MM:SS,mmm --> HH:MM:SS,mmm
        # Text follows on subsequent lines until blank line
        pattern = re.compile(
            r'(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*\n(.*?)(?=\n\n|\n$|\Z)',
            re.DOTALL
        )
        
        for match in pattern.finditer(content):
            start_h, start_m, start_s, start_ms = int(match.group(1)), int(match.group(2)), int(match.group(3)), int(match.group(4))
            end_h, end_m, end_s, end_ms = int(match.group(5)), int(match.group(6)), int(match.group(7)), int(match.group(8))
            text = match.group(9).strip()
            
            start_sec = start_h * 3600 + start_m * 60 + start_s + start_ms / 1000.0
            end_sec = end_h * 3600 + end_m * 60 + end_s + end_ms / 1000.0
            
            temp_segments.insert(start_sec, end_sec, can_merge=False, data=text)
        
        temp_segments._dirty = False
        return temp_segments