"""Core functionality for semantic text alignment and containment analysis."""

import asyncio
import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
from rich.console import Console
from rich.logging import RichHandler
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

# Configure rich logging
console = Console()
logging.basicConfig(
    level=logging.INFO,
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(console=console, rich_tracebacks=True)]
)
logger = logging.getLogger(__name__)


@dataclass
class AlignmentPair:
    """Represents a single alignment pair between two text segments."""
    
    type: str  # "match", "gap_a", "gap_b", "unaligned_a", "unaligned_b"
    text_a: Optional[str] = None
    text_b: Optional[str] = None
    score: Optional[float] = None


@dataclass
class AlignmentResult:
    """Complete result of text alignment analysis."""
    
    pairs: List[AlignmentPair]
    containment_score: float
    total_paragraphs_a: int
    total_paragraphs_b: int
    matched_pairs: int
    gap_pairs: int
    unaligned_pairs: int


class SemanticAligner:
    """Advanced semantic text alignment using Smith-Waterman algorithm."""
    
    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        gap_penalty: float = 0.3,
        similarity_threshold: float = 0.5
    ) -> None:
        """
        Initialize the semantic aligner.
        
        Args:
            model_name: Name of the sentence transformer model to use
            gap_penalty: Penalty for gaps in alignment
            similarity_threshold: Minimum similarity score for matches
        """
        self.model_name = model_name
        self.gap_penalty = gap_penalty
        self.similarity_threshold = similarity_threshold
        self._model: Optional[SentenceTransformer] = None
        
    async def _load_model(self) -> SentenceTransformer:
        """Load the sentence transformer model asynchronously."""
        if self._model is None:
            logger.info(f"Loading model: {self.model_name}")
            # Run model loading in thread pool to avoid blocking
            loop = asyncio.get_event_loop()
            self._model = await loop.run_in_executor(
                None, SentenceTransformer, self.model_name
            )
            logger.info("Model loaded successfully")
        return self._model
    
    async def _get_embeddings(self, paragraphs: List[str]) -> np.ndarray:
        """Generate embeddings for paragraphs asynchronously."""
        if not paragraphs:
            return np.array([])
            
        model = await self._load_model()
        logger.info(f"Generating embeddings for {len(paragraphs)} paragraphs")
        
        # Run encoding in thread pool
        loop = asyncio.get_event_loop()
        embeddings = await loop.run_in_executor(
            None, lambda: model.encode(paragraphs, show_progress_bar=True)
        )
        
        logger.info("Embeddings generated successfully")
        return embeddings
    
    def _calculate_similarity_matrix(
        self, 
        vectors_a: np.ndarray, 
        vectors_b: np.ndarray
    ) -> np.ndarray:
        """Calculate similarity matrix between two sets of vectors."""
        if vectors_a.size == 0 or vectors_b.size == 0:
            return np.array([])
        return cosine_similarity(vectors_a, vectors_b)
    
    def _smith_waterman_align(
        self,
        paragraphs_a: List[str],
        paragraphs_b: List[str],
        vectors_a: np.ndarray,
        vectors_b: np.ndarray,
    ) -> List[AlignmentPair]:
        """
        Perform Smith-Waterman alignment between two paragraph sequences.
        
        Args:
            paragraphs_a: Paragraphs from text A
            paragraphs_b: Paragraphs from text B
            vectors_a: Embeddings for paragraphs A
            vectors_b: Embeddings for paragraphs B
            
        Returns:
            List of alignment pairs
        """
        n, m = len(paragraphs_a), len(paragraphs_b)
        logger.info(f"Starting Smith-Waterman alignment: A={n} paragraphs, B={m} paragraphs")
        
        # Initialize matrices
        score_matrix = np.zeros((n + 1, m + 1))
        path_matrix = np.zeros((n + 1, m + 1), dtype=int)
        
        max_score = 0.0
        max_pos = (0, 0)
        
        # Fill score matrix
        for i in range(1, n + 1):
            for j in range(1, m + 1):
                # Calculate similarity score
                similarity = cosine_similarity(
                    vectors_a[i - 1].reshape(1, -1),
                    vectors_b[j - 1].reshape(1, -1)
                )[0][0]
                
                # Calculate scores from different directions
                diagonal_score = score_matrix[i - 1, j - 1] + similarity
                up_score = score_matrix[i - 1, j] - self.gap_penalty
                left_score = score_matrix[i, j - 1] - self.gap_penalty
                
                current_max = max(0, diagonal_score, up_score, left_score)
                score_matrix[i, j] = current_max
                
                # Record path
                if current_max == diagonal_score:
                    path_matrix[i, j] = 1
                elif current_max == up_score:
                    path_matrix[i, j] = 2
                elif current_max == left_score:
                    path_matrix[i, j] = 3
                
                # Track global maximum
                if current_max > max_score:
                    max_score = current_max
                    max_pos = (i, j)
        
        logger.info(f"Matrix filled. Max score: {max_score:.4f} at position {max_pos}")
        
        # Backtrack to find alignment
        alignment = []
        i, j = max_pos
        aligned_indices_a = set()
        aligned_indices_b = set()
        
        while i > 0 and j > 0 and score_matrix[i, j] > 0:
            path = path_matrix[i, j]
            
            if path == 1:  # Match
                similarity = cosine_similarity(
                    vectors_a[i-1].reshape(1,-1), 
                    vectors_b[j-1].reshape(1,-1)
                )[0][0]
                
                alignment.insert(0, AlignmentPair(
                    type="match",
                    text_a=paragraphs_a[i - 1],
                    text_b=paragraphs_b[j - 1],
                    score=similarity
                ))
                aligned_indices_a.add(i - 1)
                aligned_indices_b.add(j - 1)
                i -= 1
                j -= 1
                
            elif path == 2:  # Gap in B
                alignment.insert(0, AlignmentPair(
                    type="gap_b",
                    text_a=paragraphs_a[i - 1]
                ))
                aligned_indices_a.add(i - 1)
                i -= 1
                
            elif path == 3:  # Gap in A
                alignment.insert(0, AlignmentPair(
                    type="gap_a",
                    text_b=paragraphs_b[j - 1]
                ))
                aligned_indices_b.add(j - 1)
                j -= 1
            else:
                break
        
        # Add unaligned paragraphs
        for idx, p in enumerate(paragraphs_a):
            if idx not in aligned_indices_a:
                alignment.append(AlignmentPair(type="unaligned_a", text_a=p))
        
        for idx, p in enumerate(paragraphs_b):
            if idx not in aligned_indices_b:
                alignment.append(AlignmentPair(type="unaligned_b", text_b=p))
        
        logger.info(f"Alignment complete. Found {len(alignment)} pairs")
        return alignment
    
    async def align_texts(
        self, 
        text_a: str, 
        text_b: str
    ) -> AlignmentResult:
        """
        Align two texts and calculate containment score.
        
        Args:
            text_a: First text to align
            text_b: Second text to align
            
        Returns:
            Complete alignment result with statistics
        """
        # Split into paragraphs
        paragraphs_a = split_paragraphs(text_a)
        paragraphs_b = split_paragraphs(text_b)
        
        # Generate embeddings
        vectors_a = await self._get_embeddings(paragraphs_a)
        vectors_b = await self._get_embeddings(paragraphs_b)
        
        # Perform alignment
        alignment_pairs = self._smith_waterman_align(
            paragraphs_a, paragraphs_b, vectors_a, vectors_b
        )
        
        # Calculate containment score
        containment_score = calculate_containment_score(alignment_pairs, paragraphs_a)
        
        # Calculate statistics
        matched_pairs = sum(1 for p in alignment_pairs if p.type == "match")
        gap_pairs = sum(1 for p in alignment_pairs if p.type in ("gap_a", "gap_b"))
        unaligned_pairs = sum(1 for p in alignment_pairs if p.type in ("unaligned_a", "unaligned_b"))
        
        return AlignmentResult(
            pairs=alignment_pairs,
            containment_score=containment_score,
            total_paragraphs_a=len(paragraphs_a),
            total_paragraphs_b=len(paragraphs_b),
            matched_pairs=matched_pairs,
            gap_pairs=gap_pairs,
            unaligned_pairs=unaligned_pairs
        )


