from __future__ import annotations
from typing import Literal, Optional, Iterable, Mapping, Type, Any, Self
import pydantic as pc
import networkx as nx
import itertools
import z3
import enum
from icecream import ic
import pandas as pd
from .utils import free_vars, SimpleTypesZ3, OptTypesZ3, ArbTypesModel

DECLARED_SORTS = set()


class SimpleZ3Type(ArbTypesModel):
    py_type: Type
    z3_sort: z3.SortRef

    @classmethod
    def from_py(cls, input: Type) -> Self:
        assert input in SimpleTypesZ3, f"'{input}' is not supported as simple type"
        return cls(
            py_type=input,
            z3_sort=SimpleTypesZ3[input],
        )

    def get_constraint_value(self, input: Any):
        if self.py_type is int:
            return int(input)
        elif self.py_type is float:
            return float(input)
        elif self.py_type is bool:
            return bool(input)
        elif self.py_type is str:
            return z3.StringVal(input)
        else:
            raise Exception("Unknown type")

    def from_z3(self, input: Any, m: z3.Model):
        if self.py_type is int:
            return int(str(input))
        elif self.py_type is float:
            return float(str(input))
        elif self.py_type is bool:
            return bool(z3.is_true(input))
        elif self.py_type is str:
            return input.as_string()
        else:
            raise Exception("Unknown type")

    def pretty_str(self):
        return str(self.z3_sort)


class EnumZ3Type(ArbTypesModel):
    enum_type: Type | list[str]
    z3_sort: z3.DatatypeSortRef
    z2py: dict[z3.DatatypeRef, Any]
    py2z: dict[Any, z3.DatatypeRef]
    name: str

    def get_constraint_value(self, input: Any):
        assert input in self.py2z
        return self.py2z[input]

    def from_z3(self, input: Any, m: z3.Model):
        assert input in self.z2py
        return self.z2py[input]

    def pretty_str(self):
        return str(self.z3_sort)


class OptionalZ3Type(ArbTypesModel):
    py_type: Type
    z3_sort: z3.DatatypeSortRef

    @classmethod
    def from_py(cls, input: Type) -> Self:
        assert input in OptTypesZ3, f"Optional types are not supported for '{input}'"
        return cls(
            py_type=input,
            z3_sort=OptTypesZ3[input],
        )

    def get_constraint_value(self, input: Any):
        if input is None:
            return self.z3_sort.Nil
        elif self.py_type is int:
            val = int(input)
        elif self.py_type is float:
            val = float(input)
        elif self.py_type is bool:
            val = bool(input)
        elif self.py_type is str:
            val = z3.StringVal(input)
        else:
            raise Exception("Unknown type")
        return self.z3_sort.Some(val)

    def from_z3(self, input: Any, m: z3.Model):
        if not m.eval(self.z3_sort.is_Some(input)):
            return None

        val = m.eval(self.z3_sort.val(input))
        # ic(input, val, val.sort(), val.sort().name())
        if self.py_type is int:
            res = val.as_long()
        elif self.py_type is float:
            res = float(val)
        elif self.py_type is bool:
            res = bool(val)
        elif self.py_type is str:
            res = val.as_string()
        else:
            raise Exception("Unknown type")
        return res

    def pretty_str(self):
        return str(self.z3_sort)


class Field(ArbTypesModel):
    name: str
    typ: SimpleZ3Type | EnumZ3Type | OptionalZ3Type
    type_kind: Literal['simple', 'enum', 'optional']
    z3var: z3.SeqRef | z3.BoolRef | z3.ArithRef | z3.DatatypeRef
    default: Any
    optional: bool
    doc: Optional[str]
    tracking_literal: z3.BoolRef
    kind: Literal['field'] = 'field'
    expr: Optional[z3.BoolRef] = None  # Bad design, but should work for now...

    def value_constraint(self, input: Any) -> z3.BoolRef:
        value = self.typ.get_constraint_value(input)
        return self.z3var == value

    def from_z3(self, input: Any, m: z3.Model):
        return self.typ.from_z3(input, m)


