from __future__ import annotations

from dataclasses import dataclass
from decimal import Decimal
from typing import Dict, List, Optional, Tuple, Union

from .parser import (
    Bool,
    Binary,
    Expr,
    Number,
    Program,
    Ref,
    Rule,
    Unary,
)


class VerifyError(Exception):
    def __init__(self, message: str):
        super().__init__(f"VerifyError - {message}")


Value = Union[Decimal, bool]


@dataclass
class RuleResult:
    name: str
    status: str  # "passed" | "failed" | "external"
    value: Optional[bool]
    expr: Optional[str]


@dataclass
class VerificationSummary:
    valid: bool
    invariants: Dict[str, Dict[str, bool]]
    rules: List[RuleResult]
    proof_name: str
    verified_rules: List[str]


def _expr_to_str(expr: Expr) -> str:
    if isinstance(expr, Number):
        text = format(expr.value, "f")
        if "." in text:
            text = text.rstrip("0").rstrip(".")
        return text or "0"
    if isinstance(expr, Bool):
        return "true" if expr.value else "false"
    if isinstance(expr, Ref):
        return ".".join(expr.parts)
    if isinstance(expr, Unary):
        return f"{expr.op}{_expr_to_str(expr.right)}"
    if isinstance(expr, Binary):
        left = _expr_to_str(expr.left)
        right = _expr_to_str(expr.right)
        return f"({left} {expr.op} {right})"
    return "<?>"


def verify(program: Program) -> VerificationSummary:
    resources = program.resources
    if not resources:
        raise VerifyError("At least one Resource must be declared.")

    if "before" not in program.states or "after" not in program.states:
        raise VerifyError("Both State before and State after must be declared.")

    # Validate state assignments
    for state_name, parties in program.states.items():
        for party, assets in parties.items():
            for asset, amount in assets.items():
                if asset not in resources:
                    raise VerifyError(
                        f"State {state_name} references undeclared Resource '{asset}'."
                    )
                if amount < 0:
                    raise VerifyError(
                        f"Negative amount is not allowed: {party}.{asset} in {state_name}."
                    )

    # Invariant: conservation per resource
    invariants: Dict[str, Dict[str, bool]] = {"conservation": {}}
    for asset in resources:
        before_total = _total_for_resource(program.states["before"], asset)
        after_total = _total_for_resource(program.states["after"], asset)
        ok = before_total == after_total
        invariants["conservation"][asset] = ok
        if not ok:
            raise VerifyError(
                f"Conservation failed for {asset}: before {before_total} != after {after_total}."
            )

    # Proof verifies rules
    proof = program.proof
    for rule_name in proof.verifies:
        if rule_name not in program.rules:
            raise VerifyError(f"Proof references unknown Rule '{rule_name}'.")

    for rule_name in program.rules.keys():
        if rule_name not in proof.verifies:
            raise VerifyError(f"Rule '{rule_name}' is not verified in Proof.")

    rule_results: List[RuleResult] = []
    for rule in program.rules.values():
        if rule.external:
            rule_results.append(
                RuleResult(
                    name=rule.name,
                    status="external",
                    value=None,
                    expr=None,
                )
            )
            continue
        if rule.expr is None:
            raise VerifyError(f"Rule '{rule.name}' has no expression.")
        value = _eval_expr(rule.expr, program)
        if not isinstance(value, bool):
            raise VerifyError(f"Rule '{rule.name}' did not evaluate to a boolean.")
        if not value:
            raise VerifyError(f"Rule '{rule.name}' evaluated to false.")
        rule_results.append(
            RuleResult(
                name=rule.name,
                status="passed",
                value=value,
                expr=_expr_to_str(rule.expr),
            )
        )

    return VerificationSummary(
        valid=True,
        invariants=invariants,
        rules=rule_results,
        proof_name=proof.name,
        verified_rules=list(proof.verifies),
    )


def _total_for_resource(state: Dict[str, Dict[str, Decimal]], asset: str) -> Decimal:
    total = Decimal("0")
    for assets in state.values():
        if asset in assets:
            total += assets[asset]
    return total


def _eval_expr(expr: Expr, program: Program) -> Value:
    if isinstance(expr, Number):
        return expr.value
    if isinstance(expr, Bool):
        return expr.value
    if isinstance(expr, Ref):
        return _resolve_ref(expr.parts, program)
    if isinstance(expr, Unary):
        right = _eval_expr(expr.right, program)
        if not isinstance(right, Decimal):
            raise VerifyError("Unary operator expects a numeric value.")
        if expr.op == "+":
            return right
        if expr.op == "-":
            return -right
        raise VerifyError(f"Unknown unary operator {expr.op}")
    if isinstance(expr, Binary):
        left = _eval_expr(expr.left, program)
        right = _eval_expr(expr.right, program)
        if expr.op in ("+", "-", "*", "/"):
            if not isinstance(left, Decimal) or not isinstance(right, Decimal):
                raise VerifyError("Arithmetic operators require numeric operands.")
            if expr.op == "+":
                return left + right
            if expr.op == "-":
                return left - right
            if expr.op == "*":
                return left * right
            if expr.op == "/":
                if right == 0:
                    raise VerifyError("Division by zero.")
                return left / right
        if expr.op in ("==", "!=", "<", "<=", ">", ">="):
            if type(left) is not type(right):
                raise VerifyError("Comparison operands must have the same type.")
            if expr.op == "==":
                return left == right
            if expr.op == "!=":
                return left != right
            if expr.op == "<":
                return left < right
            if expr.op == "<=":
                return left <= right
            if expr.op == ">":
                return left > right
            if expr.op == ">=":
                return left >= right
        raise VerifyError(f"Unknown operator {expr.op}")
    raise VerifyError("Invalid expression.")


def _resolve_ref(parts: List[str], program: Program) -> Decimal:
    if not parts:
        raise VerifyError("Empty reference.")

    if parts[0] in ("before", "after"):
        state = program.states.get(parts[0], {})
        if len(parts) == 2 and parts[1] == "total":
            return _total_all(state)
        if len(parts) == 3 and parts[1] == "total":
            asset = parts[2]
            if asset not in program.resources:
                raise VerifyError(f"Unknown Resource '{asset}' in total.")
            return _total_for_resource(state, asset)
        if len(parts) == 3:
            party = parts[1]
            asset = parts[2]
            if asset not in program.resources:
                raise VerifyError(f"Unknown Resource '{asset}' in reference.")
            return state.get(party, {}).get(asset, Decimal("0"))
        raise VerifyError("Invalid state reference format.")

    raise VerifyError(f"Unknown reference '{'.'.join(parts)}'.")


def _total_all(state: Dict[str, Dict[str, Decimal]]) -> Decimal:
    total = Decimal("0")
    for assets in state.values():
        for amount in assets.values():
            total += amount
    return total
