subtitle_processor.py

"""
Subtitle Processor - Removes subtitle segments matching video cuts and adjusts timestamps
"""

import re
from pathlib import Path
from typing import List, Tuple, Optional

from profanity_words import PROFANITY_WORDS

try:
    # Newer installs expose helpers for religious/context filtering.
    from profanity_words import get_profanity_words, should_filter_word
except ImportError:  # Older servers may only ship the word set.
    def get_profanity_words(include_religious: bool = False):
        return set(PROFANITY_WORDS)

    def should_filter_word(word: str, context) -> bool:
        return True


def read_subtitle_text(path: Path) -> str:
    """
    Read SRT/VTT text safely across common encodings.

    Customer SRTs are often UTF-16 (BOM starts with 0xFF), while this tool
    historically assumed UTF-8 only.
    """
    data = Path(path).read_bytes()
    if not data:
        return ''

    # BOM-based detection first (most reliable).
    if data.startswith(b'\xff\xfe') or data.startswith(b'\xfe\xff'):
        return data.decode('utf-16')
    if data.startswith(b'\xef\xbb\xbf'):
        return data.decode('utf-8-sig')

    # UTF-16 without BOM often has NUL bytes in the first line.
    if b'\x00' in data[:64]:
        for enc in ('utf-16', 'utf-16-le', 'utf-16-be'):
            try:
                return data.decode(enc)
            except UnicodeDecodeError:
                continue

    for enc in ('utf-8', 'cp1252', 'latin-1'):
        try:
            return data.decode(enc)
        except UnicodeDecodeError:
            continue

    return data.decode('utf-8', errors='replace')


