from __future__ import annotations

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

from .lexer import Token


class ParseError(Exception):
    def __init__(self, message: str, line: int, col: int):
        super().__init__(f"ParseError at {line}:{col} - {message}")
        self.line = line
        self.col = col


@dataclass(frozen=True)
class Expr:
    pass


@dataclass(frozen=True)
class Number(Expr):
    value: Decimal


@dataclass(frozen=True)
class Bool(Expr):
    value: bool


@dataclass(frozen=True)
class Ref(Expr):
    parts: List[str]


@dataclass(frozen=True)
class Unary(Expr):
    op: str
    right: Expr


@dataclass(frozen=True)
class Binary(Expr):
    op: str
    left: Expr
    right: Expr


@dataclass
class Rule:
    name: str
    expr: Optional[Expr]
    external: bool


@dataclass
class Proof:
    name: str
    verifies: List[str]


@dataclass
class Program:
    resources: List[str]
    states: Dict[str, Dict[str, Dict[str, Decimal]]]
    rules: Dict[str, Rule]
    proof: Proof


class Parser:
    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.pos = 0

    def _peek(self) -> Token:
        return self.tokens[self.pos]

    def _advance(self) -> Token:
        tok = self.tokens[self.pos]
        self.pos += 1
        return tok

    def _match(self, kind: str, value: Optional[str] = None) -> bool:
        tok = self._peek()
        if tok.kind != kind:
            return False
        if value is not None and tok.value != value:
            return False
        self._advance()
        return True

    def _expect(self, kind: str, value: Optional[str] = None) -> Token:
        tok = self._peek()
        if tok.kind != kind or (value is not None and tok.value != value):
            expected = f"{kind} {value}" if value else kind
            raise ParseError(f"Expected {expected}", tok.line, tok.col)
        return self._advance()

    def _skip_newlines(self) -> None:
        while self._match("NEWLINE"):
            pass

    def parse(self) -> Program:
        resources: List[str] = []
        states: Dict[str, Dict[str, Dict[str, Decimal]]] = {}
        rules: Dict[str, Rule] = {}
        proof: Optional[Proof] = None

        self._skip_newlines()
        while self._peek().kind != "EOF":
            tok = self._peek()
            if tok.kind == "KEYWORD" and tok.value in ("Resource", "sophiResource"):
                self._advance()
                name_tok = self._expect("IDENT")
                resources.append(name_tok.value)
                self._expect("NEWLINE")
            elif tok.kind == "KEYWORD" and tok.value == "State":
                state_name, mapping = self._parse_state_block()
                if state_name in states:
                    raise ParseError(f"Duplicate State '{state_name}'", tok.line, tok.col)
                states[state_name] = mapping
            elif tok.kind == "KEYWORD" and tok.value == "Rule":
                rule = self._parse_rule_block()
                if rule.name in rules:
                    raise ParseError(f"Duplicate Rule '{rule.name}'", tok.line, tok.col)
                rules[rule.name] = rule
            elif tok.kind == "KEYWORD" and tok.value == "Proof":
                if proof is not None:
                    raise ParseError("Only one Proof block is allowed", tok.line, tok.col)
                proof = self._parse_proof_block()
            elif tok.kind == "NEWLINE":
                self._advance()
            else:
                raise ParseError(f"Unexpected token {tok.kind}", tok.line, tok.col)
            self._skip_newlines()

        if proof is None:
            last = self.tokens[max(self.pos - 1, 0)]
            raise ParseError("Missing Proof block", last.line, last.col)

        return Program(resources=resources, states=states, rules=rules, proof=proof)

    def _parse_state_block(self):
        self._expect("KEYWORD", "State")
        name_tok = self._expect("IDENT")
        self._expect("OP", ":")
        self._expect("NEWLINE")
        self._expect("INDENT")
        mapping: Dict[str, Dict[str, Decimal]] = {}
        while not self._match("DEDENT"):
            if self._match("NEWLINE"):
                continue
            party_tok = self._expect("IDENT")
            self._expect("OP", ".")
            res_tok = self._expect("IDENT")
            self._expect("OP", "=")
            num_tok = self._expect("NUMBER")
            amount = Decimal(num_tok.value)
            mapping.setdefault(party_tok.value, {})[res_tok.value] = amount
            self._expect("NEWLINE")
        return name_tok.value, mapping

    def _parse_rule_block(self) -> Rule:
        self._expect("KEYWORD", "Rule")
        name_tok = self._expect("IDENT")
        self._expect("OP", ":")
        self._expect("NEWLINE")
        self._expect("INDENT")

        if self._match("KEYWORD", "external"):
            self._expect("NEWLINE")
            self._expect("DEDENT")
            return Rule(name=name_tok.value, expr=None, external=True)

        expr = self._parse_expr()
        self._expect("NEWLINE")
        self._expect("DEDENT")
        return Rule(name=name_tok.value, expr=expr, external=False)

    def _parse_proof_block(self) -> Proof:
        self._expect("KEYWORD", "Proof")
        name_tok = self._expect("IDENT")
        self._expect("OP", ":")
        self._expect("NEWLINE")
        self._expect("INDENT")
        verifies: List[str] = []
        while not self._match("DEDENT"):
            if self._match("NEWLINE"):
                continue
            self._expect("KEYWORD", "verify")
            self._expect("OP", "(")
            rule_tok = self._expect("IDENT")
            self._expect("OP", ")")
            self._expect("NEWLINE")
            verifies.append(rule_tok.value)
        return Proof(name=name_tok.value, verifies=verifies)

    def _parse_expr(self, min_bp: int = 0) -> Expr:
        tok = self._advance()
        if tok.kind == "NUMBER":
            left: Expr = Number(Decimal(tok.value))
        elif tok.kind == "KEYWORD" and tok.value in ("true", "false"):
            left = Bool(tok.value == "true")
        elif tok.kind in ("IDENT", "KEYWORD"):
            parts = [tok.value]
            while self._match("OP", "."):
                part_tok = self._expect("IDENT")
                parts.append(part_tok.value)
            left = Ref(parts)
        elif tok.kind == "OP" and tok.value in ("+", "-"):
            right = self._parse_expr(40)
            left = Unary(tok.value, right)
        elif tok.kind == "OP" and tok.value == "(":
            left = self._parse_expr()
            self._expect("OP", ")")
        else:
            raise ParseError("Invalid expression", tok.line, tok.col)

        while True:
            next_tok = self._peek()
            if next_tok.kind != "OP":
                break
            op = next_tok.value
            lbp, rbp = self._binding_power(op)
            if lbp < min_bp:
                break
            self._advance()
            right = self._parse_expr(rbp)
            left = Binary(op, left, right)
        return left

    def _binding_power(self, op: str):
        if op in ("==", "!=", "<", "<=", ">", ">="):
            return 10, 11
        if op in ("+", "-"):
            return 20, 21
        if op in ("*", "/"):
            return 30, 31
        return -1, -1


def parse(tokens: List[Token]) -> Program:
    return Parser(tokens).parse()
