from __future__ import annotations

import json
from dataclasses import dataclass
from decimal import Decimal
from typing import Any, Dict

from .lexer import LexError, tokenise
from .parser import ParseError, Program, parse
from .verifier import VerificationSummary, VerifyError, verify


class CompileError(Exception):
    pass


@dataclass
class CompiledProof:
    program: Program
    summary: VerificationSummary

    def to_dict(self) -> Dict[str, Any]:
        return {
            "resources": list(self.program.resources),
            "states": {
                "before": _state_to_dict(self.program.states.get("before", {})),
                "after": _state_to_dict(self.program.states.get("after", {})),
            },
            "invariants": self.summary.invariants,
            "rules": [
                {
                    "name": r.name,
                    "status": r.status,
                    "expr": r.expr,
                    "value": r.value,
                }
                for r in self.summary.rules
            ],
            "proof": {
                "name": self.summary.proof_name,
                "verified_rules": list(self.summary.verified_rules),
                "status": "valid" if self.summary.valid else "invalid",
            },
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=2, sort_keys=True)


def compile_source(source: str) -> CompiledProof:
    try:
        tokens = tokenise(source)
        program = parse(tokens)
        summary = verify(program)
        return CompiledProof(program=program, summary=summary)
    except (LexError, ParseError, VerifyError) as exc:
        raise CompileError(str(exc)) from exc


def compile_file(path: str) -> CompiledProof:
    with open(path, "r", encoding="utf-8") as f:
        return compile_source(f.read())


def _state_to_dict(state: Dict[str, Dict[str, Decimal]]) -> Dict[str, Dict[str, str]]:
    return {
        party: {asset: _decimal_to_str(amount) for asset, amount in assets.items()}
        for party, assets in state.items()
    }


def _decimal_to_str(value: Decimal) -> str:
    text = format(value, "f")
    if "." in text:
        text = text.rstrip("0").rstrip(".")
    return text or "0"


# Backwards-friendly aliases
compile_swap = compile_source
