import io
import itertools
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 import DUMMY_SORT
from fest.base import AnyFormula, AnyTerm
from fest.theories import TheoryCache, TheoryName, get_theory
from fest.z3_utils import find_order_relation

type Node = str
type NodeRef = int | Node

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


@dataclass(eq=True, frozen=True)
class SymElement(TheoryCache):
    node: Node
    sort: z3.SortRef
    theory_name: TheoryName
    index: z3.ExprRef | 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(), "lia", index)

    @classmethod
    def create(
        cls,
        node: Node,
        sort: z3.SortRef,
        theory: TheoryName,
        is_summary: bool,
        index: str = "i",
    ) -> Self:
        if is_summary:
            return cls(node, sort, theory, get_theory(theory).to_base_const(index))
        else:
            return cls(node, sort, theory)

    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.ExprRef:
        if self.index is None:
            return self.theory.regular_index()
        else:
            return self.index

    def with_index(self, index: z3.ExprRef) -> 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

    def __str__(self) -> str:
        buffer = io.StringIO()
        for key, element in self.items():
            print(f"{key}: {element}", file=buffer)

        return buffer.getvalue()


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


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


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


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

    def wrapper(*indices: z3.ExprRef) -> z3.ExprRef:
        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.ExprRef) -> 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(TheoryCache):
    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)
    theory_name: TheoryName

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

    def assignments(self, variables: Iterable[z3.Const]) -> Iterable[Assignment]:
        elements = [list(self.elements_for(variable)) for variable in variables]
        for assignment in itertools.product(*elements):
            yield Assignment(assignment)

    def elements_for(self, variable: z3.Const) -> Iterable[tuple[str, SymElement]]:
        sort = variable.sort()
        for node in self.domain[sort]:
            yield str(variable), SymElement.create(
                node, sort, self.theory_name, node in self.bounds, str(variable)
            )

    @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):
            if tuple(nodes) not in self.relations[decl]:
                raise KeyError(f"{decl}{nodes}")
            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: ...

    @overload
    def __setitem__(self, key: z3.Const, value: SymElement) -> None: ...

    @overload
    def __setitem__(
        self, key: z3.Const, value: tuple[int, z3.ExprRef | None]
    ) -> 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 z3.is_const(key):
            sort = key.sort()
            if not isinstance(value, SymElement):
                node, term = value
                value = SymElement(auto_node(sort, node), sort, self.theory_name, term)
            self.constants[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()) or sort.name() == DUMMY_SORT:
                continue

            print(
                f"{sort}: {", ".join(
                    map(
                        lambda n: self._node_to_str(n, show_bound=True),
                        sorted(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 = list(
                map(
                    lambda i: self.theory.to_base_const(i),
                    _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
            and self.theory.name() == "lia"
        ):
            _show_int_order_relation(
                order_relation,
                buffer,
                self.relations[order_relation],
                self.domain[order_relation.domain(0)],
            )

        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 = list(
                    map(
                        lambda i: self.theory.to_base_const(i),
                        _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 _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(self.theory.to_base_const(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[str]:
    return [f"i{i + 1}" for i in range(decl.arity())]


def _show_int_order_relation(
    order_relation: Relation,
    buffer: io.StringIO,
    interpretation: dict[tuple[Node, ...], SymFormula],
    domain: set[Node],
) -> None:
    order_relation.domain(0)
    i1, i2 = z3.Ints("i1 i2")

    def cmp_nodes(n1: Node, n2: Node) -> int:
        cmp_formula = interpretation[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(domain, 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 = interpretation[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)
