import io
from collections import defaultdict
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field, replace
from functools import cmp_to_key
from inspect import signature
from typing import Any, Literal, Protocol, Self, Unpack, overload

import z3

from common import config
from common.relations import Relation, is_relation
from fest.z3_utils import find_order_relation

type Node = str
type NodeRef = int | Node


@dataclass(eq=True, frozen=True, slots=True)
class SymElement:
    node: Node
    sort: z3.SortRef
    index: z3.ArithRef | None = None

    @classmethod
    def int(cls, index: z3.ArithRef | None = None) -> Self:
        if index is None:
            index = z3.Int("i")

        return cls("int", z3.IntSort(), index)

    @classmethod
    def create(cls, node: Node, sort: z3.SortRef, is_summary: bool) -> Self:
        if is_summary:
            return cls(node, sort, z3.Int("i"))
        else:
            return cls(node, sort)

    def is_regular(self) -> bool:
        return self.index is None

    def is_summary(self) -> bool:
        return self.index is not None

    def to_index(self) -> z3.ArithRef:
        if self.index is None:
            return z3.IntVal(0)
        else:
            return self.index

    def with_index(self, index: z3.ArithRef) -> Self:
        return replace(self, index=index)

    def __str__(self) -> str:
        if self.is_regular():
            return self.node
        else:
            return f"<{self.node}, {self.index}>"


class Assignment(dict[str, SymElement]):
    def override(
        self, variables: Iterable[z3.Const], values: Iterable[SymElement]
    ) -> Self:
        overrides = zip((str(var) for var in variables), values)
        overridden_assignment = self.__class__(self)
        overridden_assignment.update(overrides)
        return overridden_assignment


type BoundsFormula = Callable[[z3.ArithRef], z3.BoolRef]

type Fun = Literal["fun"]
type Rel = Literal["rel"]


type _Term0 = Callable[[], z3.ArithRef]
type _Term1 = Callable[[z3.ArithRef], z3.ArithRef]
type _Term2 = Callable[[z3.ArithRef, z3.ArithRef], z3.ArithRef]
type _AnyTerm = _Term0 | _Term1 | _Term2

type _Formula0 = Callable[[], z3.BoolRef]
type _Formula1 = Callable[[z3.ArithRef], z3.BoolRef]
type _Formula2 = Callable[[z3.ArithRef, z3.ArithRef], z3.BoolRef]
type _Formula3 = Callable[[z3.ArithRef, z3.ArithRef, z3.ArithRef], z3.BoolRef]
type _AnyFormula = _Formula0 | _Formula1 | _Formula2 | _Formula3


class SymTerm(Protocol):
    def __call__(self, *indices: z3.ArithRef) -> z3.ArithRef: ...


class SymFormula(Protocol):
    def __call__(self, *indices: z3.ArithRef) -> z3.BoolRef: ...


def sym_term(fun: _AnyTerm) -> SymTerm:
    arity = len(signature(fun).parameters)

    def wrapper(*indices: z3.ArithRef) -> z3.ArithRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity: expected {arity} got {len(indices)}"
        return fun(*indices)

    return wrapper


def sym_formula(fun: _AnyFormula) -> SymFormula:
    arity = len(signature(fun).parameters)

    def wrapper(*indices: z3.ArithRef) -> z3.BoolRef:
        assert (
            len(indices) == arity
        ), f"Mismatch arity: expected {arity} got {len(indices)}"
        return fun(*indices)

    return wrapper


def auto_node(sort: z3.SortRef, node: NodeRef) -> Node:
    if isinstance(node, int):
        return f"{sort}_{node}"

    return node


def _default_domain() -> dict[z3.SortRef, set[Node]]:
    domain: dict[z3.SortRef, set[Node]] = defaultdict(set)
    if config.fest.experimental.int_support(bool):
        domain[z3.IntSort()] = {SymElement.int().node}

    return domain


def _default_bounds() -> dict[Node, BoundsFormula]:
    bounds = {}
    if config.fest.experimental.int_support(bool):
        bounds[SymElement.int().node] = lambda i: z3.BoolVal(True)

    return bounds


