"""
Prolog-style reasoning engine for system optimization
Supports facts, rules, and queries with unification and backtracking
"""

from typing import Any, Dict, List, Optional, Tuple, Generator, Union
from dataclasses import dataclass, field
import re


class Variable:
    """Represents a Prolog variable (starts with uppercase)"""

    def __init__(self, name: str):
        self.name = name

    def __repr__(self):
        return f"Var({self.name})"

    def __eq__(self, other):
        return isinstance(other, Variable) and self.name == other.name

    def __hash__(self):
        return hash(("var", self.name))


@dataclass
class Fact:
    """A Prolog fact - a predicate with arguments"""
    predicate: str
    args: Tuple[Any, ...]

    def __init__(self, predicate: str, *args):
        self.predicate = predicate
        self.args = args

    def __repr__(self):
        if not self.args:
            return f"{self.predicate}"
        args_str = ", ".join(str(a) for a in self.args)
        return f"{self.predicate}({args_str})"


@dataclass
class Rule:
    """A Prolog rule - head :- body"""
    head: Fact
    body: List[Fact] = field(default_factory=list)

    def __repr__(self):
        if not self.body:
            return f"{self.head}."
        body_str = ", ".join(str(b) for b in self.body)
        return f"{self.head} :- {body_str}."


class PrologEngine:
    """
    A simple Prolog-style reasoning engine with:
    - Unification
    - Backtracking via generators
    - Facts and rules
    - Arithmetic comparisons
    """

    def __init__(self):
        self.facts: List[Fact] = []
        self.rules: List[Rule] = []

    def add_fact(self, predicate: str, *args) -> None:
        """Add a fact to the knowledge base"""
        self.facts.append(Fact(predicate, *args))

    def add_facts(self, facts: List[Tuple]) -> None:
        """Add multiple facts from tuples"""
        for fact in facts:
            if isinstance(fact, tuple) and len(fact) >= 1:
                self.add_fact(fact[0], *fact[1:])

    def add_rule(self, head: Fact, *body: Fact) -> None:
        """Add a rule to the knowledge base"""
        self.rules.append(Rule(head, list(body)))

    def clear(self) -> None:
        """Clear all facts and rules"""
        self.facts = []
        self.rules = []

    def unify(self, x: Any, y: Any, bindings: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Unify two terms, returning updated bindings or None if unification fails
        """
        # Apply current bindings
        x = self._deref(x, bindings)
        y = self._deref(y, bindings)

        # Same value
        if x == y:
            return bindings

        # Variable bindings
        if isinstance(x, Variable):
            return self._bind(x, y, bindings)
        if isinstance(y, Variable):
            return self._bind(y, x, bindings)

        # Facts/compound terms
        if isinstance(x, Fact) and isinstance(y, Fact):
            if x.predicate != y.predicate or len(x.args) != len(y.args):
                return None
            for ax, ay in zip(x.args, y.args):
                bindings = self.unify(ax, ay, bindings)
                if bindings is None:
                    return None
            return bindings

        # Lists
        if isinstance(x, (list, tuple)) and isinstance(y, (list, tuple)):
            if len(x) != len(y):
                return None
            for ax, ay in zip(x, y):
                bindings = self.unify(ax, ay, bindings)
                if bindings is None:
                    return None
            return bindings

        # Type coercion for numbers
        if isinstance(x, (int, float)) and isinstance(y, (int, float)):
            if float(x) == float(y):
                return bindings

        return None

    def _deref(self, term: Any, bindings: Dict[str, Any]) -> Any:
        """Dereference a variable through bindings"""
        while isinstance(term, Variable) and term.name in bindings:
            term = bindings[term.name]
        return term

    def _bind(self, var: Variable, value: Any, bindings: Dict[str, Any]) -> Dict[str, Any]:
        """Create new bindings with var bound to value"""
        new_bindings = dict(bindings)
        new_bindings[var.name] = value
        return new_bindings

    def query(self, goal: Fact) -> Generator[Dict[str, Any], None, None]:
        """
        Query the knowledge base, yielding all solutions
        """
        yield from self._solve([goal], {})

    def query_all(self, goal: Fact) -> List[Dict[str, Any]]:
        """Query and return all solutions as a list"""
        return list(self.query(goal))

    def query_one(self, goal: Fact) -> Optional[Dict[str, Any]]:
        """Query and return first solution or None"""
        for solution in self.query(goal):
            return solution
        return None

    def ask(self, predicate: str, *args) -> bool:
        """Simple yes/no query"""
        goal = Fact(predicate, *args)
        return self.query_one(goal) is not None

    def find(self, predicate: str, *args) -> List[Any]:
        """
        Find all values for variables in a query
        Returns list of binding dicts
        """
        goal = Fact(predicate, *args)
        return self.query_all(goal)

    def get_value(self, predicate: str, default: Any = None) -> Any:
        """Get the value of a single-argument fact"""
        for fact in self.facts:
            if fact.predicate == predicate and len(fact.args) == 1:
                return fact.args[0]
        return default

    def _solve(self, goals: List[Fact], bindings: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
        """Solve a list of goals with backtracking"""
        if not goals:
            yield bindings
            return

        goal = self._substitute(goals[0], bindings)
        remaining = goals[1:]

        # Handle built-in predicates
        if goal.predicate == ">=":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._eval_arithmetic(a, bindings)
                b = self._eval_arithmetic(b, bindings)
                if a is not None and b is not None and a >= b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "<=":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._eval_arithmetic(a, bindings)
                b = self._eval_arithmetic(b, bindings)
                if a is not None and b is not None and a <= b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == ">":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._eval_arithmetic(a, bindings)
                b = self._eval_arithmetic(b, bindings)
                if a is not None and b is not None and a > b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "<":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._eval_arithmetic(a, bindings)
                b = self._eval_arithmetic(b, bindings)
                if a is not None and b is not None and a < b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "==":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._deref(a, bindings)
                b = self._deref(b, bindings)
                if a == b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "\\==" or goal.predicate == "!=":
            if len(goal.args) == 2:
                a, b = goal.args
                a = self._deref(a, bindings)
                b = self._deref(b, bindings)
                if a != b:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "not" or goal.predicate == "\\+":
            if len(goal.args) == 1 and isinstance(goal.args[0], Fact):
                # Negation as failure
                inner_goal = goal.args[0]
                solutions = list(self._solve([inner_goal], bindings))
                if not solutions:
                    yield from self._solve(remaining, bindings)
            return

        if goal.predicate == "member":
            if len(goal.args) == 2:
                elem, lst = goal.args
                lst = self._deref(lst, bindings)
                if isinstance(lst, (list, tuple)):
                    for item in lst:
                        new_bindings = self.unify(elem, item, bindings)
                        if new_bindings is not None:
                            yield from self._solve(remaining, new_bindings)
            return

        # Try facts
        for fact in self.facts:
            new_bindings = self.unify(goal, fact, bindings)
            if new_bindings is not None:
                yield from self._solve(remaining, new_bindings)

        # Try rules
        for rule in self.rules:
            # Create fresh variables for the rule
            renamed_rule = self._rename_variables(rule)
            new_bindings = self.unify(goal, renamed_rule.head, bindings)
            if new_bindings is not None:
                yield from self._solve(renamed_rule.body + remaining, new_bindings)

    def _substitute(self, fact: Fact, bindings: Dict[str, Any]) -> Fact:
        """Substitute variables in a fact with their bindings"""
        new_args = tuple(self._deref(arg, bindings) for arg in fact.args)
        return Fact(fact.predicate, *new_args)

    def _eval_arithmetic(self, expr: Any, bindings: Dict[str, Any]) -> Optional[Union[int, float]]:
        """Evaluate an arithmetic expression"""
        expr = self._deref(expr, bindings)
        if isinstance(expr, (int, float)):
            return expr
        if isinstance(expr, Variable):
            return None
        return None

    _var_counter = 0

    def _rename_variables(self, rule: Rule) -> Rule:
        """Rename variables in a rule to avoid capture"""
        PrologEngine._var_counter += 1
        suffix = f"_{PrologEngine._var_counter}"

        mapping = {}

        def rename(term):
            if isinstance(term, Variable):
                if term.name not in mapping:
                    mapping[term.name] = Variable(term.name + suffix)
                return mapping[term.name]
            if isinstance(term, Fact):
                return Fact(term.predicate, *[rename(a) for a in term.args])
            if isinstance(term, (list, tuple)):
                return type(term)(rename(x) for x in term)
            return term

        new_head = rename(rule.head)
        new_body = [rename(b) for b in rule.body]
        return Rule(new_head, new_body)


# Convenience functions for creating terms
def var(name: str) -> Variable:
    """Create a variable"""
    return Variable(name)


def fact(predicate: str, *args) -> Fact:
    """Create a fact"""
    return Fact(predicate, *args)


def rule(head: Fact, *body: Fact) -> Rule:
    """Create a rule"""
    return Rule(head, list(body))


# Parser for Prolog-like syntax
def parse_term(text: str) -> Union[Fact, Variable, str, int, float]:
    """Parse a simple term from text"""
    text = text.strip()

    # Number
    if re.match(r'^-?\d+\.?\d*$', text):
        return float(text) if '.' in text else int(text)

    # Variable (starts with uppercase or _)
    if text[0].isupper() or text[0] == '_':
        return Variable(text)

    # Compound term/fact
    match = re.match(r'^(\w+)\s*\((.*)\)$', text)
    if match:
        predicate = match.group(1)
        args_str = match.group(2)
        # Simple split by comma (doesn't handle nested commas)
        args = [parse_term(a.strip()) for a in args_str.split(',')] if args_str else []
        return Fact(predicate, *args)

    # Atom (lowercase identifier or quoted string)
    if text.startswith('"') and text.endswith('"'):
        return text[1:-1]
    if text.startswith("'") and text.endswith("'"):
        return text[1:-1]

    return text


if __name__ == "__main__":
    # Demo
    engine = PrologEngine()

    # Add facts
    engine.add_fact("cpu_cores", 8)
    engine.add_fact("ram_gb", 16)
    engine.add_fact("gpu_tier", "high")
    engine.add_fact("task", "gaming")

    # Add rules
    X = var("X")
    Y = var("Y")

    engine.add_rule(
        fact("can_handle", "gaming"),
        fact("gpu_tier", "high"),
        fact(">=", var("R"), 8),
    )

    engine.add_rule(
        fact("recommend", "high_settings"),
        fact("can_handle", "gaming")
    )

    # Queries
    print("Facts in knowledge base:")
    for f in engine.facts:
        print(f"  {f}")

    print("\nRules in knowledge base:")
    for r in engine.rules:
        print(f"  {r}")

    print("\nQuery: gpu_tier(X)?")
    for solution in engine.query(fact("gpu_tier", var("X"))):
        print(f"  X = {solution.get('X')}")

    print("\nAsk: cpu_cores(8)?", engine.ask("cpu_cores", 8))
    print("Ask: cpu_cores(4)?", engine.ask("cpu_cores", 4))
