from __future__ import annotations
from dataclasses import dataclass, field
import yaml
from inference.forward_chain import Rule


@dataclass
class BackwardChainResult:
    goal: str
    proved: bool
    proof_tree: dict
    proof_trace: list[str]
    confidence: float


@dataclass
class DifferentialResult:
    rankings: list[dict]
    primary: str
    primary_confidence: float
    trace: list[str]


class BackwardChainingEngine:
    """
    Recursive backward chaining for differential dosha diagnosis.

    Algorithm:
    - To prove a goal: check if it is already a known fact.
    - If not, find all rules whose conclusions contain the goal.
    - Recursively attempt to prove each condition of those rules.
    - A rule fires if ALL its conditions are proved.
    - Returns a proof tree showing exactly how the goal was established.
    """

    def __init__(self, rules_path: str = "vrikshayurveda_rules.yaml") -> None:
        raw_rules = self._load_rules(rules_path)
        self.rules = raw_rules
        # goal_index: conclusion → list of rules that can prove it
        self.goal_index: dict[str, list[Rule]] = {}
        for rule in raw_rules:
            for conc in rule.conclusions:
                self.goal_index.setdefault(conc, []).append(rule)

    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 prove(
        self,
        goal: str,
        known_facts: set[str],
        depth: int = 0,
        max_depth: int = 6,
        visited: set[str] | None = None,
    ) -> BackwardChainResult:
        if visited is None:
            visited = set()

        trace: list[str] = []
        indent = "  " * depth

        # Base case: goal already in working memory
        if goal in known_facts:
            trace.append(f"{indent}✓ '{goal}' is a known fact")
            return BackwardChainResult(
                goal=goal,
                proved=True,
                proof_tree={"goal": goal, "grounded": True},
                proof_trace=trace,
                confidence=1.0,
            )

        # Depth limit guard
        if depth >= max_depth or goal in visited:
            trace.append(f"{indent}✗ Cannot prove '{goal}' (depth limit or cycle)")
            return BackwardChainResult(
                goal=goal,
                proved=False,
                proof_tree={"goal": goal, "grounded": False},
                proof_trace=trace,
                confidence=0.0,
            )

        visited = visited | {goal}
        candidate_rules = sorted(
            self.goal_index.get(goal, []),
            key=lambda r: (r.priority, len(r.conditions)),
            reverse=True,
        )

        trace.append(f"{indent}? Trying to prove '{goal}'")

        for rule in candidate_rules:
            trace.append(f"{indent}  Trying rule {rule.id}: {rule.name}")
            sub_proofs: list[dict] = []
            all_proved = True
            conditions_proved = 0

            for cond in rule.conditions:
                sub_result = self.prove(
                    cond, known_facts, depth + 1, max_depth, visited
                )
                trace.extend(sub_result.proof_trace)
                sub_proofs.append(sub_result.proof_tree)
                if sub_result.proved:
                    conditions_proved += 1
                else:
                    all_proved = False
                    break

            if all_proved:
                trace.append(
                    f"{indent}  ✓ Rule {rule.id} fires → '{goal}' proved"
                )
                return BackwardChainResult(
                    goal=goal,
                    proved=True,
                    proof_tree={
                        "goal": goal,
                        "rule_used": rule.id,
                        "rule_name": rule.name,
                        "sub_proofs": sub_proofs,
                    },
                    proof_trace=trace,
                    confidence=1.0,
                )

        trace.append(f"{indent}✗ '{goal}' cannot be proved from known facts")
        return BackwardChainResult(
            goal=goal,
            proved=False,
            proof_tree={"goal": goal, "grounded": False},
            proof_trace=trace,
            confidence=0.0,
        )

    def _partial_confidence(self, goal: str, known_facts: set[str]) -> float:
        """
        Score how strongly the known facts support a goal,
        even if the goal cannot be fully proved.
        Returns fraction of best-matching rule's conditions that are satisfied.
        """
        if goal in known_facts:
            return 1.0
        candidates = self.goal_index.get(goal, [])
        if not candidates:
            return 0.0
        best = max(
            sum(1 for c in rule.conditions if c in known_facts) / len(rule.conditions)
            for rule in candidates
        )
        return round(best, 2)

    def differential_diagnosis(self, known_facts: set[str]) -> DifferentialResult:
        """
        Attempt to prove each primary dosha hypothesis.
        Score each by confidence and rank descending.
        """
        hypotheses = [
            ("vata", "vata_disorder"),
            ("pitta", "pitta_imbalance"),
            ("kapha", "kapha_obstruction"),
        ]
        rankings: list[dict] = []
        all_trace: list[str] = []

        for dosha, goal in hypotheses:
            result = self.prove(goal, known_facts)
            all_trace.extend(result.proof_trace)
            confidence = (
                result.confidence
                if result.proved
                else self._partial_confidence(goal, known_facts)
            )
            rankings.append({
                "dosha": dosha,
                "goal": goal,
                "proved": result.proved,
                "confidence": confidence,
                "proof_tree": result.proof_tree,
            })

        rankings.sort(key=lambda x: x["confidence"], reverse=True)
        primary = rankings[0]

        return DifferentialResult(
            rankings=rankings,
            primary=primary["dosha"],
            primary_confidence=primary["confidence"],
            trace=all_trace,
        )