"""
Kolmogorov Engine — Proof State & Primitives
==============================================
The proof environment for equational reasoning over recursive datatypes.

A proof state tracks:
  - goals: equations we need to prove
  - context: known facts (equations, lemmas)
  - definitions: recursive function definitions (rewrite rules)

Primitives operate on proof states and produce new proof states.
Every primitive is small, typed, and verifiable.
"""

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

from nexus_gate_pkg.engine.terms import (
    Term, TermKind, Var, Const, App,
    Substitution, EMPTY_SUB,
    unify, match, rewrite_once, rewrite_all, normalize,
    RewriteRule, Equation, free_vars, term_size, fresh_var,
)


# ══════════════════════════════════════════════════════════════
# Goal: an equation to prove
# ══════════════════════════════════════════════════════════════

@dataclass
class Goal:
    """An equation we need to prove: lhs = rhs."""
    lhs: Term
    rhs: Term
    name: str = ""

    def is_trivial(self) -> bool:
        """Goal is solved if lhs and rhs are syntactically equal."""
        return self.lhs == self.rhs

    def __repr__(self):
        prefix = f"[{self.name}] " if self.name else ""
        return f"{prefix}{self.lhs} = {self.rhs}"


# ══════════════════════════════════════════════════════════════
# Proof State
# ══════════════════════════════════════════════════════════════

@dataclass
class ProofState:
    """
    The complete state of a proof attempt.
    
    goals:       list of equations still to prove
    context:     known equations (axioms, IH, proven lemmas)
    definitions: rewrite rules from function definitions
    solved:      whether all goals are discharged
    """
    goals: list[Goal] = field(default_factory=list)
    context: list[Equation] = field(default_factory=list)
    definitions: list[RewriteRule] = field(default_factory=list)
    solved: bool = False

    def current_goal(self) -> Optional[Goal]:
        """The first unsolved goal."""
        for g in self.goals:
            if not g.is_trivial():
                return g
        return None

    def is_complete(self) -> bool:
        """All goals trivially true?"""
        return all(g.is_trivial() for g in self.goals)

    def clone(self) -> ProofState:
        return ProofState(
            goals=[Goal(g.lhs, g.rhs, g.name) for g in self.goals],
            context=list(self.context),
            definitions=list(self.definitions),
            solved=self.solved,
        )

    def __repr__(self):
        lines = []
        for i, g in enumerate(self.goals):
            mark = "✓" if g.is_trivial() else "?"
            lines.append(f"  {mark} Goal {i}: {g}")
        ctx = ", ".join(eq.name for eq in self.context)
        return f"ProofState({len(self.goals)} goals, ctx=[{ctx}])\n" + "\n".join(lines)


# ══════════════════════════════════════════════════════════════
# Proof Step (trace element)
# ══════════════════════════════════════════════════════════════

@dataclass
class ProofStep:
    """A single primitive application in a proof trace."""
    primitive: str          # Name of the primitive
    args: dict              # Arguments (for replay and abstraction)
    goal_before: Goal       # State before
    goal_after: list[Goal]  # State after (may be 0, 1, or multiple goals)
    cost: int = 1           # Cost of this step

    def __repr__(self):
        return f"{self.primitive}({self.args}) : {self.goal_before} ⟶ {self.goal_after}"


# ══════════════════════════════════════════════════════════════
# Primitives
# ══════════════════════════════════════════════════════════════

class PrimitiveResult:
    """Result of applying a primitive."""
    def __init__(self, state: ProofState, steps: list[ProofStep], success: bool = True):
        self.state = state
        self.steps = steps
        self.success = success

    @staticmethod
    def fail():
        return PrimitiveResult(ProofState(), [], success=False)


def _check_complete(state: ProofState) -> ProofState:
    """Update solved flag."""
    state.solved = state.is_complete()
    return state


