"""
Belief Revision System - Mesh and Similarity

Cross-domain pattern matching using pluggable similarity metrics.
The mesh enables discovery of analogies between domains.
"""
from __future__ import annotations

import re
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import (
    Any, Callable, Dict, FrozenSet, Iterable, List, Optional,
    Protocol, Sequence, Set, Tuple, TypeVar
)

from .storage import CASStore


# =============================================================================
# Tokenization
# =============================================================================

TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
ROLE_RE = re.compile(r"\b[A-Z][A-Z0-9_]{2,}\b")


def tokenize(text: str) -> Set[str]:
    """Extract lowercase tokens from text."""
    return {t.lower() for t in TOKEN_RE.findall(text or "")}


def extract_roles(text: str) -> Set[str]:
    """
    Extract role identifiers (ALLCAPS tokens).

    These often represent domain-specific entities like:
    - ARAMAIC, BRAHMI in writing systems
    - HOMO_SAPIENS in biology
    - BROWN_V_BOARD in legal cases
    """
    return set(ROLE_RE.findall(text or ""))


# =============================================================================
# Pattern Signatures
# =============================================================================

@dataclass(frozen=True)
class PatternSignature:
    """
    Domain-agnostic signature for cross-domain matching.

    Captures the "shape" of a pattern through:
    - tokens: Key terms from the pattern definition
    - roles: Domain-specific entity references
    - operators: Transformation or constraint operators
    """
    domain_id: str
    pattern_id: str
    name: str
    tokens: FrozenSet[str]
    roles: FrozenSet[str]
    operators: FrozenSet[str]

    @classmethod
    def from_pattern(cls, domain_id: str, pattern_json: Dict[str, Any]) -> "PatternSignature":
        """Build a signature from a pattern dict."""
        name = pattern_json.get("name", "")
        summary = pattern_json.get("summary", "")
        invariants = " ".join(pattern_json.get("invariants", []) or [])
        operators = pattern_json.get("operators", []) or []

        # Combine text sources
        blob = " ".join([name, summary, invariants, " ".join(operators)])

        # Extract components
        toks = tokenize(blob)
        ops = {o.lower() for o in operators}
        roles = extract_roles(" ".join([name, summary, invariants]))

        return cls(
            domain_id=domain_id,
            pattern_id=pattern_json.get("id", ""),
            name=name,
            tokens=frozenset(toks),
            roles=frozenset(roles),
            operators=frozenset(ops),
        )


# =============================================================================
# Similarity Metrics (Pluggable)
# =============================================================================

class SimilarityMetric(Protocol):
    """Protocol for similarity metrics between pattern signatures."""

    def __call__(self, a: PatternSignature, b: PatternSignature) -> float:
        """
        Compute similarity score between two signatures.

        Returns:
            Float in [0, 1] where 1 is identical
        """
        ...


def jaccard(a: Set[str], b: Set[str]) -> float:
    """Jaccard similarity coefficient."""
    if not a and not b:
        return 1.0
    inter = len(a.intersection(b))
    union = len(a.union(b))
    return inter / union if union else 0.0


@dataclass
class WeightedJaccardSimilarity:
    """
    Weighted Jaccard similarity across tokens, operators, and roles.

    Default weights favor tokens (general semantics) over operators
    (structural similarity) over roles (domain-specific entities).
    """
    token_weight: float = 0.65
    operator_weight: float = 0.25
    role_weight: float = 0.10

    def __call__(self, a: PatternSignature, b: PatternSignature) -> float:
        tok = jaccard(set(a.tokens), set(b.tokens))
        op = jaccard(set(a.operators), set(b.operators))
        rl = jaccard(set(a.roles), set(b.roles))
        return (
            self.token_weight * tok +
            self.operator_weight * op +
            self.role_weight * rl
        )


# Default similarity metric
DEFAULT_SIMILARITY = WeightedJaccardSimilarity()


@dataclass
class CosineSimilarity:
    """
    Cosine similarity using term frequency vectors.

    Alternative to Jaccard for cases where term frequency matters.
    """

    def _term_freq(self, sig: PatternSignature) -> Dict[str, float]:
        """Build term frequency vector from all signature components."""
        terms: Dict[str, float] = {}
        for t in sig.tokens:
            terms[f"tok:{t}"] = terms.get(f"tok:{t}", 0) + 1
        for o in sig.operators:
            terms[f"op:{o}"] = terms.get(f"op:{o}", 0) + 1
        for r in sig.roles:
            terms[f"role:{r}"] = terms.get(f"role:{r}", 0) + 1
        return terms

    def __call__(self, a: PatternSignature, b: PatternSignature) -> float:
        va = self._term_freq(a)
        vb = self._term_freq(b)
        all_keys = set(va.keys()) | set(vb.keys())

        dot = sum(va.get(k, 0) * vb.get(k, 0) for k in all_keys)
        mag_a = sum(v**2 for v in va.values()) ** 0.5
        mag_b = sum(v**2 for v in vb.values()) ** 0.5

        if mag_a == 0 or mag_b == 0:
            return 0.0
        return dot / (mag_a * mag_b)