class Rule(ArbTypesModel):
    expr: z3.BoolRef
    name: str
    tags: tuple[str, ...]
    tracking_literal: z3.BoolRef
    kind: Literal['rule'] = 'rule'


class Config:

    def __init__(
        self,
        name: Optional[str] = None,
        solver_options: Optional[dict] = None,
    ):
        self._fields: dict[str, Field] = {}
        self._rules: dict[str, Rule] = {}
        self._enum_sorts: dict[str, EnumZ3Type] = {}
        self._rule_count = 0
        self._cast: dict = {}

    def is_nil(self, field_name: str) -> z3.BoolRef:
        assert type(field_name) is str, (
            "Make sure you are passing field name, not Z3 constant")
        field = self._fields[field_name]
        assert field.optional, (
            f"Attempted check for nil state for non-optional field '{field_name}'"
        )
        return field.z3var == field.typ.z3_sort.Nil

    def not_nil(self, field_name: str) -> z3.BoolRef:
        assert type(field_name) is str, (
            "Make sure you are passing field name, not Z3 constant")
        field = self._fields[field_name]
        assert field.optional, (
            f"Attempted check for nil state for non-optional field '{field_name}'"
        )
        return field.z3var != field.typ.z3_sort.Nil

    def get_some(self, field_name: str, strict: bool = False):
        assert type(field_name) is str, (
            "Make sure you are passing field name, not Z3 constant")
        field = self._fields[field_name]
        if strict and not field.optional:
            assert field.optional, (
                f"Attempted check for nil state for non-optional field '{field_name}'"
            )
        elif not field.optional:
            return field.z3var
        else:
            return field.typ.z3_sort.val(field.z3var)

    def equivalent(self, *clauses):
        if len(clauses) < 2:
            return []
        return (a == b for i, a in enumerate(clauses)
                for j, b in enumerate(clauses) if i < j)

    def distinct(self, *clauses):
        if len(clauses) < 2:
            return []
        return (a != b for i, a in enumerate(clauses)
                for j, b in enumerate(clauses) if i < j)

    def exclusive(self, *clauses):
        if len(clauses) < 2:
            return tuple()
        return (z3.Implies(true_clause, z3.Not(false_clause))
                for i, true_clause in enumerate(clauses)
                for j, false_clause in enumerate(clauses) if i != j)

    def _mk_field_type(
        self,
        typ: Type,
        optional: bool,
    ) -> SimpleZ3Type | OptionalZ3Type:
        if optional:
            return OptionalZ3Type.from_py(typ)
        else:
            return SimpleZ3Type.from_py(typ)

    def _mk_z3_var(self, name: str, sort: z3.SortRef):
        return z3.Const(name, sort)

    def add_field(
        self,
        name: str,
        typ: Type | list[str],
        *,
        default: Any = None,
        optional: bool = False,
        doc: Optional[str] = None,
    ) -> Field:
        if name in self._fields:
            raise ValueError(f"Field '{name}' already exists")

        if isinstance(typ, list) or issubclass(typ, enum.Enum):
            return self._add_enum(
                name=name,
                enum_cls=typ,
                default=default,
                optional=optional,
                doc=doc,
            )
        else:
            return self._add_basic_field(
                name=name,
                typ=typ,
                default=default,
                optional=optional,
                doc=doc,
            )

    def _add_basic_field(
        self,
        name: str,
        typ: Type,
        *,
        default: Any = None,
        optional: bool = False,
        doc: Optional[str] = None,
    ) -> Field:
        type_obj = self._mk_field_type(typ, optional)
        # ic(name, type_obj)
        z3var = self._mk_z3_var(name, type_obj.z3_sort)
        tracking_literal = z3.Bool(f"__var_track_{name}")
        field = Field(name=name,
                      typ=type_obj,
                      type_kind='simple',
                      z3var=z3var,
                      default=default,
                      optional=optional,
                      doc=doc,
                      tracking_literal=tracking_literal)
        self._fields[name] = field
        return field

    def _mk_enum_sort(
        self,
        enum_cls: Type | list[str],
        sort_name: Optional[str] = None,
    ) -> EnumZ3Type:
        if isinstance(enum_cls, type) and issubclass(enum_cls, enum.Enum):
            mode = 'enum'
        elif isinstance(enum_cls, list) and all(
                type(x) is str for x in enum_cls):
            mode = 'list'
        else:
            raise TypeError(
                f"Enum types created from either a subclass of enum.Enum or a list of string values. Got {enum_cls}"
            )

        enum_names: list[str]
        if mode == 'enum':
            enum_names = [m.name for m in enum_cls]
        else:
            enum_names = list(enum_cls)
        assert len(enum_names) == len(set(enum_names)), (
            f"Duplicates are found in the list of options: {enum_names}")

        signature = self._get_enum_signature(enum_cls, enum_names, mode)
        if signature in self._enum_sorts:
            return self._enum_sorts[signature]

        if sort_name is None and mode == 'enum':
            sort_name = enum_cls.__name__
        elif sort_name is None and mode == 'list':
            i = 0
            while True:
                i += 1
                sort_name = f"Enum{i}"
                if all(sort.name != sort_name for sort in self._enum_sorts.
                       values()) and sort_name not in DECLARED_SORTS:
                    break

        # ic(sort_name)
        DECLARED_SORTS.add(sort_name)
        sort, z3_vals = z3.EnumSort(f"{sort_name}_EnumSort", enum_names)

        # Complicated updates for convenient casting
        py2z = {m: z3_vals[i] for i, m in enumerate(enum_cls)}
        z2py = {v: k for k, v in py2z.items()}
        self._cast.update(py2z)
        self._cast.update(z2py)
        if mode == 'enum':
            py2z.update({k.name: v for k, v in py2z.items()})
            self._cast.update({
                (enum_cls.__name__, k.name): v
                for k, v in py2z.items()
            })

        enum_type = EnumZ3Type(
            enum_type=enum_cls,
            z3_sort=sort,
            z2py=z2py,
            py2z=py2z,
            name=sort_name,
        )
        self._enum_sorts[signature] = enum_type
        return enum_type

    def _get_enum_signature(
        self,
        enum_cls: Type | list[str],
        enum_names: list[str],
        mode: Literal['enum', 'list'],
    ) -> str:
        if mode == 'enum':
            sig = enum_cls.__name__
            if sig in self._enum_sorts:
                orig_sort = self._enum_sorts[sig]
                assert sorted([x.name
                               for x in orig_sort.enum_type]) == enum_names
            return sig
        else:
            return tuple(sorted(enum_names))

    def _add_enum(
        self,
        name: str,
        enum_cls: Type | list[str],
        *,
        default: Any = None,
        sort_name: Optional[str] = None,
        optional: bool = False,
        doc: Optional[str] = None,
    ) -> Field:
        assert not optional, "Optional mode is not supported for enums. Include null case into your enum."
        assert default is None or default in enum_cls, (
            f"default value '{default}' must be a member of {enum_cls.__name__}"
        )

        ctx = z3.main_ctx()
        enum_sort = self._mk_enum_sort(enum_cls, sort_name)
        z3var = z3.Const(name, enum_sort.z3_sort)
        tracking_literal = z3.Bool(f"__var_track_{name}")

        field = Field(
            name=name,
            typ=enum_sort,
            type_kind='enum',
            z3var=z3var,
            default=enum_sort.py2z[default] if default is not None else None,
            optional=optional,
            doc=doc,
            tracking_literal=tracking_literal,
        )
        self._fields[name] = field
        return field

    def __getitem__(self, key: str) -> z3.ExprRef:
        return self._fields[key].z3var

    def cast(self, x):
        if x in self._cast:
            return self._cast[x]
        else:
            return x

    def add_rules(
            self,
            exprs: Iterable[z3.BoolRef],
            name: Optional[str] = None,
            tags: Iterable[str] = (),
    ) -> list[Rule]:
        return [
            self.add_rule(expr,
                          name=f"{name}-{i}" if name is not None else None)
            for i, expr in enumerate(exprs)
        ]

    def add_rule(
            self,
            expr: z3.BoolRef,
            *,
            name: Optional[str] = None,
            tags: Iterable[str] = (),
    ) -> Rule:
        assert z3.is_bool(expr), (
            f"expr must be a boolean z3 expression; got sort {getattr(expr, 'sort', lambda: None)()}"
        )

        unique_name = name or f"rule_{self._rule_count}"
        assert unique_name not in self._rules, f"Rule '{unique_name}' already exists"
        self._rule_count += 1

        tracking_literal = z3.Bool(f"__rule_track_{unique_name}")
        tags_tuple = tuple(map(str, tags)) if tags is not None else tuple()
        rule = Rule(
            expr=expr,
            name=unique_name,
            tags=tags_tuple,
            tracking_literal=tracking_literal,
        )
        self._rules[unique_name] = rule
        return rule

    def rules_on_fields(self, field_names: list[str]) -> list[Rule]:
        field_names_set = {self._fields[name].z3var for name in field_names}
        res = [
            rule for rule in self._rules.values()
            if len(free_vars(rule.expr).intersection(field_names_set)) > 0
        ]
        return res

    def new(
        self,
        mapping: Mapping[str, Any],
        *,
        partial: bool = True,
    ) -> "ConfigInstance":
        # TODO: Find redundant input params
        res = self._basic_new(
            mapping=mapping,
            partial=partial,
            get_core=True,
        )

        if res.status.sat() and partial:
            uncertain_fields = {}
            for name, field in self._fields.items():
                if name in mapping or field.default is not None:
                    continue

                elim_res = self._basic_new(
                    mapping=mapping,
                    partial=partial,
                    extra_constraints=[
                        (z3.Not(field.value_constraint(res.solved[name])),
                         field.tracking_literal)
                    ],
                    get_core=False)
                if elim_res.status.sat():
                    uncertain_fields[name] = elim_res
            res.uncertain_fields = uncertain_fields
        elif res.status.unsat():
            assert res.conflicts is not None
            do_next_step = True
            while do_next_step:
                tmp_res = self._basic_new(
                    mapping=mapping,
                    partial=partial,
                    drop_rules=[
                        rule_name for core in res.conflicts
                        for rule_name in core.rules.keys()
                    ],
                    get_core=True,
                )

                do_next_step = tmp_res.status.unsat()
                if do_next_step:
                    assert tmp_res.conflicts is not None
                    res.conflicts += tmp_res.conflicts
        return res

    def _basic_new(
        self,
        mapping: Mapping[str, Any],
        *,
        extra_constraints: list[tuple[z3.ExprRef, z3.BoolRef]] = [],
        drop_rules: list[str] = [],
        partial: bool = True,
        get_core: bool = True,
    ) -> "ConfigInstance":
        solver = z3.Solver()
        provided = dict(mapping)

        def add(expr, track):
            if get_core:
                solver.assert_and_track(expr, track)
            else:
                solver.assert_exprs(expr)

        for name, field in self._fields.items():
            if name in provided:
                val = provided[name]
            elif field.default is not None:
                val = field.default
            elif not partial and not field.optional:
                raise Exception(f"Missing required field '{name}'")
            else:
                field.expr = None
                continue
            expr = field.value_constraint(val)
            add(expr, field.tracking_literal)
            field.expr = expr

        for rule in self._rules.values():
            if rule.name in drop_rules:
                continue
            add(rule.expr, rule.tracking_literal)

        for constraint, lit in extra_constraints:
            add(constraint, lit)

        if get_core:
            assert len(extra_constraints) == 0, (
                "Cannot get unsat core with extra constraints yet")
            tracking = {
                x.tracking_literal: x
                for x in itertools.chain(self._rules.values(),
                                         self._fields.values())
                # if x.name not in drop_rules
            }

        result = solver.check()
        if result == z3.sat:
            model = solver.model()
            solved = {}
            for name, field in self._fields.items():
                new_val = model.eval(field.z3var, model_completion=True)
                # ic(name, field, new_val)
                solved[name] = field.from_z3(new_val, model)

            return ConfigInstance(
                config=self,
                provided=provided,
                solved=solved,
                status=SMTResult.SAT,
                model=model,
            )
        elif result == z3.unsat:
            # TODO: Conflict resolution strats:
            # constraint removal, max-smt, interactive
            if get_core:
                unsat_core = [tracking[c] for c in solver.unsat_core()]
                disjoint_cores = self._get_disjoint_cores(unsat_core)
                # NOTE: It seems that unsat_core always returns only one conn comp
                assert len(disjoint_cores) == 1, (
                    "Expected one conn comp from unsat core")
                conflicts = [
                    ConflictingConstraints(
                        values={c.name: c
                                for c in core if c.kind == 'field'},
                        rules={c.name: c
                               for c in core if c.kind == 'rule'},
                    ) for core in disjoint_cores
                ]
            else:
                unsat_core = {}
                conflicts = None

            return ConfigInstance(config=self,
                                  provided=provided,
                                  status=SMTResult.UNSAT,
                                  conflicts=conflicts)
        else:
            return ConfigInstance(config=self,
                                  provided=provided,
                                  status=SMTResult.UNKNOWN)

    def _get_disjoint_cores(
            self, unsat_core: list[Rule | Field]) -> list[list[Rule | Field]]:
        gr = nx.Graph()

        class NodeKind(enum.Enum):
            EXPR = enum.auto()
            SETVAR = enum.auto()
            FREEVAR = enum.auto()

        vartypes = {}
        for x in unsat_core:
            if x.kind == 'rule':
                continue
            vars = free_vars(x.expr)
            assert len(vars) == 1
            cur_var = next(iter(vars))
            vartypes[x.name] = cur_var

        for x in unsat_core:
            if x.kind == 'field':
                continue
            ic(x.expr, free_vars(x.expr))
            vars = free_vars(x.expr)
            gr.add_node(x.name)
            gr.nodes[x.name]['kind'] = NodeKind.EXPR
            gr.add_nodes_from([(str(v), {
                'kind':
                NodeKind.SETVAR if str(v) in vartypes else NodeKind.FREEVAR
            }) for v in vars])
            gr.add_edges_from([(x.name, str(i)) for i in vars])

        # ic(list(gr.nodes(data=True)))
        # ic(list(gr.edges))

        subg = gr.subgraph([
            n for n, data in gr.nodes(data=True)
            if data['kind'] != NodeKind.SETVAR
        ])

        # ic(list(subg.nodes(data=True)))
        # ic(list(subg.edges))

        ind_cores = []
        for comp in nx.connected_components(subg):
            ind_gr_nodes = set(comp)
            for rule_node in comp:
                if gr.nodes[rule_node]['kind'] != NodeKind.EXPR:
                    continue
                for nb in gr.neighbors(rule_node):
                    if gr.nodes[nb]['kind'] == NodeKind.SETVAR:
                        ind_gr_nodes.add(nb)
            # ind_gr = gr.subgraph(ind_gr_nodes)
            # ic(list(ind_gr.nodes(data=True)))
            # ic(list(ind_gr.edges))
            ind_cores.append([x for x in unsat_core if x.name in ind_gr_nodes])

        return ind_cores


