"""
Extraction Pipeline - pluggable, configurable entity extraction.

This module provides:
1. Pluggable extractor chain
2. Always-on normalization
3. Automatic cross-document linking
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional, Protocol
import networkx as nx
import os

from .normalizer import EntityNormalizer, NormalizedEntity, get_normalizer
from .linker import EntityLinker, LinkedEntity, get_linker


class EntityExtractor(Protocol):
    """Protocol for entity extractors."""

    def extract(self, text: str) -> tuple[list[str], list[tuple[str, str, str]]]:
        """
        Extract entities and relations from text.

        Args:
            text: Input text

        Returns:
            Tuple of (entities, relations) where relations are (subj, rel, obj)
        """
        ...


@dataclass
class ExtractionResult:
    """Result of extraction pipeline."""
    entities: list[NormalizedEntity]
    relations: list[tuple[NormalizedEntity, str, NormalizedEntity]]
    linked_entities: list[LinkedEntity]
    graph: nx.DiGraph
    raw_entities: list[str]
    raw_relations: list[tuple[str, str, str]]


class ExtractionPipeline:
    """
    Pluggable extraction pipeline with always-on normalization.

    Features:
    - Chain multiple extractors (merge results)
    - Automatic normalization (can't be disabled)
    - Cross-document entity linking
    - Graph construction
    """

    def __init__(
        self,
        extractors: Optional[list] = None,
        normalizer: Optional[EntityNormalizer] = None,
        linker: Optional[EntityLinker] = None,
    ):
        """
        Initialize the extraction pipeline.

        Args:
            extractors: List of extractors to chain (uses regex if empty)
            normalizer: Entity normalizer (uses global if not provided)
            linker: Entity linker (uses global if not provided)
        """
        self.extractors: list = extractors or []
        self.normalizer = normalizer or get_normalizer()
        self.linker = linker or get_linker()

    def add_extractor(self, extractor) -> None:
        """Add an extractor to the pipeline."""
        self.extractors.append(extractor)

    def extract(
        self,
        text: str,
        episode_id: Optional[str] = None,
        link_entities: bool = True,
    ) -> ExtractionResult:
        """
        Extract entities and relations from text.

        Args:
            text: Input text
            episode_id: ID for cross-document linking (optional)
            link_entities: Whether to register entities in global linker

        Returns:
            ExtractionResult with normalized entities, relations, and graph
        """
        # Step 1: Run all extractors and merge results
        all_entities: list[str] = []
        all_relations: list[tuple[str, str, str]] = []

        for extractor in self.extractors:
            try:
                entities, relations = extractor.extract(text)
                all_entities.extend(entities)
                all_relations.extend(relations)
            except Exception as e:  # nosec B112 - intentional: skip failed extractors
                import warnings
                warnings.warn(f"Extractor {type(extractor).__name__} failed: {e}")
                continue

        # Step 2: Normalize entities (ALWAYS ON - cannot skip)
        normalized_entities = self.normalizer.normalize_many(all_entities)

        # Step 3: Normalize relations
        normalized_relations = self.normalizer.normalize_relations(all_relations)

        # Step 4: Link entities across documents
        linked_entities: list[LinkedEntity] = []
        if link_entities and episode_id:
            entity_strings = [e.canonical for e in normalized_entities]
            linked_entities = self.linker.link(entity_strings, episode_id)

        # Step 5: Build graph
        graph = self._build_graph(normalized_entities, normalized_relations)

        return ExtractionResult(
            entities=normalized_entities,
            relations=normalized_relations,
            linked_entities=linked_entities,
            graph=graph,
            raw_entities=all_entities,
            raw_relations=all_relations,
        )

    def _build_graph(
        self,
        entities: list[NormalizedEntity],
        relations: list[tuple[NormalizedEntity, str, NormalizedEntity]],
    ) -> nx.DiGraph:
        """Build NetworkX graph from normalized entities and relations."""
        G = nx.DiGraph()

        # Add all entities as nodes
        for entity in entities:
            G.add_node(
                entity.canonical,
                original=entity.original,
                type=entity.entity_type
            )

        # Add relations as edges
        for subj, rel, obj in relations:
            G.add_edge(
                subj.canonical,
                obj.canonical,
                relation=rel,
                relation_type=rel,
            )

        return G


# Built-in simple regex extractor (no dependencies)

class SimpleRegexExtractor:
    """
    Simple regex-based entity extractor with no external dependencies.

    Extracts:
    - Capitalized phrases (likely proper nouns/entities)
    - Simple relation patterns (X verb Y)
    """

    # Common verbs for relation extraction
    RELATION_VERBS = [
        'founded', 'created', 'acquired', 'bought', 'sold',
        'works at', 'works for', 'employed by', 'leads', 'manages',
        'located in', 'based in', 'headquartered in',
        'married to', 'parent of', 'child of',
        'studied at', 'graduated from', 'attended',
        'owns', 'developed', 'invented', 'discovered',
        'is a', 'is the', 'was a', 'was the',
    ]

    def __init__(self):
        import re
        # Pattern for capitalized phrases (entities)
        self._entity_pattern = re.compile(
            r'\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b'
        )
        # Build relation patterns
        self._relation_patterns = []
        for verb in self.RELATION_VERBS:
            pattern = re.compile(
                rf'([A-Z][a-zA-Z\s]+?)\s+{verb}\s+([A-Z][a-zA-Z\s]+)',
                re.IGNORECASE
            )
            self._relation_patterns.append((verb.replace(' ', '_'), pattern))

    def extract(
        self, text: str
    ) -> tuple[list[str], list[tuple[str, str, str]]]:
        """Extract entities and relations using simple regex."""
        import re

        # Extract entities (capitalized phrases)
        entities = list(set(self._entity_pattern.findall(text)))

        # Extract relations
        relations: list[tuple[str, str, str]] = []
        for rel_type, pattern in self._relation_patterns:
            for match in pattern.finditer(text):
                subj = match.group(1).strip()
                obj = match.group(2).strip()
                # Clean up - remove trailing punctuation
                obj = re.sub(r'[.,;:!?]+$', '', obj)
                if subj and obj and len(subj) > 1 and len(obj) > 1:
                    relations.append((subj, rel_type, obj))

        return entities, relations


# Factory functions for common pipeline configurations

def create_regex_pipeline() -> ExtractionPipeline:
    """Create pipeline with simple regex extractor (no dependencies)."""
    pipeline = ExtractionPipeline()
    pipeline.add_extractor(SimpleRegexExtractor())
    return pipeline


def create_llm_pipeline(
    provider: str = "openai",
    api_key: Optional[str] = None
) -> ExtractionPipeline:
    """Create pipeline with LLM extractor."""
    from ..llm_extractor import LLMExtractor

    class LLMExtractorAdapter:
        def __init__(self, provider: str, api_key: Optional[str]):
            self._extractor = LLMExtractor(
                provider=provider,
                api_key=api_key,
                cache_extractions=True,
            )

        def extract(
            self, text: str
        ) -> tuple[list[str], list[tuple[str, str, str]]]:
            result = self._extractor.extract(text)
            return result.entities, result.relations

    pipeline = ExtractionPipeline()
    pipeline.add_extractor(LLMExtractorAdapter(provider, api_key))
    return pipeline


def create_hybrid_pipeline(
    provider: str = "openai",
    api_key: Optional[str] = None,
    model: Optional[str] = None,
) -> ExtractionPipeline:
    """
    Create pipeline with both simple regex and optionally LLM extractors.
    LLM is primary (if API key provided), regex is fallback.
    """
    pipeline = ExtractionPipeline()

    # Add LLM if API key provided or found in env
    if not api_key:
        if provider == "openai":
            api_key = os.getenv("OPENAI_API_KEY")
        # Add other providers here if needed

    if api_key:
        try:
            from ..llm_extractor import LLMExtractor

            class LLMExtractorAdapter:
                def __init__(self, provider: str, api_key: str, model: Optional[str] = None):
                    self._extractor = LLMExtractor(
                        provider=provider,
                        api_key=api_key,
                        model=model,
                        cache_extractions=True,
                    )

                def extract(
                    self, text: str
                ) -> tuple[list[str], list[tuple[str, str, str]]]:
                    result = self._extractor.extract(text)
                    return result.entities, result.relations

            pipeline.add_extractor(LLMExtractorAdapter(provider, api_key, model))
        except ImportError:
            pass

    # Always add simple regex as fallback (works without dependencies)
    pipeline.add_extractor(SimpleRegexExtractor())
    return pipeline