# ─── Reflexivity ──────────────────────────────────────────────
def prim_refl(state: ProofState) -> PrimitiveResult:
    """
    If current goal is t = t, close it.
    This is the base case — the only way to actually close a goal.
    """
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()
    if not goal.is_trivial():
        return PrimitiveResult.fail()

    new_state = state.clone()
    step = ProofStep("refl", {}, goal, [])
    # Goal stays but is trivial, so it's "solved"
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Rewrite LHS ─────────────────────────────────────────────
def prim_rewrite_lhs(state: ProofState, rule: RewriteRule) -> PrimitiveResult:
    """
    Rewrite the LHS of the current goal using a rule.
    goal: f(x) = rhs  --[rule: f(x) → g(x)]--> g(x) = rhs
    """
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    new_lhs = rewrite_once(goal.lhs, rule)
    if new_lhs is None:
        return PrimitiveResult.fail()

    new_state = state.clone()
    new_goal = Goal(new_lhs, goal.rhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals[idx] = new_goal

    step = ProofStep("rewrite_lhs", {"rule": rule.name}, old_goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Rewrite RHS ─────────────────────────────────────────────
def prim_rewrite_rhs(state: ProofState, rule: RewriteRule) -> PrimitiveResult:
    """Rewrite the RHS of the current goal using a rule."""
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    new_rhs = rewrite_once(goal.rhs, rule)
    if new_rhs is None:
        return PrimitiveResult.fail()

    new_state = state.clone()
    new_goal = Goal(goal.lhs, new_rhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals[idx] = new_goal

    step = ProofStep("rewrite_rhs", {"rule": rule.name}, old_goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Normalize LHS ───────────────────────────────────────────
def prim_normalize_lhs(state: ProofState) -> PrimitiveResult:
    """Normalize LHS using all definitions."""
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    new_lhs = normalize(goal.lhs, state.definitions)
    if new_lhs == goal.lhs:
        return PrimitiveResult.fail()  # No change

    new_state = state.clone()
    new_goal = Goal(new_lhs, goal.rhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals[idx] = new_goal

    step = ProofStep("normalize_lhs", {}, old_goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Normalize RHS ───────────────────────────────────────────
def prim_normalize_rhs(state: ProofState) -> PrimitiveResult:
    """Normalize RHS using all definitions."""
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    new_rhs = normalize(goal.rhs, state.definitions)
    if new_rhs == goal.rhs:
        return PrimitiveResult.fail()

    new_state = state.clone()
    new_goal = Goal(goal.lhs, new_rhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals[idx] = new_goal

    step = ProofStep("normalize_rhs", {}, old_goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Symmetry ────────────────────────────────────────────────
def prim_symmetry(state: ProofState) -> PrimitiveResult:
    """Swap LHS and RHS of current goal."""
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    new_state = state.clone()
    new_goal = Goal(goal.rhs, goal.lhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals[idx] = new_goal

    step = ProofStep("symmetry", {}, old_goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Congruence ──────────────────────────────────────────────
def prim_congruence(state: ProofState) -> PrimitiveResult:
    """
    If goal is f(a1,..,an) = f(b1,..,bn), split into a1=b1, ..., an=bn.
    """
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    if (goal.lhs.kind != TermKind.APP or goal.rhs.kind != TermKind.APP):
        return PrimitiveResult.fail()
    if goal.lhs.name != goal.rhs.name:
        return PrimitiveResult.fail()
    if len(goal.lhs.args) != len(goal.rhs.args):
        return PrimitiveResult.fail()

    new_goals = []
    for i, (a, b) in enumerate(zip(goal.lhs.args, goal.rhs.args)):
        if a != b:
            new_goals.append(Goal(a, b, f"{goal.name}_arg{i}"))

    if not new_goals:
        # All args equal — goal is trivially true
        new_goals = []

    new_state = state.clone()
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    old_goal = new_state.goals[idx]
    new_state.goals = new_state.goals[:idx] + new_goals + new_state.goals[idx+1:]

    step = ProofStep("congruence", {}, old_goal, new_goals)
    return PrimitiveResult(_check_complete(new_state), [step])


# ─── Apply known equation ────────────────────────────────────
def prim_apply_equation(state: ProofState, eq_name: str, side: str = "lhs") -> PrimitiveResult:
    """
    Find an equation in context by name and apply it as a rewrite.
    """
    eq = None
    for e in state.context:
        if e.name == eq_name:
            eq = e
            break
    if eq is None:
        return PrimitiveResult.fail()

    rule = eq.forward()
    if side == "lhs":
        return prim_rewrite_lhs(state, rule)
    else:
        return prim_rewrite_rhs(state, rule)


# ─── Custom Induction Schemes ────────────────────────────────
# Skills register new types here. Each entry maps:
#   type_name → {
#     "base_constructor": str,         # e.g. "Empty", "Leaf"
#     "rec_constructor": str,          # e.g. "Position", "Grant"
#     "rec_arg_count": int,            # total args to the constructor
#     "rec_positions": [int],          # which args are recursive (get IHs)
#     "non_rec_positions": [int],      # which args are non-recursive constants
#     "var_hints": set[str],           # variable names that hint at this type
#   }

CUSTOM_INDUCTION: dict[str, dict] = {}

def register_induction_scheme(type_name: str, base_constructor: str,
                                rec_constructor: str, rec_arg_count: int,
                                rec_positions: list[int],
                                var_hints: set[str] = None):
    """Register a new induction scheme for a custom datatype."""
    non_rec = [i for i in range(rec_arg_count) if i not in rec_positions]
    CUSTOM_INDUCTION[type_name] = {
        "base_constructor": base_constructor,
        "rec_constructor": rec_constructor,
        "rec_arg_count": rec_arg_count,
        "rec_positions": rec_positions,
        "non_rec_positions": non_rec,
        "var_hints": var_hints or set(),
    }


# ─── Induction ───────────────────────────────────────────────
def prim_induction(state: ProofState, var_name: str, 
                   datatype: str = "Nat") -> PrimitiveResult:
    """
    Apply structural induction on a variable.
    
    For Nat: split into base case (var = Z) and step case (var = S(n)).
    For List: split into base case (var = Nil) and step case (var = Cons(h, t)).
    
    The induction hypothesis is added to context for the step case.
    """
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    var = Var(var_name)

    if datatype == "Nat":
        # Base case: substitute var with Z
        base_sub = Substitution({var_name: Const("Z")})
        base_goal = Goal(
            base_sub.apply(goal.lhs),
            base_sub.apply(goal.rhs),
            f"{goal.name}_base"
        )

        # Step case: substitute var with S(n) where n is a LOCAL CONSTANT
        # n_pred is ground — it's a specific value, not a universally quantified variable.
        # This is critical: the IH must only match the exact predecessor term,
        # not arbitrary expressions. Using Const ensures pattern matching treats
        # it as ground (only matches itself).
        n = Const(f"{var_name}_pred")
        step_sub = Substitution({var_name: App("S", [n])})
        step_goal = Goal(
            step_sub.apply(goal.lhs),
            step_sub.apply(goal.rhs),
            f"{goal.name}_step"
        )

        # Induction hypothesis: the goal holds for n (ground)
        ih_sub = Substitution({var_name: n})
        ih = Equation(
            f"IH_{var_name}",
            ih_sub.apply(goal.lhs),
            ih_sub.apply(goal.rhs),
        )

        new_state = state.clone()
        idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
        new_state.goals = (
            new_state.goals[:idx] + 
            [base_goal, step_goal] + 
            new_state.goals[idx+1:]
        )
        new_state.context = new_state.context + [ih]

        step = ProofStep("induction", {"var": var_name, "type": "Nat"}, 
                        goal, [base_goal, step_goal])
        return PrimitiveResult(_check_complete(new_state), [step])

    elif datatype == "List":
        # Base case: var = Nil
        base_sub = Substitution({var_name: Const("Nil")})
        base_goal = Goal(
            base_sub.apply(goal.lhs),
            base_sub.apply(goal.rhs),
            f"{goal.name}_base"
        )

        # Step case: var = Cons(h, t) where h, t are LOCAL CONSTANTS
        h = Const(f"{var_name}_head")
        t = Const(f"{var_name}_tail")
        step_sub = Substitution({var_name: App("Cons", [h, t])})
        step_goal = Goal(
            step_sub.apply(goal.lhs),
            step_sub.apply(goal.rhs),
            f"{goal.name}_step"
        )

        # IH on the tail (ground)
        ih_sub = Substitution({var_name: t})
        ih = Equation(
            f"IH_{var_name}",
            ih_sub.apply(goal.lhs),
            ih_sub.apply(goal.rhs),
        )

        new_state = state.clone()
        idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
        new_state.goals = (
            new_state.goals[:idx] + 
            [base_goal, step_goal] + 
            new_state.goals[idx+1:]
        )
        new_state.context = new_state.context + [ih]

        step = ProofStep("induction", {"var": var_name, "type": "List"},
                        goal, [base_goal, step_goal])
        return PrimitiveResult(_check_complete(new_state), [step])

    elif datatype == "Tree":
        # Base case: var = Leaf
        base_sub = Substitution({var_name: Const("Leaf")})
        base_goal = Goal(
            base_sub.apply(goal.lhs),
            base_sub.apply(goal.rhs),
            f"{goal.name}_base"
        )

        # Step case: var = Node(l, v, r) where l, v, r are LOCAL CONSTANTS
        l_child = Const(f"{var_name}_left")
        v_val = Const(f"{var_name}_val")
        r_child = Const(f"{var_name}_right")
        step_sub = Substitution({var_name: App("Node", [l_child, v_val, r_child])})
        step_goal = Goal(
            step_sub.apply(goal.lhs),
            step_sub.apply(goal.rhs),
            f"{goal.name}_step"
        )

        # TWO induction hypotheses: one for left subtree, one for right
        ih_l_sub = Substitution({var_name: l_child})
        ih_l = Equation(
            f"IH_{var_name}_l",
            ih_l_sub.apply(goal.lhs),
            ih_l_sub.apply(goal.rhs),
        )
        ih_r_sub = Substitution({var_name: r_child})
        ih_r = Equation(
            f"IH_{var_name}_r",
            ih_r_sub.apply(goal.lhs),
            ih_r_sub.apply(goal.rhs),
        )

        new_state = state.clone()
        idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
        new_state.goals = (
            new_state.goals[:idx] +
            [base_goal, step_goal] +
            new_state.goals[idx+1:]
        )
        new_state.context = new_state.context + [ih_l, ih_r]

        step = ProofStep("induction", {"var": var_name, "type": "Tree"},
                        goal, [base_goal, step_goal])
        return PrimitiveResult(_check_complete(new_state), [step])

    elif datatype in CUSTOM_INDUCTION:
        scheme = CUSTOM_INDUCTION[datatype]
        
        # Base case: var = BaseConstructor (e.g. Empty, NoPerms)
        base_sub = Substitution({var_name: Const(scheme["base_constructor"])})
        base_goal = Goal(
            base_sub.apply(goal.lhs),
            base_sub.apply(goal.rhs),
            f"{goal.name}_base"
        )
        
        # Step case: var = RecConstructor(arg0, arg1, ...)
        # Recursive args get local constants + IHs
        # Non-recursive args get local constants only
        args = []
        for i in range(scheme["rec_arg_count"]):
            args.append(Const(f"{var_name}_arg{i}"))
        
        step_term = App(scheme["rec_constructor"], args)
        step_sub = Substitution({var_name: step_term})
        step_goal = Goal(
            step_sub.apply(goal.lhs),
            step_sub.apply(goal.rhs),
            f"{goal.name}_step"
        )
        
        # IHs for recursive positions
        ihs = []
        for pos in scheme["rec_positions"]:
            ih_sub = Substitution({var_name: args[pos]})
            ih = Equation(
                f"IH_{var_name}_{pos}",
                ih_sub.apply(goal.lhs),
                ih_sub.apply(goal.rhs),
            )
            ihs.append(ih)
        
        new_state = state.clone()
        idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
        new_state.goals = (
            new_state.goals[:idx] +
            [base_goal, step_goal] +
            new_state.goals[idx+1:]
        )
        new_state.context = new_state.context + ihs
        
        step = ProofStep("induction", {"var": var_name, "type": datatype},
                        goal, [base_goal, step_goal])
        return PrimitiveResult(_check_complete(new_state), [step])

    return PrimitiveResult.fail()


# ─── Focus ───────────────────────────────────────────────────
def prim_focus(state: ProofState, goal_idx: int) -> PrimitiveResult:
    """Move a specific goal to the front (make it current)."""
    if goal_idx < 0 or goal_idx >= len(state.goals):
        return PrimitiveResult.fail()

    new_state = state.clone()
    goal = new_state.goals.pop(goal_idx)
    new_state.goals.insert(0, goal)

    step = ProofStep("focus", {"idx": goal_idx}, goal, [goal])
    return PrimitiveResult(new_state, [step])


# ─── Substitute known equality ────────────────────────────────
def prim_substitute(state: ProofState, eq_name: str) -> PrimitiveResult:
    """
    Use a known equation to substitute in current goal.
    If context has `a = b`, replace all occurrences of a with b in goal.
    """
    goal = state.current_goal()
    if goal is None:
        return PrimitiveResult.fail()

    eq = None
    for e in state.context:
        if e.name == eq_name:
            eq = e
            break
    if eq is None:
        return PrimitiveResult.fail()

    rule = eq.forward()
    new_lhs = rewrite_all(goal.lhs, rule)
    new_rhs = rewrite_all(goal.rhs, rule)

    if new_lhs == goal.lhs and new_rhs == goal.rhs:
        return PrimitiveResult.fail()

    new_state = state.clone()
    new_goal = Goal(new_lhs, new_rhs, goal.name)
    idx = next(i for i, g in enumerate(new_state.goals) if not g.is_trivial())
    new_state.goals[idx] = new_goal

    step = ProofStep("substitute", {"eq": eq_name}, goal, [new_goal])
    return PrimitiveResult(_check_complete(new_state), [step])


# ═══════════════════════════════════════════════════════════════
# Macro replay — execute a learned macro on a new proof state
# ═══════════════════════════════════════════════════════════════

def replay_macro(macro_name: str, macro_steps: list[ProofStep],
                 state: ProofState) -> PrimitiveResult:
    """
    Try to replay a macro's step sequence on a new proof state.
    Each step is re-executed (not just copied) so it must actually succeed.
    Returns the final state if ALL steps succeed, otherwise fails.
    """
    current = state
    collected_steps = []

    for template_step in macro_steps:
        prim = template_step.primitive
        args = template_step.args

        result = None

        if prim == "induction":
            # Adapt: find the right variable to induct on
            goal = current.current_goal()
            if goal is None:
                return PrimitiveResult.fail()
            from nexus_gate_pkg.engine.terms import free_vars
            goal_vars = free_vars(goal.lhs) | free_vars(goal.rhs)

            dtype = args.get("type", "Nat")
            # Try each variable with the same datatype
            for var_name in sorted(goal_vars):
                r = prim_induction(current, var_name, dtype)
                if r.success:
                    result = r
                    break

        elif prim == "normalize_lhs":
            result = prim_normalize_lhs(current)

        elif prim == "normalize_rhs":
            result = prim_normalize_rhs(current)

        elif prim == "rewrite_lhs":
            rule_name = args.get("rule", "")
            if rule_name.startswith("IH_"):
                # Find any IH in context and apply it
                for eq in current.context:
                    if eq.name.startswith("IH_"):
                        fwd = eq.forward()
                        r = prim_rewrite_lhs(current, fwd)
                        if r.success:
                            result = r
                            break
            else:
                # Try the exact rule from definitions
                for defn in current.definitions:
                    if defn.name == rule_name:
                        r = prim_rewrite_lhs(current, defn)
                        if r.success:
                            result = r
                            break

        elif prim == "rewrite_rhs":
            rule_name = args.get("rule", "")
            if rule_name.startswith("IH_"):
                for eq in current.context:
                    if eq.name.startswith("IH_"):
                        bwd = eq.backward()
                        r = prim_rewrite_rhs(current, bwd)
                        if r.success:
                            result = r
                            break
            else:
                for defn in current.definitions:
                    if defn.name == rule_name:
                        r = prim_rewrite_rhs(current, defn)
                        if r.success:
                            result = r
                            break

        elif prim == "symmetry":
            result = prim_symmetry(current)

        elif prim == "congruence":
            result = prim_congruence(current)

        else:
            return PrimitiveResult.fail()

        if result is None or not result.success:
            return PrimitiveResult.fail()

        current = result.state
        collected_steps.extend(result.steps)

    # All steps succeeded — wrap as single macro step
    goal_before = state.current_goal()
    macro_step = ProofStep(
        primitive=f"MACRO:{macro_name}",
        args={"macro": macro_name},
        goal_before=goal_before,
        goal_after=[g for g in current.goals if not g.is_trivial()],
        cost=1,
    )
    return PrimitiveResult(_check_complete(current), [macro_step])


# ══════════════════════════════════════════════════════════════
# Registry of all primitives
# ══════════════════════════════════════════════════════════════

PRIMITIVES = {
    "refl": prim_refl,
    "rewrite_lhs": prim_rewrite_lhs,
    "rewrite_rhs": prim_rewrite_rhs,
    "normalize_lhs": prim_normalize_lhs,
    "normalize_rhs": prim_normalize_rhs,
    "symmetry": prim_symmetry,
    "congruence": prim_congruence,
    "apply_equation": prim_apply_equation,
    "induction": prim_induction,
    "focus": prim_focus,
    "substitute": prim_substitute,
}