@dataclass(kw_only=True)
class SymStruct:
    domain: dict[z3.SortRef, set[Node]] = field(default_factory=_default_domain)
    bounds: dict[Node, BoundsFormula] = field(default_factory=_default_bounds)
    functions: dict[
        tuple[z3.FuncDeclRef, Unpack[tuple[Node, ...]]], tuple[Node, SymTerm]
    ] = field(default_factory=dict)
    relations: dict[Relation, dict[tuple[Node, ...], SymFormula]] = field(
        default_factory=lambda: defaultdict(dict)
    )
    constants: dict[z3.Const, SymElement] = field(default_factory=dict)

    @property
    def elements(self) -> dict[z3.SortRef, set[SymElement]]:
        return {
            sort: {SymElement.create(node, sort, node in self.bounds) for node in nodes}
            for sort, nodes in self.domain.items()
        }

    @overload
    def __getitem__(self, item: z3.SortRef) -> frozenset[Node]: ...

    @overload
    def __getitem__(self, item: Node) -> BoundsFormula | None: ...

    @overload
    def __getitem__(
        self, item: tuple[Rel, Relation, Unpack[tuple[Node, ...]]]
    ) -> SymFormula: ...

    @overload
    def __getitem__(
        self, item: tuple[Fun, z3.FuncDeclRef, Unpack[tuple[Node, ...]]]
    ) -> tuple[Node, SymTerm]: ...

    def __getitem__(self, item: Any) -> Any:
        if z3.is_sort(item):
            return self.domain[item]
        elif not isinstance(item, tuple):
            return self.bounds.get(item, None)

        kind, decl, *nodes = item

        if kind == "rel" and is_relation(decl):
            return self.relations[decl][tuple(nodes)]
        elif kind == "fun":
            return self.functions[(decl, *nodes)]

        raise KeyError(item)

    @overload
    def __setitem__(self, key: z3.SortRef, value: int | set[Node]) -> None: ...

    @overload
    def __setitem__(self, key: Node, value: BoundsFormula) -> None: ...

    @overload
    def __setitem__(
        self, key: tuple[z3.SortRef, int], value: BoundsFormula
    ) -> None: ...

    @overload
    def __setitem__(
        self,
        key: tuple[Fun, z3.FuncDeclRef, Unpack[tuple[NodeRef, ...]]],
        value: tuple[NodeRef, _AnyTerm],
    ) -> None: ...

    @overload
    def __setitem__(
        self,
        key: tuple[Rel, Relation, Unpack[tuple[NodeRef, ...]]],
        value: _AnyFormula,
    ) -> None: ...

    def __setitem__(self, key: Any, value: Any) -> None:
        if z3.is_sort(key):
            if isinstance(value, int):
                self.domain[key] = {auto_node(key, i) for i in range(value)}
            else:
                self.domain[key] = value
            return
        elif not isinstance(key, tuple):
            self.bounds[key] = value
            return

        head, *rest = key

        if z3.is_sort(head):
            (i,) = rest
            self[auto_node(head, i)] = value
            return

        decl, *nodes = rest

        assert decl.arity() == len(
            nodes
        ), f"Mismatched arity, expected {decl.arity()} but got {len(nodes)}"

        nodes = [auto_node(decl.domain(i), node) for i, node in enumerate(nodes)]

        if head == "rel" and is_relation(decl):
            self.relations[decl][tuple(nodes)] = sym_formula(value)
            return
        elif head == "fun":
            node, term = value
            node = auto_node(decl.range(), node)

            self.functions[(decl, *nodes)] = (node, sym_term(term))
            return

        assert False, f"Unsupported key {key}"

    def __str__(self) -> str:
        buffer = io.StringIO()
        for sort, nodes in self.domain.items():
            if sort.eq(z3.IntSort()):
                continue

            print(
                f"{sort}: {", ".join(map(lambda n: self._node_to_str(n, show_bound=True), nodes))}",
                file=buffer,
            )

        for constant, element in self.constants.items():
            print(f"{constant}: {element}", file=buffer)

        for (fun, *args), (node, term) in self.functions.items():
            indices = _indices_for_decl(fun)
            expanded_args = [
                self._node_to_str(arg, index=str(index))
                for arg, index in zip(args, indices)
            ]
            print(
                f"{fun}({", ".join(expanded_args)}) -> <{node}, {term(*indices)}>",
                file=buffer,
            )

        order_relation = find_order_relation(self.relations.keys())
        if config.fest.experimental.explain_order and order_relation is not None:
            self._show_order_relation(order_relation, buffer)

        for rel, defs in self.relations.items():
            if (
                not config.fest.experimental.show_order_interpretation
                and rel is order_relation
            ):
                continue

            for rel_args, formula in defs.items():
                indices = _indices_for_decl(rel)
                expanded_args = [
                    self._node_to_str(arg, index=str(index))
                    for arg, index in zip(rel_args, indices)
                ]
                print(
                    f"{rel}({", ".join(expanded_args)}): {formula(*indices)}",
                    file=buffer,
                )

        return buffer.getvalue()

    def _show_order_relation(
        self, order_relation: Relation, buffer: io.StringIO
    ) -> None:
        sort = order_relation.domain(0)
        i1, i2 = z3.Ints("i1 i2")

        def cmp_nodes(n1: Node, n2: Node) -> int:
            cmp_formula = self.relations[order_relation][n1, n2](i1, i2)
            if (
                z3.is_true(cmp_formula)
                or cmp_formula.eq(i1 <= i2)
                or cmp_formula.eq(i2 >= i1)
            ):
                return -1
            elif (
                z3.is_false(cmp_formula)
                or cmp_formula.eq(i1 < i2)
                or cmp_formula.eq(i2 > i1)
            ):
                return 1
            else:
                return 0

        ordered_nodes = sorted(self.domain[sort], key=cmp_to_key(cmp_nodes))
        last: Node | None = None
        print(f"{order_relation}: ", end="", file=buffer)
        for node in ordered_nodes:
            if last is not None:
                order_formula = self.relations[order_relation][last, node](i1, i2)
                if z3.is_true(order_formula):
                    print(" < ", end="", file=buffer)
                else:
                    print(" ~ ", end="", file=buffer)
            print(node, end="", file=buffer)
            last = node
        print(file=buffer)

    def _node_to_str(
        self, node: Node, *, index: str = "i", show_bound: bool = False
    ) -> str:
        bound = self[node]
        if bound is None:
            return str(node)
        elif show_bound:
            return f"<{node}, {index}: {bound(z3.Int(index))}>"
        else:
            return f"<{node}, {index}>"

    def eval_relation(self, relation: Relation, *args: SymElement) -> z3.BoolRef:
        nodes = tuple(arg.node for arg in args)
        indices = [arg.to_index() for arg in args]
        return self.relations[relation][nodes](*indices)


def _indices_for_decl(decl: z3.FuncDeclRef) -> list[z3.Int]:
    return [z3.Int(f"i_{i + 1}") for i in range(decl.arity())]
