"""
Kolmogorov Engine — Skill System
==================================
A Skill is a pluggable domain pack that teaches the engine
new datatypes, functions, rules, and theorems.

The core engine never changes. Skills define:
  1. Types      — new datatypes with induction schemes
  2. Functions   — rewrite rules (definitions)
  3. Theorems    — properties to verify
  4. Vocabulary  — human-readable names for the chat UI
  5. Examples    — sample queries for onboarding

Skills are plain Python files. Drop one into skills/ and it's live.
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable, Optional
from nexus_gate_pkg.engine.terms import RewriteRule, Equation, Var, Const, App, Term
from nexus_gate_pkg.engine.proof import ProofState, Goal


@dataclass
class InductionScheme:
    """Defines how to do structural induction on a custom type."""
    type_name: str
    # Variable name hints (so the search engine knows which type to induct on)
    var_hints: set[str] = field(default_factory=set)
    # Base constructors: name → term (e.g. "Leaf" → Const("Leaf"))
    base_cases: dict[str, Callable] = field(default_factory=dict)
    # Recursive constructors: name → (term_builder, [recursive_arg_names])
    # e.g. "Node" → (lambda l,v,r: App("Node",[l,v,r]), ["l","r"])
    recursive_cases: dict[str, tuple] = field(default_factory=dict)


@dataclass
class TheoremDef:
    """A theorem that the engine can prove."""
    name: str
    description: str           # Human-readable
    lhs: Term
    rhs: Term
    context: list[Equation] = field(default_factory=list)
    difficulty: str = "medium"  # trivial, easy, medium, hard
    tags: list[str] = field(default_factory=list)


@dataclass
class Skill:
    """A pluggable domain pack."""
    id: str                          # e.g. "finance", "sets", "strings"
    name: str                        # e.g. "Financial Logic"
    description: str                 # One-liner for the UI
    version: str = "1.0.0"
    author: str = ""
    
    # Domain content
    definitions: list[RewriteRule] = field(default_factory=list)
    theorems: list[TheoremDef] = field(default_factory=list)
    induction_schemes: list[InductionScheme] = field(default_factory=list)
    
    # Vocabulary for the chat UI
    vocabulary: dict[str, str] = field(default_factory=dict)
    # e.g. {"add": "addition", "S": "successor", "Nil": "empty list"}
    
    # Example queries for onboarding
    examples: list[str] = field(default_factory=list)
    
    # Dependencies on other skills
    depends_on: list[str] = field(default_factory=list)
    
    def make_proof_state(self, thm: TheoremDef, all_defs=None) -> ProofState:
        """Create a proof state for a theorem in this skill."""
        defs = list(all_defs or self.definitions)
        return ProofState(
            goals=[Goal(thm.lhs, thm.rhs, thm.name)],
            context=list(thm.context),
            definitions=defs,
        )
    
    def to_dict(self) -> dict:
        """Serialize for the frontend."""
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "version": self.version,
            "author": self.author,
            "num_definitions": len(self.definitions),
            "num_theorems": len(self.theorems),
            "theorems": [
                {
                    "name": t.name,
                    "description": t.description,
                    "formula": f"{t.lhs} = {t.rhs}",
                    "difficulty": t.difficulty,
                    "tags": t.tags,
                }
                for t in self.theorems
            ],
            "vocabulary": self.vocabulary,
            "examples": self.examples,
            "depends_on": self.depends_on,
        }


class SkillRegistry:
    """Central registry — load, combine, and query skills."""
    
    def __init__(self):
        self._skills: dict[str, Skill] = {}
    
    def register(self, skill: Skill):
        self._skills[skill.id] = skill
    
    def get(self, skill_id: str) -> Optional[Skill]:
        return self._skills.get(skill_id)
    
    def list_skills(self) -> list[Skill]:
        return list(self._skills.values())
    
    def get_all_definitions(self, skill_ids: list[str] = None) -> list[RewriteRule]:
        """Merge definitions from multiple skills (with dependency resolution)."""
        if skill_ids is None:
            skill_ids = list(self._skills.keys())
        
        seen = set()
        defs = []
        
        def collect(sid):
            if sid in seen:
                return
            seen.add(sid)
            skill = self._skills.get(sid)
            if not skill:
                return
            for dep in skill.depends_on:
                collect(dep)
            defs.extend(skill.definitions)
        
        for sid in skill_ids:
            collect(sid)
        return defs
    
    def get_all_theorems(self, skill_ids: list[str] = None) -> list[tuple[Skill, TheoremDef]]:
        """Get all theorems, tagged with their source skill."""
        if skill_ids is None:
            skill_ids = list(self._skills.keys())
        
        result = []
        for sid in skill_ids:
            skill = self._skills.get(sid)
            if skill:
                for thm in skill.theorems:
                    result.append((skill, thm))
        return result
    
    def prove_all(self, skill_ids=None, beam_width=100, max_depth=40,
                  max_nodes=80000, timeout=15.0, verbose=False):
        """Prove all theorems across selected skills."""
        from nexus_gate_pkg.engine.search import beam_search
        
        all_defs = self.get_all_definitions(skill_ids)
        all_thms = self.get_all_theorems(skill_ids)
        results = []
        
        for skill, thm in all_thms:
            state = skill.make_proof_state(thm, all_defs=all_defs)
            trace = beam_search(state, beam_width=beam_width, max_depth=max_depth,
                              max_nodes=max_nodes, timeout=timeout, verbose=False)
            success = trace is not None and trace.success
            results.append({
                "skill": skill.id,
                "theorem": thm.name,
                "description": thm.description,
                "success": success,
                "steps": trace.length if success else 0,
                "nodes": trace.nodes_expanded if success else 0,
                "trace": trace if success else None,
            })
            if verbose:
                mark = "✓" if success else "✗"
                steps = trace.length if success else "-"
                print(f"  {mark} [{skill.id}] {thm.name}: {steps} steps")
        
        return results
    
    def to_dict(self) -> dict:
        return {
            "skills": [s.to_dict() for s in self._skills.values()],
            "total_theorems": sum(len(s.theorems) for s in self._skills.values()),
        }


# Global registry
REGISTRY = SkillRegistry()


def register_skill(skill: Skill):
    """Decorator-friendly registration."""
    REGISTRY.register(skill)
    return skill