# =============================================================================
# Analog Proposals
# =============================================================================

@dataclass(frozen=True)
class AnalogProposal:
    """
    A proposed cross-domain analogy between two patterns.

    Analog proposals suggest that knowledge from one domain can
    inform belief revision in another domain.
    """
    left_domain: str
    left_pattern_id: str
    left_name: str
    right_domain: str
    right_pattern_id: str
    right_name: str
    score: float
    shared_tokens: Tuple[str, ...]
    shared_operators: Tuple[str, ...]


def propose_analogs(
    signatures: Sequence[PatternSignature],
    top_k: int = 20,
    min_score: float = 0.35,
    similarity: SimilarityMetric = DEFAULT_SIMILARITY
) -> List[AnalogProposal]:
    """
    Generate cross-domain analog proposals from pattern signatures.

    Args:
        signatures: All patterns to consider
        top_k: Maximum proposals to return
        min_score: Minimum similarity threshold
        similarity: Similarity metric to use

    Returns:
        Top-scoring analog proposals (cross-domain only)
    """
    props: List[AnalogProposal] = []
    sigs = list(signatures)

    for i in range(len(sigs)):
        for j in range(i + 1, len(sigs)):
            a, b = sigs[i], sigs[j]

            # Skip same-domain comparisons
            if a.domain_id == b.domain_id:
                continue

            s = similarity(a, b)
            if s < min_score:
                continue

            shared_tokens = tuple(sorted(set(a.tokens).intersection(b.tokens)))[:40]
            shared_ops = tuple(sorted(set(a.operators).intersection(b.operators)))

            props.append(AnalogProposal(
                left_domain=a.domain_id,
                left_pattern_id=a.pattern_id,
                left_name=a.name,
                right_domain=b.domain_id,
                right_pattern_id=b.pattern_id,
                right_name=b.name,
                score=s,
                shared_tokens=shared_tokens,
                shared_operators=shared_ops
            ))

    props.sort(key=lambda x: x.score, reverse=True)
    return props[:top_k]


# =============================================================================
# Mesh Index
# =============================================================================

@dataclass(frozen=True)
class MeshIndex:
    """
    Index of pattern signatures across domains.

    The mesh enables efficient cross-domain analogy discovery.
    """
    signatures: Tuple[PatternSignature, ...]

    def find_similar(
        self,
        target: PatternSignature,
        top_k: int = 10,
        min_score: float = 0.3,
        similarity: SimilarityMetric = DEFAULT_SIMILARITY,
        exclude_same_domain: bool = True
    ) -> List[Tuple[PatternSignature, float]]:
        """
        Find signatures similar to a target.

        Args:
            target: Pattern to match against
            top_k: Maximum results
            min_score: Minimum similarity threshold
            similarity: Similarity metric
            exclude_same_domain: If True, exclude same-domain matches

        Returns:
            List of (signature, score) tuples
        """
        results: List[Tuple[PatternSignature, float]] = []

        for sig in self.signatures:
            if exclude_same_domain and sig.domain_id == target.domain_id:
                continue
            if sig.pattern_id == target.pattern_id:
                continue

            s = similarity(target, sig)
            if s >= min_score:
                results.append((sig, s))

        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]

    def domains(self) -> Set[str]:
        """Get all domains in the mesh."""
        return {s.domain_id for s in self.signatures}

    def patterns_for_domain(self, domain_id: str) -> List[PatternSignature]:
        """Get all signatures for a domain."""
        return [s for s in self.signatures if s.domain_id == domain_id]


def build_mesh_index(
    store: CASStore,
    domains: Sequence[Tuple[str, str]]  # (domain_id, world_label) pairs
) -> MeshIndex:
    """
    Build a mesh index from domain worlds.

    Args:
        store: Storage backend
        domains: List of (domain_id, version_label) pairs to include

    Returns:
        MeshIndex containing signatures for all patterns in specified worlds
    """
    # Scan all Pattern objects
    cur = store._conn.cursor()
    rows = cur.execute("SELECT json FROM objects WHERE kind='Pattern'").fetchall()
    import json as _json
    all_patterns = [_json.loads(r[0]) for r in rows]
    by_id = {p.get("id"): p for p in all_patterns if p.get("id")}

    sigs: List[PatternSignature] = []

    for domain_id, world_label in domains:
        try:
            world = store.get_world(domain_id, world_label)["json"]
        except KeyError:
            continue

        for pid in world.get("pattern_ids", []) or []:
            p = by_id.get(pid)
            if p:
                sigs.append(PatternSignature.from_pattern(domain_id, p))

    return MeshIndex(signatures=tuple(sigs))