class SubtitleProcessor:
    """Processes subtitles to align with cut video and filter profanity"""
    
    # Use shared profanity word list
    PROFANITY_WORDS = PROFANITY_WORDS

    def __init__(self, profanity_words=None, include_religious: bool = False):
        if profanity_words is not None:
            self.PROFANITY_WORDS = set(profanity_words)
        else:
            self.PROFANITY_WORDS = get_profanity_words(include_religious=include_religious)
    
    def process_srt(self, input_srt: Path, output_srt: Path, 
                    removed_segments: List[Tuple[float, float]]) -> bool:
        """
        Process SRT subtitle file to align with cut video.
        
        Args:
            input_srt: Input SRT file path
            output_srt: Output SRT file path
            removed_segments: List of (start, end) segments removed from video
            
        Returns:
            True if successful, False otherwise
        """
        try:
            content = read_subtitle_text(input_srt)
            
            # Parse SRT entries
            entries = self._parse_srt(content)
            
            # Filter profanity from ALL subtitle text (remove profanity words, keep sentences)
            # This removes profanity words but keeps the rest of the sentence
            entries = self._filter_profanity(entries)
            
            # Process entries: clip overlapping subtitles to preserve content
            processed_entries = []
            for entry in entries:
                # Skip empty entries after profanity filtering
                if not entry['text'].strip():
                    continue
                
                entry_start = entry['start']
                entry_end = entry['end']
                
                # Check if entry overlaps with any removed segment
                overlaps_removal = False
                for remove_start, remove_end in removed_segments:
                    if entry_start < remove_end and entry_end > remove_start:
                        overlaps_removal = True
                        break
                
                if overlaps_removal:
                    # Clip the entry to non-overlapping segments
                    clipped_entries = self._clip_entry_to_keep_segments(entry, removed_segments)
                    processed_entries.extend(clipped_entries)
                else:
                    # Entry doesn't overlap - keep as-is for timestamp adjustment
                    processed_entries.append(entry)
            
            # Adjust timestamps for all processed entries
            adjusted_entries = self._adjust_timestamps(processed_entries, removed_segments)
            
            # Write output SRT
            self._write_srt(output_srt, adjusted_entries)
            
            return True
            
        except Exception as e:
            print(f"  Error processing subtitles: {e}")
            return False
    
    def _parse_srt(self, content: str) -> List[dict]:
        """Parse SRT content into list of subtitle entries"""
        entries = []
        
        # Split by double newlines (SRT entry separator)
        blocks = re.split(r'\n\s*\n', content.strip())
        
        for block in blocks:
            if not block.strip():
                continue
            
            lines = block.strip().split('\n')
            if len(lines) < 2:
                continue
            
            # First line is index
            try:
                index = int(lines[0])
            except ValueError:
                continue
            
            # Second line is timestamp
            timestamp_line = lines[1]
            time_match = re.match(r'(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*(\d{2}):(\d{2}):(\d{2}),(\d{3})', 
                                 timestamp_line)
            if not time_match:
                continue
            
            # Parse timestamps to seconds
            start_seconds = self._srt_time_to_seconds(time_match.groups()[:4])
            end_seconds = self._srt_time_to_seconds(time_match.groups()[4:])
            
            # Remaining lines are subtitle text
            text = '\n'.join(lines[2:]).strip()
            # Clean VTT timing tags and formatting tags
            text = self._clean_subtitle_text(text)
            
            entries.append({
                'index': index,
                'start': start_seconds,
                'end': end_seconds,
                'text': text
            })
        
        return entries
    
    def _clean_subtitle_text(self, text: str) -> str:
        """Remove VTT timing tags and formatting tags from subtitle text"""
        # Remove VTT timing tags like <00:00:37,260> or <00:00:37.260>
        text = re.sub(r'<00:\d{2}:\d{2}[,.]\d{3}>', '', text)
        # Remove VTT cue tags like <c>...</c>
        text = re.sub(r'<c>', '', text)
        text = re.sub(r'</c>', '', text)
        # Remove any other HTML-like tags
        text = re.sub(r'<[^>]+>', '', text)
        # Clean up multiple spaces
        text = re.sub(r'\s+', ' ', text)
        # Clean up spaces at start/end of lines
        lines = [line.strip() for line in text.split('\n')]
        text = '\n'.join(line for line in lines if line)
        return text.strip()
    
    def detect_profanity_segments(self, subtitle_path: Path, srt_window: float = None, pad: float = 0.2) -> List[Tuple[float, float, str]]:
        """
        Detect profanity in subtitle file and return segments with timestamps.
        
        Args:
            subtitle_path: Path to subtitle file (SRT or VTT)
            
        Returns:
            List of (start_time, end_time, words) tuples for profanity segments
        """
        try:
            content = read_subtitle_text(subtitle_path)

            # Parse based on file extension
            if subtitle_path.suffix.lower() == '.srt':
                entries = self._parse_srt(content)
            elif subtitle_path.suffix.lower() == '.vtt':
                entries = self._parse_vtt(content)
            else:
                # Try SRT format as default
                entries = self._parse_srt(content)
            
            profanity_segments = []
            
            for entry in entries:
                text = entry['text'].lower()
                # Check if entry contains any profanity
                found_profanity = []
                
                # First check for multi-word phrases (e.g., "fuck you", "ass hole")
                multi_word_phrases = [
                    ('fuck you', 'fuckyou'),
                    ('fuck off', 'fuckoff'),
                    ('ass hole', 'asshole'),
                    ('mother fucker', 'motherfucker'),
                    ('mother fuck', 'motherfuck'),
                ]
                
                for phrase, word_key in multi_word_phrases:
                    if word_key in self.PROFANITY_WORDS:
                        pattern = r'\b' + re.escape(phrase) + r'\b'
                        if re.search(pattern, text, re.IGNORECASE):
                            found_profanity.append(word_key)
                
                # Then check individual words
                for word in self.PROFANITY_WORDS:
                    # Use word boundary to match whole words only
                    pattern = r'\b' + re.escape(word) + r'\b'
                    if re.search(pattern, text, re.IGNORECASE):
                        if (
                            should_filter_word(word, text)
                            and word not in found_profanity
                        ):
                            found_profanity.append(word)
                
                if found_profanity:
                    # Determine segment bounds for removal
                    if srt_window is not None and srt_window > 0:
                        duration = max(0.0, entry['end'] - entry['start'])
                        window = min(srt_window, duration if duration > 0 else srt_window)
                        center = (entry['start'] + entry['end']) / 2.0 if duration > 0 else entry['start']
                        start = max(0.0, center - window / 2.0)
                        end = start + window
                    else:
                        # Default behavior: use the full cue with small padding
                        start = max(0.0, entry['start'] - pad)
                        end = entry['end'] + pad
                    profanity_segments.append((
                        start,
                        end,
                        ', '.join(sorted(set(found_profanity)))
                    ))
            
            return profanity_segments
            
        except Exception as e:
            print(f"  Error detecting profanity in subtitles: {e}")
            import traceback
            traceback.print_exc()
            return []
    
    def _filter_profanity(self, entries: List[dict]) -> List[dict]:
        """
        Filter profanity words from subtitle text entries.
        
        Args:
            entries: List of subtitle entry dictionaries
            
        Returns:
            List of entries with profanity filtered
        """
        filtered = []
        profanity_count = 0
        
        for entry in entries:
            original_text = entry['text']
            filtered_text = self._filter_text_profanity(original_text)
            
            if filtered_text != original_text:
                profanity_count += 1
            
            filtered.append({
                'index': entry['index'],
                'start': entry['start'],
                'end': entry['end'],
                'text': filtered_text
            })
        
        if profanity_count > 0:
            print(f"  Filtered profanity from {profanity_count} subtitle entry/entries")
        
        return filtered
    
    def _filter_text_profanity(self, text: str) -> str:
        """
        Filter profanity words from text, removing only the profanity words (not entire sentences).
        
        Args:
            text: Input text to filter
            
        Returns:
            Text with profanity words completely removed, sentences preserved
        """
        if not text:
            return text
        
        # First handle multi-word phrases (e.g., "fuck you", "ass hole")
        multi_word_phrases = [
            ('fuck you', 'fuckyou'),
            ('fuck off', 'fuckoff'),
            ('ass hole', 'asshole'),
            ('mother fucker', 'motherfucker'),
            ('mother fuck', 'motherfuck'),
        ]
        
        filtered_text = text
        for phrase, word_key in multi_word_phrases:
            if word_key in self.PROFANITY_WORDS:
                # Match phrase with word boundaries, case-insensitive
                pattern = r'\b' + re.escape(phrase) + r'\b'
                filtered_text = re.sub(pattern, '', filtered_text, flags=re.IGNORECASE)
        
        # Sort profanity by length (longest first) to match longer words first
        # e.g., "fucking" before "fuck", "motherfucker" before "fuck"
        sorted_profanity = sorted(self.PROFANITY_WORDS, key=len, reverse=True)
        
        # Process each profanity word
        for profanity in sorted_profanity:
            if not should_filter_word(profanity, text):
                continue

            # Create regex pattern that matches the profanity word as a whole word
            # \b ensures word boundaries, so "class" won't match "ass"
            # Match word with optional punctuation after it
            # Use word boundary to ensure we only match complete words
            pattern = r'\b' + re.escape(profanity) + r'\b'
            
            # Remove only the profanity word itself (case-insensitive)
            # This preserves the sentence structure
            filtered_text = re.sub(pattern, '', filtered_text, flags=re.IGNORECASE)
        
        # Clean up extra spaces that might result from word removal
        # Replace multiple spaces with single space
        filtered_text = re.sub(r' +', ' ', filtered_text)
        # Remove space before punctuation
        filtered_text = re.sub(r'\s+([.,!?;:])', r'\1', filtered_text)
        # Remove space after punctuation if followed by space
        filtered_text = re.sub(r'([.,!?;:])\s+', r'\1 ', filtered_text)
        # Clean up any double spaces that might remain
        filtered_text = re.sub(r'  +', ' ', filtered_text)
        # Remove leading/trailing spaces
        filtered_text = filtered_text.strip()
        
        return filtered_text
    
    def _clip_entry_to_keep_segments(self, entry: dict, 
                                     removed_segments: List[Tuple[float, float]]) -> List[dict]:
        """
        Clip an entry to only keep parts that are NOT in removed segments.
        Returns a list of entry dicts (may be empty, or 1-2 entries if split).
        
        Improved logic:
        - If overlap is small (<30% of entry), keep whole entry
        - If overlap is large, clip but only keep segments >= 0.5s duration
        - This prevents misaligned or missing subtitles
        """
        entry_start = entry['start']
        entry_end = entry['end']
        entry_text = entry['text']
        entry_index = entry['index']
        entry_duration = entry_end - entry_start
        
        # Sort removed segments
        sorted_removed = sorted(removed_segments, key=lambda x: x[0])
        
        # Calculate total overlap with removals
        total_removed_in_entry = 0.0
        for remove_start, remove_end in sorted_removed:
            if entry_start < remove_end and entry_end > remove_start:
                overlap_start = max(entry_start, remove_start)
                overlap_end = min(entry_end, remove_end)
                total_removed_in_entry += (overlap_end - overlap_start)
        
        # If less than 40% of entry is removed, keep whole entry
        # Timestamp adjustment will handle the small overlap
        if total_removed_in_entry < entry_duration * 0.4:
            return [entry]
        
        # If most of entry is removed, clip it but only keep segments with sufficient duration
        keep_segments = []
        current_time = entry_start
        
        for remove_start, remove_end in sorted_removed:
            if entry_start < remove_end and entry_end > remove_start:
                # Keep part before removal (if any)
                if current_time < remove_start:
                    keep_segments.append((current_time, remove_start))
                # Move current_time past the removal
                current_time = max(current_time, remove_end)
        
        # Keep part after last removal (if any)
        if current_time < entry_end:
            keep_segments.append((current_time, entry_end))
        
        # If no keep segments, entry is completely removed
        if not keep_segments:
            return []
        
        # Create entry dicts for each keep segment (only if duration >= 0.5s)
        result = []
        MIN_DURATION = 0.2  # Minimum duration for a subtitle entry (reduced to preserve more subtitles)
        
        for keep_start, keep_end in keep_segments:
            duration = keep_end - keep_start
            if duration >= MIN_DURATION:  # Only keep segments with sufficient duration
                result.append({
                    'index': entry_index,
                    'start': keep_start,
                    'end': keep_end,
                    'text': entry_text
                })
        
        # If all segments were too short, return empty (entry effectively removed)
        # But if entry had significant non-overlapping content,
        # keep the largest segment even if slightly below MIN_DURATION
        if not result and keep_segments:
            # Find the longest segment
            longest_segment = max(keep_segments, key=lambda x: x[1] - x[0])
            longest_duration = longest_segment[1] - longest_segment[0]
            # If it's at least 0.1s, keep it (very short but better than missing)
            if longest_duration >= 0.1:
                result.append({
                    'index': entry_index,
                    'start': longest_segment[0],
                    'end': longest_segment[1],
                    'text': entry_text
                })
        
        return result
    
    def _srt_time_to_seconds(self, time_parts: Tuple[str, str, str, str]) -> float:
        """Convert SRT time format (HH:MM:SS,mmm) to seconds"""
        hours, minutes, seconds, milliseconds = map(int, time_parts)
        return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0
    
    def _seconds_to_srt_time(self, seconds: float) -> str:
        """Convert seconds to SRT time format (HH:MM:SS,mmm)"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
    
    def _filter_entries(self, entries: List[dict], 
                       removed_segments: List[Tuple[float, float]]) -> List[dict]:
        """Remove subtitle entries that overlap with removed segments"""
        filtered = []
        
        # Sort removed segments for efficient checking
        sorted_removed = sorted(removed_segments, key=lambda x: x[0])
        
        for entry in entries:
            entry_start = entry['start']
            entry_end = entry['end']
            
            # Check if entry overlaps with any removed segment
            should_remove = False
            for remove_start, remove_end in sorted_removed:
                # Remove if entry overlaps with removed segment
                # Overlap condition: entry starts before remove ends AND entry ends after remove starts
                # Also remove if entry is completely within removed segment
                # Also remove if entry starts within removed segment (even if it extends beyond)
                if (entry_start < remove_end and entry_end > remove_start):
                    should_remove = True
                    break
                # Also check: if entry starts very close after removed segment, it might be affected
                # But we'll handle this in adjustment - if adjusted time falls in removed segment, skip it
            
            if not should_remove:
                filtered.append(entry)
        
        return filtered
    
    def _adjust_timestamps(self, entries: List[dict], 
                          removed_segments: List[Tuple[float, float]]) -> List[dict]:
        """
        Adjust subtitle timestamps to match a cut video timeline.

        For each original timestamp, subtract only removals that occur strictly
        before it. Timestamps that land inside a removed interval are snapped to
        the cleaned join point (where that cut is stitched).
        """
        if not removed_segments:
            return entries
        
        sorted_removed = sorted(removed_segments, key=lambda x: x[0])
        
        def map_timestamp(original_time: float) -> float:
            """Map original timestamp to cleaned video timestamp."""
            time_removed_before = 0.0
            for remove_start, remove_end in sorted_removed:
                if remove_end <= original_time:
                    # Removal fully before this timestamp.
                    time_removed_before += (remove_end - remove_start)
                    continue
                if remove_start <= original_time < remove_end:
                    # Inside a removed interval: snap to the cut join point.
                    return max(0.0, remove_start - time_removed_before)
                # Later removals do not affect this timestamp.
                break
            return max(0.0, original_time - time_removed_before)
        
        adjusted = []
        
        for entry in entries:
            orig_start = entry['start']
            orig_end = entry['end']
            orig_duration = max(0.1, orig_end - orig_start)
            
            new_start = map_timestamp(orig_start)
            new_end = map_timestamp(orig_end)
            
            # Keep a usable cue duration after mapping.
            if new_end <= new_start:
                new_end = new_start + orig_duration
            
            duration = new_end - new_start
            if duration < 0.1:
                continue
            
            adjusted.append({
                'index': entry['index'],
                'start': new_start,
                'end': new_end,
                'text': entry['text']
            })
        
        # Ensure cues stay chronological and avoid pathological overlaps from snaps.
        adjusted.sort(key=lambda e: (e['start'], e['end']))
        cleaned = []
        last_end = -1.0
        for entry in adjusted:
            start = entry['start']
            end = entry['end']
            if start < last_end:
                start = last_end
            if end <= start:
                end = start + 0.1
            cleaned.append({
                'index': entry['index'],
                'start': start,
                'end': end,
                'text': entry['text']
            })
            last_end = end
        
        return cleaned
    
    def _write_srt(self, output_path: Path, entries: List[dict]):
        """Write subtitle entries to SRT file"""
        with open(output_path, 'w', encoding='utf-8') as f:
            for i, entry in enumerate(entries, 1):
                f.write(f"{i}\n")
                f.write(f"{self._seconds_to_srt_time(entry['start'])} --> {self._seconds_to_srt_time(entry['end'])}\n")
                f.write(f"{entry['text']}\n")
                f.write("\n")
    
    def process_vtt(self, input_vtt: Path, output_vtt: Path,
                    removed_segments: List[Tuple[float, float]]) -> bool:
        """
        Process VTT subtitle file to align with cut video.
        
        Args:
            input_vtt: Input VTT file path
            output_vtt: Output VTT file path
            removed_segments: List of (start, end) segments removed from video
            
        Returns:
            True if successful, False otherwise
        """
        try:
            content = read_subtitle_text(input_vtt)

            # Parse VTT entries
            entries = self._parse_vtt(content)
            
            # Filter profanity from subtitle text
            entries = self._filter_profanity(entries)
            
            # Process entries: clip overlapping subtitles to preserve content
            processed_entries = []
            for entry in entries:
                # Skip empty entries after profanity filtering
                if not entry['text'].strip():
                    continue
                
                entry_start = entry['start']
                entry_end = entry['end']
                
                # Check if entry overlaps with any removed segment
                overlaps_removal = False
                for remove_start, remove_end in removed_segments:
                    if entry_start < remove_end and entry_end > remove_start:
                        overlaps_removal = True
                        break
                
                if overlaps_removal:
                    # Clip the entry to non-overlapping segments
                    clipped_entries = self._clip_entry_to_keep_segments(entry, removed_segments)
                    processed_entries.extend(clipped_entries)
                else:
                    # Entry doesn't overlap - keep as-is for timestamp adjustment
                    processed_entries.append(entry)
            
            # Adjust timestamps to account for removed time
            adjusted_entries = self._adjust_timestamps(processed_entries, removed_segments)
            
            # Write output VTT
            self._write_vtt(output_vtt, adjusted_entries)
            
            return True
            
        except Exception as e:
            print(f"  Error processing VTT subtitles: {e}")
            return False
    
    def _parse_vtt(self, content: str) -> List[dict]:
        """Parse VTT content into list of subtitle entries"""
        entries = []
        
        # Remove WEBVTT header and metadata
        lines = content.split('\n')
        start_idx = 0
        for i, line in enumerate(lines):
            if line.strip().startswith('WEBVTT'):
                start_idx = i + 1
                break
        
        # Find cue blocks
        current_cue = None
        for line in lines[start_idx:]:
            line = line.strip()
            if not line:
                if current_cue:
                    entries.append(current_cue)
                    current_cue = None
                continue
            
            # Check if line is a timestamp
            time_match = re.match(r'(\d{2}):(\d{2}):(\d{2})\.(\d{3})\s*-->\s*(\d{2}):(\d{2}):(\d{2})\.(\d{3})', 
                                 line)
            if time_match:
                if current_cue:
                    entries.append(current_cue)
                
                start_seconds = self._vtt_time_to_seconds(time_match.groups()[:4])
                end_seconds = self._vtt_time_to_seconds(time_match.groups()[4:])
                
                current_cue = {
                    'index': len(entries) + 1,
                    'start': start_seconds,
                    'end': end_seconds,
                    'text': ''
                }
            elif current_cue:
                # Add text to current cue
                if current_cue['text']:
                    current_cue['text'] += '\n' + line
                else:
                    current_cue['text'] = line
        
        if current_cue:
            entries.append(current_cue)
        
        # Clean all entries
        for entry in entries:
            entry['text'] = self._clean_subtitle_text(entry['text'])
        
        return entries
    
    def _vtt_time_to_seconds(self, time_parts: Tuple[str, str, str, str]) -> float:
        """Convert VTT time format (HH:MM:SS.mmm) to seconds"""
        hours, minutes, seconds, milliseconds = map(int, time_parts)
        return hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0
    
    def _seconds_to_vtt_time(self, seconds: float) -> str:
        """Convert seconds to VTT time format (HH:MM:SS.mmm)"""
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}"
    
    def _write_vtt(self, output_path: Path, entries: List[dict]):
        """Write subtitle entries to VTT file"""
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write("WEBVTT\n\n")
            for entry in entries:
                f.write(f"{self._seconds_to_vtt_time(entry['start'])} --> {self._seconds_to_vtt_time(entry['end'])}\n")
                f.write(f"{entry['text']}\n\n")