class ConflictingConstraints(ArbTypesModel):
    values: dict[str, Field]
    rules: dict[str, Rule]

    def __repr__(self):
        values_str = "\n".join([
            f"{i}) {field.expr}"
            for i, (name, field) in enumerate(self.values.items(), 1)
        ])
        rule_str = "\n".join([
            f"{i}) {rule.expr}"
            for i, (name, rule) in enumerate(self.rules.items(), 1)
        ])
        return f"""\
Field values:
{values_str}
are in conflict with rules:
{rule_str}
"""


class SMTResult(enum.Enum):
    SAT = enum.auto()
    UNSAT = enum.auto()
    UNKNOWN = enum.auto()

    def __str__(self):
        return str(self._name_)

    def __repr__(self):
        return str(self._name_)

    def sat(self):
        return self == SMTResult.SAT

    def unsat(self):
        return self == SMTResult.UNSAT


class FieldStatus(enum.Enum):
    PROVIDED = enum.auto()
    SOLVED = enum.auto()
    CONFLICT = enum.auto()
    UNCERTAIN = enum.auto()
    UNSOLVED = enum.auto()

    def __str__(self):
        return str(self._name_).lower()


class FieldInst(ArbTypesModel):
    status: FieldStatus
    conflict_ids: Optional[list[int]] = None

    @classmethod
    def new(cls, status: FieldStatus, ids: Optional[list[int]] = None) -> Self:
        return cls(
            status=status,
            conflict_ids=ids,
        )

    @pc.model_validator(mode='after')
    def validate_conflict_ids(self) -> Self:
        if self.status != FieldStatus.CONFLICT and self.conflict_ids is not None:
            raise pc.ValidationError(
                f"conflict_ids must be None when status is {self.status}, "
                f"but got {self.conflict_ids}")

        if self.status == FieldStatus.CONFLICT and (
                self.conflict_ids is None or len(self.conflict_ids) == 0):
            raise pc.ValidationError(
                "conflict_ids must be provided when status is CONFLICT")

        return self

    def __str__(self):
        if self.status == FieldStatus.CONFLICT:
            assert self.conflict_ids is not None
            return f"{self.status} ({','.join(str(x) for x in self.conflict_ids)})"
        else:
            return str(self.status)


