"""
Caption-based analysis tools for sign language research.

Provides analysis tools for caption/subtitle-based datasets
that don't have gloss-level annotations (like YTSL-25, YT-ASL).

These tools work with continuous text segments rather than
individual sign labels.
"""

from __future__ import annotations

import re
from collections import Counter
from dataclasses import dataclass, field
from typing import Any

import numpy as np


@dataclass
class CaptionSegment:
    """A single caption/subtitle segment."""

    file_id: str
    start_time: float
    end_time: float
    text: str
    tier: str = "caption"
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def duration(self) -> float:
        return self.end_time - self.start_time

    @property
    def words(self) -> list[str]:
        """Get words from caption text."""
        return re.findall(r"\b\w+\b", self.text.lower())

    @property
    def word_count(self) -> int:
        return len(self.words)


@dataclass
class CaptionStatistics:
    """Statistics for caption-based analysis."""

    total_captions: int
    total_words: int
    unique_words: int
    total_duration: float
    mean_caption_duration: float
    std_caption_duration: float
    mean_words_per_caption: float
    word_frequencies: dict[str, int]
    caption_count_by_file: dict[str, int]


def analyze_caption_corpus(
    captions_by_file: dict[str, list[dict[str, Any]]],
    min_word_freq: int = 1,
    exclude_stopwords: bool = False,
    stopwords: set | None = None,
) -> CaptionStatistics:
    """
    Analyze a corpus of caption-based annotations.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists
        min_word_freq: Minimum word frequency to include in results
        exclude_stopwords: Whether to exclude common stopwords
        stopwords: Custom set of stopwords to exclude

    Returns:
        CaptionStatistics with corpus-level statistics
    """
    if stopwords is None and exclude_stopwords:
        stopwords = {
            "the",
            "a",
            "an",
            "and",
            "or",
            "but",
            "in",
            "on",
            "at",
            "to",
            "for",
            "of",
            "with",
            "is",
            "are",
            "was",
            "were",
            "be",
            "been",
            "being",
            "have",
            "has",
            "had",
            "do",
            "does",
            "did",
            "will",
            "would",
            "could",
            "should",
            "may",
            "might",
            "must",
            "shall",
            "this",
            "that",
            "these",
            "those",
            "i",
            "you",
            "he",
            "she",
            "it",
            "we",
            "they",
            "my",
            "your",
            "his",
            "her",
            "its",
            "our",
            "their",
        }
    elif stopwords is None:
        stopwords = set()

    word_freq: Counter[str] = Counter()
    total_captions = 0
    total_words = 0
    durations = []
    words_per_caption = []
    caption_counts = {}

    for file_id, captions in captions_by_file.items():
        caption_counts[file_id] = len(captions)
        total_captions += len(captions)

        for cap in captions:
            # Get text content
            text = cap.get("text", "") or cap.get("label", "")
            start = cap.get("start", 0) or cap.get("start_time", 0)
            end = cap.get("end", 0) or cap.get("end_time", 0)

            # Duration
            duration = end - start
            if duration > 0:
                durations.append(duration)

            # Word analysis
            words = re.findall(r"\b\w+\b", text.lower())
            filtered_words = [w for w in words if w not in stopwords]

            total_words += len(filtered_words)
            words_per_caption.append(len(filtered_words))
            word_freq.update(filtered_words)

    # Filter by minimum frequency
    filtered_freq = {w: f for w, f in word_freq.items() if f >= min_word_freq}

    durations_arr = np.array(durations) if durations else np.array([0])

    return CaptionStatistics(
        total_captions=total_captions,
        total_words=total_words,
        unique_words=len(filtered_freq),
        total_duration=float(durations_arr.sum()),
        mean_caption_duration=float(durations_arr.mean()) if len(durations) > 0 else 0,
        std_caption_duration=float(durations_arr.std()) if len(durations) > 1 else 0,
        mean_words_per_caption=float(np.mean(words_per_caption)) if words_per_caption else 0,
        word_frequencies=filtered_freq,
        caption_count_by_file=caption_counts,
    )


def extract_caption_keywords(
    captions_by_file: dict[str, list[dict[str, Any]]],
    top_n: int = 100,
    min_freq: int = 2,
    exclude_stopwords: bool = True,
) -> list[tuple[str, int]]:
    """
    Extract top keywords from caption corpus.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists
        top_n: Number of top keywords to return
        min_freq: Minimum frequency threshold
        exclude_stopwords: Whether to exclude common stopwords

    Returns:
        List of (word, frequency) tuples
    """
    stats = analyze_caption_corpus(
        captions_by_file,
        min_word_freq=min_freq,
        exclude_stopwords=exclude_stopwords,
    )

    sorted_words = sorted(stats.word_frequencies.items(), key=lambda x: -x[1])

    return sorted_words[:top_n]


def search_captions(
    captions_by_file: dict[str, list[dict[str, Any]]],
    query: str,
    case_sensitive: bool = False,
    whole_word: bool = False,
    max_results: int = 100,
) -> list[dict[str, Any]]:
    """
    Search for a query string in captions.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists
        query: Search query string
        case_sensitive: Whether to match case
        whole_word: Whether to match whole words only
        max_results: Maximum number of results to return

    Returns:
        List of matching caption results with context
    """
    results = []

    if not case_sensitive:
        query = query.lower()

    if whole_word:
        pattern = rf"\b{re.escape(query)}\b"
        regex = re.compile(pattern, 0 if case_sensitive else re.IGNORECASE)

    for file_id, captions in captions_by_file.items():
        for i, cap in enumerate(captions):
            text = cap.get("text", "") or cap.get("label", "")
            search_text = text if case_sensitive else text.lower()

            match = False
            if whole_word:
                match = regex.search(text) is not None
            else:
                match = query in search_text

            if match:
                # Get context (previous and next captions)
                prev_cap = captions[i - 1] if i > 0 else None
                next_cap = captions[i + 1] if i < len(captions) - 1 else None

                results.append(
                    {
                        "file_id": file_id,
                        "index": i,
                        "text": text,
                        "start": cap.get("start", 0),
                        "end": cap.get("end", 0),
                        "context_before": prev_cap.get("text", "") if prev_cap else "",
                        "context_after": next_cap.get("text", "") if next_cap else "",
                    }
                )

                if len(results) >= max_results:
                    return results

    return results


