from typing import Dict, List, Optional, Tuple, Union, Any
from spacy.language import Language
from spacy.tokens import Doc, Span
from spacy.matcher import Matcher
from spacy.pipeline import Pipe
import logging

logger = logging.getLogger("context-enhancer")

ContextPattern = Dict[str, Union[str, bool, List[Dict[str, Any]]]]

@Language.factory("context_enhancer")
class ContextEnhancer(Pipe):
    """spaCy pipeline component that enhances or invalidates entities based on context."""
    
    def __init__(
        self,
        nlp: Language,
        name: str = "context_enhancer",
        patterns: List[ContextPattern] = None,
        added_context_words: Optional[List[str]] = None,
        confidence_boost: float = 0.35,
        min_enhanced_score: float = 0.4,
        context_window: Tuple[int, int] = (5, 2),
        allow_dependency_link: bool = True,
        style: str = "span",
        spans_key: str = "sc",
        ignore_recognizer_context: bool = False,
        ignore_sources: List[str] = ["unknown"]
    ):
        """Initialize ContextEnhancer.
        
        Args:
            nlp: The spaCy Language object
            name: Name of the component
            patterns: List of context patterns for enhancement/invalidation
            added_context_words: Additional context words to append to doc
            confidence_boost: Score increase for enhanced spans
            min_enhanced_score: Minimum score after enhancement
            context_window: (before, after) token window for context matching
            allow_dependency_link: If True, allows dependency-based context matching
            style: "ent" or "span" for where to store results
            spans_key: Key for doc.spans if style="span"
            ignore_recognizer_context: If True, ignores context patterns from recognizers
            ignore_sources: List of sources to ignore when enhancing spans
        """
        self.nlp = nlp
        self.name = name
        self.added_context_words = added_context_words if added_context_words else []
        self.confidence_boost = confidence_boost
        self.min_enhanced_score = min_enhanced_score
        self.context_before, self.context_after = context_window
        self.allow_dependency_link = allow_dependency_link
        self.style = style
        self.spans_key = spans_key
        self.ignore_recognizer_context = ignore_recognizer_context
        self.ignore_sources = ignore_sources

        # Store our own patterns
        self._patterns = patterns if patterns else []

        # Ensure extensions exist
        if not Span.has_extension("score"):
            Span.set_extension("score", default=0.0)
        if not Span.has_extension("supportive_context"):
            Span.set_extension("supportive_context", default=[])
    
    def __call__(self, doc: Doc) -> Doc:
        """Process document."""
        from anonymacy.recognizer import BaseRecognizer
        # Get spans
        spans = self._get_spans(doc)
        if not spans:
            return doc
        
        # Collect all patterns
        all_patterns = []
        all_patterns.extend(self._patterns)
        
        # Collect patterns from recognizers unless ignoring them
        if not self.ignore_recognizer_context:
            for name, pipe in self.nlp.pipeline:
                if isinstance(pipe, BaseRecognizer) and pipe.context_patterns:
                    all_patterns.extend(pipe.context_patterns)
        
        if not all_patterns:
            return doc
        
        # Build matcher with all patterns
        matcher = Matcher(self.nlp.vocab)
        pattern_map = {}  # Maps pattern_id to pattern config
        
        for idx, pattern_config in enumerate(all_patterns):
            try:
                label = pattern_config["label"]
                pattern_list = pattern_config["pattern"]
                invalidate = pattern_config.get("invalidate", False)
                
                pattern_id = f"{label}_{idx}"
                matcher.add(pattern_id, [pattern_list])
                pattern_map[pattern_id] = {
                    "label": label,
                    "pattern": pattern_list,
                    "invalidate": invalidate
                }
            except KeyError as e:
                logger.warning(f"Skipping invalid pattern {pattern_config}: missing {e}")
                continue
        
        # Create extended doc if needed
        if self.added_context_words:
            extended_text = doc.text + " " + " ".join(self.added_context_words)
            extended_doc = self.nlp.make_doc(extended_text)
        else:
            extended_doc = doc
        
        # Find all matches
        matches = matcher(extended_doc)
        
        # Process each span
        processed_spans = []
        for span in spans:

            # Skip if source is in ignore list
            source = getattr(span._, "source", "unknown")
            if source in self.ignore_sources:
                processed_spans.append(span)
                continue

            # Skip if already enhanced
            if hasattr(span._, "supportive_context") and span._.supportive_context:
                processed_spans.append(span)
                continue
         
            should_invalidate = False
            supportive_context = []
            
            # Check all matches for this span
            for match_id, start, end in matches:
                match = doc[start:end]
                if not self._in_context(span, match, len(doc)):
                    continue
                
                pattern_id = self.nlp.vocab.strings[match_id]
                pattern_config = pattern_map.get(pattern_id)
                
                if pattern_config and pattern_config["label"] == span.label_:
                    if pattern_config["invalidate"]:
                        should_invalidate = True
                        break  # No need to continue if invalidating
                    else:
                        context_text = extended_doc[start:end].text
                        supportive_context.append(context_text)
            
            # Apply results
            if should_invalidate:
                continue  # Skip invalidated spans
            
            if supportive_context:
                enhanced_span = self._create_enhanced_span(span, supportive_context)
                processed_spans.append(enhanced_span)
            else:
                processed_spans.append(span)
        
        # Update spans
        self._set_spans(doc, processed_spans)
        return doc
    
    def _get_spans(self, doc: Doc) -> List[Span]:
        """Get current spans."""
        return (
            list(doc.ents) if self.style == "ent" 
            else list(doc.spans.get(self.spans_key, []))
        )
    
    def _set_spans(self, doc: Doc, spans: List[Span]) -> None:
        """Set spans on document."""
        if self.style == "ent":
            doc.ents = spans
        else:
            doc.spans[self.spans_key] = spans

    def _has_dependency_link(self, span: Span, match_tokens: Span) -> bool:
        """Check if span and match_tokens are directly related via dependency (head or child)."""
        if not span.doc.has_annotation("DEP"):
            return False

        for token in span:
            for match_token in match_tokens:
                if token.head == match_token or match_token.head == token:
                    return True
        return False
    
    def _in_context(self, span: Span, match: Span, doc_length: int) -> bool:
        """Check if match is in span's context window."""
        # Check if match is one of the added_context_words

        if match.start >= doc_length:
            return True
        
        context_window_start = max(0, span.start - self.context_before)
        context_window_end = min(doc_length, span.end + self.context_after)

        # Before context - check if match ends in the before-window
        if (context_window_start < match.end <= span.start):
            return True
        
        # After context - check if match starts in the after-window
        if (span.end <= match.start < context_window_end):
            return True

        # Outside of window: try dependency-based relation as fallback
        if self.allow_dependency_link:
            return self._has_dependency_link(span, match)
        
        return False
    
    def _create_enhanced_span(self, span: Span, context: List[str]) -> Span:
        """Create enhanced version of span."""
        logger.debug(f"Enhancing span: {span.text} [{span.label_}] with contexts: {context}")
        enhanced = Span(span.doc, span.start, span.end, label=span.label_)
        
        # Calculate new score
        original_score = getattr(span._, "score", 1.0)
        new_score = min(
            max(original_score + self.confidence_boost, self.min_enhanced_score),
            1.0
        )
        
        enhanced._.score = new_score
        enhanced._.supportive_context = context
        
        return enhanced