"""
Lemma Library — Proven theorems become rewrite rules for future proofs.

A lemma is a proven equation: lhs = rhs
It can be applied as a rewrite rule in either direction:
  - forward:  lhs → rhs
  - backward: rhs → lhs

The library accumulates lemmas as the engine proves theorems.
Each new proof has access to ALL previously proven lemmas.
"""

from dataclasses import dataclass, field
from nexus_gate_pkg.engine.terms import Term, RewriteRule, Equation, free_vars, match

@dataclass
class Lemma:
    """A proven theorem stored for reuse."""
    name: str
    lhs: Term       # Left side of proven equation
    rhs: Term       # Right side of proven equation
    domain: str     # Which skill/domain proved it
    complexity: int  # How many steps the original proof took
    
    def forward_rule(self) -> RewriteRule:
        """Use as lhs → rhs rewrite."""
        return RewriteRule(f"lemma:{self.name}→", self.lhs, self.rhs)
    
    def backward_rule(self) -> RewriteRule:
        """Use as rhs → lhs rewrite."""
        return RewriteRule(f"lemma:{self.name}←", self.rhs, self.lhs)
    
    def as_equation(self) -> Equation:
        return Equation(self.name, self.lhs, self.rhs)


class LemmaLibrary:
    """
    Accumulates proven theorems for use in future proofs.
    
    Usage:
        lib = LemmaLibrary()
        
        # After proving length(append(A,B)) = length(A) + length(B):
        lib.add("append_length", lhs, rhs, domain="list", complexity=6)
        
        # When proving reverse:
        extra_rules = lib.get_applicable_rules(goal)
        # → includes the append_length lemma as a rewrite rule
    """
    
    def __init__(self):
        self.lemmas: list[Lemma] = []
        self._by_name: dict[str, Lemma] = {}
        self._by_head: dict[str, list[Lemma]] = {}  # index by top function symbol
        self._universal: list[Lemma] = []  # always-included lemmas (arithmetic etc)
    
    def add(self, name: str, lhs: Term, rhs: Term, 
            domain: str = "", complexity: int = 0) -> Lemma:
        """Add a proven theorem to the library."""
        lemma = Lemma(name, lhs, rhs, domain, complexity)
        self.lemmas.append(lemma)
        self._by_name[name] = lemma
        
        # Index by head symbol of lhs and rhs for fast lookup
        for term in [lhs, rhs]:
            head = self._head_symbol(term)
            if head:
                if head not in self._by_head:
                    self._by_head[head] = []
                self._by_head[head].append(lemma)
        
        return lemma
    
    def _head_symbol(self, term: Term) -> str:
        """Extract the top-level function symbol."""
        from nexus_gate_pkg.engine.terms import TermKind
        if term.kind == TermKind.APP:
            return term.name
        elif term.kind == TermKind.CONST:
            return term.name
        return ""
    
    def get_all_rules(self) -> list[RewriteRule]:
        """Get all lemmas as rewrite rules (both directions)."""
        rules = []
        for lem in self.lemmas:
            rules.append(lem.forward_rule())
            rules.append(lem.backward_rule())
        return rules
    
    def mark_universal(self, name: str):
        """Mark a lemma as universal — always included in relevant rules."""
        if name in self._by_name:
            self._universal.append(self._by_name[name])
    
    def get_relevant_rules(self, goal_lhs: Term, goal_rhs: Term, 
                           max_rules: int = 20) -> list[RewriteRule]:
        """
        Get lemma rules that are likely relevant to a goal.
        
        Relevance heuristic: a lemma is relevant if its head symbols
        appear in the goal. Universal lemmas are always included.
        """
        # Collect all function symbols in the goal
        goal_symbols = set()
        self._collect_symbols(goal_lhs, goal_symbols)
        self._collect_symbols(goal_rhs, goal_symbols)
        
        # Find lemmas whose head symbols overlap with goal symbols
        relevant = []
        seen = set()
        
        # Always include universals first
        for lem in self._universal:
            if lem.name not in seen:
                seen.add(lem.name)
                relevant.append(lem)
        
        for sym in goal_symbols:
            for lem in self._by_head.get(sym, []):
                if lem.name not in seen:
                    seen.add(lem.name)
                    relevant.append(lem)
        
        # Sort by complexity (simpler lemmas first — more likely to be useful)
        relevant.sort(key=lambda l: l.complexity)
        
        # Convert to rules (both directions)
        rules = []
        for lem in relevant[:max_rules]:
            rules.append(lem.forward_rule())
            rules.append(lem.backward_rule())
        
        return rules
    
    def _collect_symbols(self, term: Term, symbols: set):
        """Recursively collect all function symbols in a term."""
        from nexus_gate_pkg.engine.terms import TermKind
        if term.kind == TermKind.APP:
            symbols.add(term.name)
            for arg in term.args:
                self._collect_symbols(arg, symbols)
        elif term.kind == TermKind.CONST:
            symbols.add(term.name)
    
    def size(self) -> int:
        return len(self.lemmas)
    
    def __repr__(self):
        return f"LemmaLibrary({len(self.lemmas)} lemmas, {len(self._by_head)} head symbols indexed)"