def compute_caption_alignment_stats(
    captions_by_file: dict[str, list[dict[str, Any]]],
) -> dict[str, Any]:
    """
    Compute statistics about caption timing and alignment.

    Useful for understanding the temporal distribution of captions
    and detecting potential alignment issues.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists

    Returns:
        Dictionary with alignment statistics
    """
    gaps = []
    overlaps = []
    durations = []

    for file_id, captions in captions_by_file.items():
        if len(captions) < 2:
            continue

        # Sort by start time
        sorted_caps = sorted(captions, key=lambda x: x.get("start", 0))

        for i in range(len(sorted_caps) - 1):
            curr = sorted_caps[i]
            next_cap = sorted_caps[i + 1]

            curr_end = curr.get("end", 0)
            next_start = next_cap.get("start", 0)

            gap = next_start - curr_end

            if gap > 0:
                gaps.append(gap)
            elif gap < 0:
                overlaps.append(-gap)

            duration = curr_end - curr.get("start", 0)
            if duration > 0:
                durations.append(duration)

    gaps_arr = np.array(gaps) if gaps else np.array([0])
    overlaps_arr = np.array(overlaps) if overlaps else np.array([0])
    durations_arr = np.array(durations) if durations else np.array([0])

    return {
        "num_gaps": len(gaps),
        "num_overlaps": len(overlaps),
        "mean_gap": float(gaps_arr.mean()) if len(gaps) > 0 else 0,
        "max_gap": float(gaps_arr.max()) if len(gaps) > 0 else 0,
        "mean_overlap": float(overlaps_arr.mean()) if len(overlaps) > 0 else 0,
        "max_overlap": float(overlaps_arr.max()) if len(overlaps) > 0 else 0,
        "mean_duration": float(durations_arr.mean()) if len(durations) > 0 else 0,
        "std_duration": float(durations_arr.std()) if len(durations) > 1 else 0,
        "total_files": len(captions_by_file),
    }


def build_caption_concordance(
    captions_by_file: dict[str, list[dict[str, Any]]],
    target_word: str,
    context_words: int = 5,
    case_sensitive: bool = False,
    max_results: int = 100,
) -> list[dict[str, Any]]:
    """
    Build a concordance view for a target word in captions.

    Unlike gloss concordance, this shows word context within
    caption text rather than segment context.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists
        target_word: Word to search for
        context_words: Number of context words on each side
        case_sensitive: Whether to match case
        max_results: Maximum results to return

    Returns:
        List of concordance entries
    """
    results = []
    target = target_word if case_sensitive else target_word.lower()

    for file_id, captions in captions_by_file.items():
        for cap_idx, cap in enumerate(captions):
            text = cap.get("text", "") or cap.get("label", "")
            words = text.split()

            for i, word in enumerate(words):
                word_compare = word if case_sensitive else word.lower()
                # Remove punctuation for comparison
                word_compare = re.sub(r"[^\w\s]", "", word_compare)

                if word_compare == target:
                    left = words[max(0, i - context_words) : i]
                    right = words[i + 1 : i + 1 + context_words]

                    results.append(
                        {
                            "file_id": file_id,
                            "caption_idx": cap_idx,
                            "word_idx": i,
                            "word": word,
                            "left_context": " ".join(left),
                            "right_context": " ".join(right),
                            "full_text": text,
                            "start": cap.get("start", 0),
                            "end": cap.get("end", 0),
                        }
                    )

                    if len(results) >= max_results:
                        return results

    return results


def extract_caption_ngrams(
    captions_by_file: dict[str, list[dict[str, Any]]],
    n: int = 2,
    min_freq: int = 2,
    max_ngrams: int = 100,
    exclude_stopwords: bool = True,
) -> list[tuple[tuple[str, ...], int]]:
    """
    Extract word n-grams from captions.

    Args:
        captions_by_file: Dictionary mapping file IDs to caption lists
        n: N-gram size
        min_freq: Minimum frequency threshold
        max_ngrams: Maximum n-grams to return
        exclude_stopwords: Whether to exclude n-grams containing stopwords

    Returns:
        List of (ngram_tuple, frequency) sorted by frequency
    """
    stopwords = (
        {
            "the",
            "a",
            "an",
            "and",
            "or",
            "but",
            "in",
            "on",
            "at",
            "to",
            "for",
            "of",
            "with",
            "is",
            "are",
            "was",
            "were",
            "be",
            "been",
        }
        if exclude_stopwords
        else set()
    )

    ngram_counts: Counter[tuple[str, ...]] = Counter()

    for captions in captions_by_file.values():
        for cap in captions:
            text = cap.get("text", "") or cap.get("label", "")
            words = re.findall(r"\b\w+\b", text.lower())

            for i in range(len(words) - n + 1):
                ngram = tuple(words[i : i + n])

                # Skip if contains stopword
                if exclude_stopwords and any(w in stopwords for w in ngram):
                    continue

                ngram_counts[ngram] += 1

    # Filter and sort
    result = [
        (ng, count) for ng, count in ngram_counts.most_common(max_ngrams) if count >= min_freq
    ]

    return result
