"""
Kolmogorov Engine — Term Algebra
=================================
The foundation. Everything else builds on this.

Terms are the objects we reason about:
  - Variables: x, y, xs, ys
  - Constants: Z (zero), Nil (empty list)
  - Applications: S(Z), Cons(x, Nil), append(xs, ys)

This is a standard first-order term algebra with:
  - Substitution
  - Unification (Robinson's algorithm)
  - Pattern matching (one-directional unification)
  - Rewriting (apply equations as directed rules)
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum, auto


# ══════════════════════════════════════════════════════════════
# Terms
# ══════════════════════════════════════════════════════════════

class TermKind(Enum):
    VAR = auto()
    CONST = auto()
    APP = auto()


@dataclass(frozen=True)
class Term:
    """
    A first-order term.
    
    Var("x")           — a variable
    Const("Z")         — a constant (nullary constructor)
    App("S", [Const("Z")])  — application: S(Z)
    App("append", [xs, ys]) — application: append(xs, ys)
    """
    kind: TermKind
    name: str
    args: tuple[Term, ...] = ()

    def __repr__(self):
        if self.kind == TermKind.VAR:
            return f"?{self.name}"
        elif self.kind == TermKind.CONST:
            return self.name
        elif not self.args:
            return self.name
        else:
            args_str = ", ".join(repr(a) for a in self.args)
            return f"{self.name}({args_str})"

    def __hash__(self):
        return hash((self.kind, self.name, self.args))

    def __eq__(self, other):
        if not isinstance(other, Term):
            return False
        return self.kind == other.kind and self.name == other.name and self.args == other.args


def Var(name: str) -> Term:
    return Term(TermKind.VAR, name)

def Const(name: str) -> Term:
    return Term(TermKind.CONST, name)

def App(name: str, args: list[Term]) -> Term:
    return Term(TermKind.APP, name, tuple(args))


# ══════════════════════════════════════════════════════════════
# Substitutions
# ══════════════════════════════════════════════════════════════

class Substitution:
    """A mapping from variable names to terms."""

    def __init__(self, mapping: Optional[dict[str, Term]] = None):
        self.mapping: dict[str, Term] = mapping or {}

    def apply(self, term: Term) -> Term:
        """Apply this substitution to a term (single pass, no chaining)."""
        if term.kind == TermKind.VAR:
            return self.mapping.get(term.name, term)
        elif term.kind == TermKind.CONST:
            return term
        else:  # APP
            new_args = tuple(self.apply(a) for a in term.args)
            if new_args == term.args:
                return term
            return Term(TermKind.APP, term.name, new_args)

    def compose(self, other: Substitution) -> Substitution:
        """Compose two substitutions: (self ∘ other)(t) = self(other(t))"""
        result = {}
        for var, term in other.mapping.items():
            result[var] = self.apply(term)
        for var, term in self.mapping.items():
            if var not in result:
                result[var] = term
        return Substitution(result)

    def __repr__(self):
        items = [f"?{k} ↦ {v}" for k, v in self.mapping.items()]
        return "{" + ", ".join(items) + "}"

    def __bool__(self):
        return True  # Empty sub is still valid; use None for failure


EMPTY_SUB = Substitution()


# ══════════════════════════════════════════════════════════════
# Free variables
# ══════════════════════════════════════════════════════════════

def free_vars(term: Term) -> set[str]:
    """Return all variable names in a term."""
    if term.kind == TermKind.VAR:
        return {term.name}
    elif term.kind == TermKind.CONST:
        return set()
    else:
        result = set()
        for a in term.args:
            result |= free_vars(a)
        return result


def occurs_in(var_name: str, term: Term) -> bool:
    """Occurs check: does variable appear in term?"""
    if term.kind == TermKind.VAR:
        return term.name == var_name
    elif term.kind == TermKind.CONST:
        return False
    else:
        return any(occurs_in(var_name, a) for a in term.args)


# ══════════════════════════════════════════════════════════════
# Unification (Robinson's algorithm)
# ══════════════════════════════════════════════════════════════

def unify(t1: Term, t2: Term, sub: Optional[Substitution] = None) -> Optional[Substitution]:
    """
    Find a most general unifier for t1 and t2.
    Returns None if no unifier exists.
    """
    sub = sub or Substitution()
    t1 = sub.apply(t1)
    t2 = sub.apply(t2)

    if t1 == t2:
        return sub

    if t1.kind == TermKind.VAR:
        return _unify_var(t1.name, t2, sub)

    if t2.kind == TermKind.VAR:
        return _unify_var(t2.name, t1, sub)

    if t1.kind == TermKind.CONST and t2.kind == TermKind.CONST:
        return sub if t1.name == t2.name else None

    if t1.kind == TermKind.APP and t2.kind == TermKind.APP:
        if t1.name != t2.name or len(t1.args) != len(t2.args):
            return None
        for a1, a2 in zip(t1.args, t2.args):
            sub = unify(a1, a2, sub)
            if sub is None:
                return None
        return sub

    return None


def _unify_var(var_name: str, term: Term, sub: Substitution) -> Optional[Substitution]:
    """Unify a variable with a term."""
    if term.kind == TermKind.VAR and term.name == var_name:
        return sub
    if occurs_in(var_name, term):
        return None  # Occurs check failure
    new_mapping = dict(sub.mapping)
    new_mapping[var_name] = term
    return Substitution(new_mapping)


# ══════════════════════════════════════════════════════════════
# Pattern matching (one-directional unification)
# ══════════════════════════════════════════════════════════════

def match(pattern: Term, target: Term, sub: Optional[Substitution] = None) -> Optional[Substitution]:
    """
    Match pattern against target.
    Only pattern variables get bound. Target is ground (or treated as ground).
    Returns None if no match.
    """
    sub = sub or Substitution()

    if pattern.kind == TermKind.VAR:
        existing = sub.mapping.get(pattern.name)
        if existing is not None:
            return sub if existing == target else None
        new_mapping = dict(sub.mapping)
        new_mapping[pattern.name] = target
        return Substitution(new_mapping)

    if pattern.kind == TermKind.CONST:
        if target.kind == TermKind.CONST and pattern.name == target.name:
            return sub
        return None

    # APP
    if target.kind != TermKind.APP:
        return None
    if pattern.name != target.name or len(pattern.args) != len(target.args):
        return None
    for p_arg, t_arg in zip(pattern.args, target.args):
        sub = match(p_arg, t_arg, sub)
        if sub is None:
            return None
    return sub


# ══════════════════════════════════════════════════════════════
# Rewriting
# ══════════════════════════════════════════════════════════════

@dataclass(frozen=True)
class RewriteRule:
    """A directed equation: lhs → rhs (under matching)."""
    name: str
    lhs: Term
    rhs: Term
    conditions: tuple[Term, ...] = ()  # Side conditions (for future use)

    def __repr__(self):
        return f"{self.name}: {self.lhs} → {self.rhs}"


def rewrite_at_root(term: Term, rule: RewriteRule) -> Optional[Term]:
    """Try to apply a rewrite rule at the root of a term."""
    sub = match(rule.lhs, term)
    if sub is None:
        return None
    return sub.apply(rule.rhs)


def rewrite_once(term: Term, rule: RewriteRule) -> Optional[Term]:
    """
    Apply a rewrite rule at the first position where it matches.
    Leftmost-outermost strategy.
    Returns None if no match anywhere.
    """
    # Try at root
    result = rewrite_at_root(term, rule)
    if result is not None:
        return result

    # Try in subterms
    if term.kind == TermKind.APP:
        for i, arg in enumerate(term.args):
            new_arg = rewrite_once(arg, rule)
            if new_arg is not None:
                new_args = list(term.args)
                new_args[i] = new_arg
                return Term(TermKind.APP, term.name, tuple(new_args))

    return None


def rewrite_all(term: Term, rule: RewriteRule) -> Term:
    """Apply a rewrite rule everywhere until no more matches."""
    while True:
        new_term = rewrite_once(term, rule)
        if new_term is None:
            return term
        term = new_term


def normalize(term: Term, rules: list[RewriteRule], max_steps: int = 100) -> Term:
    """Normalize a term by applying all rules until fixpoint."""
    for _ in range(max_steps):
        changed = False
        for rule in rules:
            new_term = rewrite_once(term, rule)
            if new_term is not None:
                term = new_term
                changed = True
                break
        if not changed:
            return term
    return term  # Max steps reached


# ══════════════════════════════════════════════════════════════
# Term size (for compression metrics)
# ══════════════════════════════════════════════════════════════

def term_size(term: Term) -> int:
    """Count nodes in a term tree."""
    if term.kind in (TermKind.VAR, TermKind.CONST):
        return 1
    return 1 + sum(term_size(a) for a in term.args)


# ══════════════════════════════════════════════════════════════
# Fresh variable generation
# ══════════════════════════════════════════════════════════════

_fresh_counter = 0

def fresh_var(prefix: str = "v") -> Term:
    global _fresh_counter
    _fresh_counter += 1
    return Var(f"{prefix}_{_fresh_counter}")

def reset_fresh():
    global _fresh_counter
    _fresh_counter = 0


# ══════════════════════════════════════════════════════════════
# Convenience: Equation (bidirectional)
# ══════════════════════════════════════════════════════════════

@dataclass(frozen=True)
class Equation:
    """An equation: lhs = rhs. Can be used as rewrite rule in either direction."""
    name: str
    lhs: Term
    rhs: Term

    def forward(self) -> RewriteRule:
        return RewriteRule(self.name + "→", self.lhs, self.rhs)

    def backward(self) -> RewriteRule:
        return RewriteRule(self.name + "←", self.rhs, self.lhs)

    def __repr__(self):
        return f"{self.name}: {self.lhs} = {self.rhs}"
