from __future__ import annotations
from dataclasses import dataclass, field
import yaml


@dataclass
class Rule:
    id: str
    name: str
    source: str
    priority: int
    conditions: list[str]
    conclusions: list[str]
    dosha: str
    explanation: str


@dataclass
class ForwardChainResult:
    final_facts: set[str]
    fired_rules: list[dict]
    proof_trace: list[str]
    dosha_scores: dict[str, int]
    new_facts: set[str]


class ForwardChainingEngine:
    """
    Fixed-point forward chaining over Vrikshayurveda rules.

    Algorithm:
    - Each pass scans all unfired rules sorted by conflict resolution order.
    - A rule fires when ALL its conditions are in working memory.
    - Conclusions are added to working memory as new facts.
    - Iteration stops when no new rules fire in a full pass (fixed point).
    """

    def __init__(self, rules_path: str = "vrikshayurveda_rules.yaml") -> None:
        self.rules: list[Rule] = self._load_rules(rules_path)

    def _load_rules(self, path: str) -> list[Rule]:
        with open(path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)
        return [
            Rule(
                id=r["id"],
                name=r["name"],
                source=r["source"],
                priority=r["priority"],
                conditions=r["conditions"],
                conclusions=r["conclusions"],
                dosha=r["dosha"],
                explanation=r["explanation"],
            )
            for r in data["rules"]
        ]

    def _resolve_conflicts(self, fireable: list[Rule]) -> list[Rule]:
        """
        Sort fireable rules by:
        1. priority descending (explicit)
        2. len(conditions) descending (specificity)
        3. source == "vrikshayurveda" before "derived"
        """
        return sorted(
            fireable,
            key=lambda r: (
                r.priority,
                len(r.conditions),
                1 if r.source == "vrikshayurveda" else 0,
            ),
            reverse=True,
        )

    def run(self, facts: set[str]) -> ForwardChainResult:
        working_memory: set[str] = set(facts)
        fired_ids: set[str] = set()
        fired_rules: list[dict] = []
        proof_trace: list[str] = []
        dosha_scores: dict[str, int] = {"vata": 0, "pitta": 0, "kapha": 0}

        changed = True
        while changed:
            changed = False
            fireable = [
                r for r in self.rules
                if r.id not in fired_ids
                and all(c in working_memory for c in r.conditions)
            ]
            ordered = self._resolve_conflicts(fireable)

            for rule in ordered:
                new_facts = [c for c in rule.conclusions if c not in working_memory]
                working_memory.update(rule.conclusions)
                fired_ids.add(rule.id)
                fired_rules.append({
                    "rule_id": rule.id,
                    "name": rule.name,
                    "conditions": rule.conditions,
                    "conclusions": rule.conclusions,
                    "explanation": rule.explanation,
                    "dosha": rule.dosha,
                })

                cond_str = " ∧ ".join(rule.conditions)
                conc_str = " + ".join(rule.conclusions)
                proof_trace.append(
                    f"[{rule.id}] {cond_str} → {conc_str}"
                )

                if rule.dosha in dosha_scores:
                    dosha_scores[rule.dosha] += 1

                if new_facts:
                    changed = True

        return ForwardChainResult(
            final_facts=working_memory,
            fired_rules=fired_rules,
            proof_trace=proof_trace,
            dosha_scores=dosha_scores,
            new_facts=working_memory - facts,
        )