class ConfigInstance(ArbTypesModel):
    config: Config
    provided: dict
    status: SMTResult
    solved: Optional[dict] = pc.Field(default=None)
    model: Optional[z3.ModelRef] = pc.Field(default=None)
    conflicts: Optional[list[ConflictingConstraints]] = pc.Field(default=None)
    uncertain_fields: Optional[dict[str,
                                    ConfigInstance]] = pc.Field(default=None)

    def build_table_dataframe(self) -> pd.DataFrame:
        rows = []

        handled_fields = set()

        def handle_field(
            key: str,
            value: str | int | bool | None,
            status: FieldInst,
        ) -> None:
            if key not in self.config._fields:
                return
            type_str = self.config._fields[key].typ.pretty_str()
            rows.append({
                "Field": key,
                "Value": str(value),
                "Type": type_str,
                "Status": status,
            })
            handled_fields.add(key)

        if self.conflicts is not None:
            for i, conflict in enumerate(self.conflicts, 1):
                for key, field in conflict.values.items():
                    handle_field(key, self.provided[key],
                                 FieldInst.new(FieldStatus.CONFLICT, [i]))

        for key, v in self.provided.items():
            if key in handled_fields:
                continue
            handle_field(key, v, FieldInst.new(FieldStatus.PROVIDED))

        if self.uncertain_fields is not None:
            assert self.solved is not None
            for key in self.uncertain_fields:
                handle_field(key, self.solved[key],
                             FieldInst.new(FieldStatus.UNCERTAIN))

        if self.solved is not None:
            for key, v in self.solved.items():
                if key in handled_fields:
                    continue
                handle_field(key, v, FieldInst.new(FieldStatus.SOLVED))

        for key, field in self.config._fields.items():
            if key in handled_fields:
                continue
            handle_field(key, "", FieldInst.new(FieldStatus.UNSOLVED))
        return pd.DataFrame(rows)