def split_paragraphs(text: str) -> List[str]:
    """
    Split text into paragraphs based on double newlines.
    
    Args:
        text: Input text to split
        
    Returns:
        List of paragraph strings
    """
    logger.info("Splitting text into paragraphs")
    paragraphs = text.strip().split('\n\n')
    # Filter out empty paragraphs
    result = [p.strip() for p in paragraphs if p.strip()]
    logger.info(f"Split into {len(result)} paragraphs")
    return result


def calculate_containment_score(
    alignment_pairs: List[AlignmentPair], 
    paragraphs_a: List[str]
) -> float:
    """
    Calculate semantic containment score of text B relative to text A.
    
    Args:
        alignment_pairs: List of alignment pairs from Smith-Waterman
        paragraphs_a: Original paragraphs from text A
        
    Returns:
        Containment score between 0 and 1
    """
    logger.info("Calculating weighted semantic containment score")
    
    total_chars_a = sum(len(p) for p in paragraphs_a)
    if total_chars_a == 0:
        return 1.0 if not paragraphs_a else 0.0
    
    weighted_matched_chars_a = 0.0
    for pair in alignment_pairs:
        if pair.type == "match":
            # Only consider positive scores to avoid negative impact
            score = max(0, pair.score or 0)
            weighted_matched_chars_a += len(pair.text_a or "") * score
    
    score = weighted_matched_chars_a / total_chars_a
    logger.info(f"Containment score: {score:.2%}")
    return score 