"""
cogniq.plugins — Extensible plugin system.

Allows anyone to hook into the RAG pipeline:
  pre_add      → transform text before chunking
  post_chunk   → modify/filter chunks
  pre_embed    → modify text before embedding  
  post_search  → rerank/filter results
  post_answer  → post-process generated answer

Usage:
    from cogniq.plugins import Plugin, register_plugin

    @register_plugin("my_reranker")
    class MyReranker(Plugin):
        def post_search(self, results, query, **kw):
            # rerank results by custom logic
            return sorted(results, key=lambda r: my_score(r, query), reverse=True)

    rag = RAG(plugins=["my_reranker"])
"""

from __future__ import annotations
import logging
from typing import List, Tuple, Dict, Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    pass

log = logging.getLogger("cogniq")

# ─────────────────────────────────────────────
# REGISTRY (singleton)
# ─────────────────────────────────────────────

_REGISTRY: Dict[str, type] = {}


def register_plugin(name: str):
    """Decorator to register a plugin class by name."""
    def decorator(cls):
        _REGISTRY[name] = cls
        log.debug(f"Plugin registered: {name} → {cls.__name__}")
        return cls
    return decorator


def get_plugin(name: str) -> Optional[type]:
    return _REGISTRY.get(name)


def list_plugins() -> List[str]:
    return list(_REGISTRY.keys())


# ─────────────────────────────────────────────
# BASE PLUGIN
# ─────────────────────────────────────────────

class Plugin:
    """
    Base plugin — override any hook you need.
    All hooks are optional; default implementation is a no-op pass-through.
    """
    name: str = "base_plugin"

    def pre_add(self, text: str, meta: Dict[str, Any]) -> Tuple[str, Dict]:
        """Called before text is chunked. Return (text, meta)."""
        return text, meta

    def post_chunk(self, chunks: List[Dict], meta: Dict[str, Any]) -> List[Dict]:
        """Called after chunking. Return modified chunk list."""
        return chunks

    def pre_embed(self, text: str) -> str:
        """Called before embedding each chunk. Return text."""
        return text

    def post_search(
        self,
        results: List[Tuple[Dict, float]],
        query: str,
        **kwargs,
    ) -> List[Tuple[Dict, float]]:
        """Called after vector search. Return (possibly reranked) results."""
        return results

    def post_answer(self, answer: str, query: str, context: str, **kwargs) -> str:
        """Called after answer generation. Return modified answer."""
        return answer

    def __repr__(self):
        return f"Plugin({self.name})"


# ─────────────────────────────────────────────
# BUILT-IN PLUGINS
# ─────────────────────────────────────────────

@register_plugin("dedup")
class DedupPlugin(Plugin):
    """
    Remove near-duplicate chunks from search results.
    Uses Jaccard similarity on word sets.
    """
    name = "dedup"

    def __init__(self, threshold: float = 0.85):
        self.threshold = threshold

    def _jaccard(self, a: str, b: str) -> float:
        sa, sb = set(a.lower().split()), set(b.lower().split())
        if not sa or not sb:
            return 0.0
        return len(sa & sb) / len(sa | sb)

    def post_search(self, results, query, **kw):
        filtered = []
        for doc, score in results:
            text = doc.get("text", "")
            is_dup = any(
                self._jaccard(text, d.get("text", "")) > self.threshold
                for d, _ in filtered
            )
            if not is_dup:
                filtered.append((doc, score))
        return filtered


@register_plugin("min_score")
class MinScorePlugin(Plugin):
    """Filter out results below a minimum similarity threshold."""
    name = "min_score"

    def __init__(self, threshold: float = 0.3):
        self.threshold = threshold

    def post_search(self, results, query, **kw):
        return [(d, s) for d, s in results if s >= self.threshold]


@register_plugin("source_filter")
class SourceFilterPlugin(Plugin):
    """Only return results from specific sources."""
    name = "source_filter"

    def __init__(self, allowed_sources: List[str]):
        self.allowed = set(allowed_sources)

    def post_search(self, results, query, **kw):
        return [
            (d, s) for d, s in results
            if d.get("source") in self.allowed
        ]


@register_plugin("text_cleaner")
class TextCleanerPlugin(Plugin):
    """Clean text before adding: strip HTML, normalize whitespace."""
    name = "text_cleaner"

    def pre_add(self, text, meta):
        import re
        text = re.sub(r'<[^>]+>', ' ', text)        # strip HTML
        text = re.sub(r'\s+', ' ', text).strip()     # normalize whitespace
        text = re.sub(r'[^\x00-\x7F]+', ' ', text)  # remove non-ASCII
        return text, meta


@register_plugin("metadata_injector")
class MetadataInjectorPlugin(Plugin):
    """Inject fixed metadata into every document."""
    name = "metadata_injector"

    def __init__(self, **meta):
        self.meta = meta

    def pre_add(self, text, meta):
        meta.update(self.meta)
        return text, meta


# ─────────────────────────────────────────────
# PLUGIN CHAIN
# ─────────────────────────────────────────────

class PluginChain:
    """Runs a list of plugins in order for each hook."""

    def __init__(self, plugins: List[Plugin]):
        self.plugins = plugins

    def pre_add(self, text, meta):
        for p in self.plugins:
            text, meta = p.pre_add(text, meta)
        return text, meta

    def post_chunk(self, chunks, meta):
        for p in self.plugins:
            chunks = p.post_chunk(chunks, meta)
        return chunks

    def pre_embed(self, text):
        for p in self.plugins:
            text = p.pre_embed(text)
        return text

    def post_search(self, results, query, **kwargs):
        for p in self.plugins:
            results = p.post_search(results, query, **kwargs)
        return results

    def post_answer(self, answer, query, context, **kwargs):
        for p in self.plugins:
            answer = p.post_answer(answer, query, context, **kwargs)
        return answer

    @classmethod
    def build(cls, specs) -> "PluginChain":
        """
        Build chain from list of: Plugin instances, registered names, or callables.
        
        Examples:
            PluginChain.build(["dedup", "min_score"])
            PluginChain.build([DedupPlugin(0.9), my_plugin_instance])
        """
        plugins = []
        for spec in (specs or []):
            if isinstance(spec, Plugin):
                plugins.append(spec)
            elif isinstance(spec, str):
                cls_ = get_plugin(spec)
                if cls_:
                    plugins.append(cls_())
                else:
                    log.warning(f"Plugin '{spec}' not found in registry. Available: {list_plugins()}")
            elif callable(spec):
                # Wrap a bare function as a post_search plugin
                class _FuncPlugin(Plugin):
                    def post_search(self_, results, query, **kw):
                        return spec(results, query, **kw)
                plugins.append(_FuncPlugin())
        return cls(plugins)
