import typing

from relationalai.util.graph import topological_sort
from relationalai.early_access.dsl.orm.constraints import Unique, Mandatory, RoleValueConstraint, \
    InclusiveSubtypeConstraint, ExclusiveSubtypeConstraint, Range, RingConstraint
from relationalai.early_access.dsl.orm.models import Model
from relationalai.early_access.dsl.orm.relationships import Role
from relationalai.early_access.dsl.orm.types import Concept
from relationalai.early_access.metamodel.util import Printer as BasePrinter, OrderedSet
import relationalai.early_access.builder as qb

class Printer(BasePrinter):
    _unique_constraints: OrderedSet[Unique]
    _mandatory_constraints: OrderedSet[Mandatory]
    _role_value_constraints: OrderedSet[RoleValueConstraint]
    _inclusive_subtype_constraints: OrderedSet[InclusiveSubtypeConstraint]
    _exclusive_subtype_constraints: OrderedSet[ExclusiveSubtypeConstraint]
    _ring_constraints: OrderedSet[RingConstraint]

    def __init__(self, io: typing.Optional[typing.IO[str]] = None):
        # Set the base class field (frozen)
        object.__setattr__(self, 'io', io)

        # Set all mutable fields using object.__setattr__
        object.__setattr__(self, '_unique_constraints', OrderedSet())
        object.__setattr__(self, '_mandatory_constraints', OrderedSet())
        object.__setattr__(self, '_role_value_constraints', OrderedSet())
        object.__setattr__(self, '_inclusive_subtype_constraints', OrderedSet())
        object.__setattr__(self, '_exclusive_subtype_constraints', OrderedSet())
        object.__setattr__(self, '_ring_constraints', OrderedSet())


    def to_string(self, model: Model, enums_enabled: bool=False, pyi_enabled: bool=False) -> None:
        self._process_constraints(model)

        self._print_nl("import relationalai.early_access.builder as qb")
        self._print_nl("from relationalai.early_access.dsl.orm.models import Model")
        if pyi_enabled:
            self._print_nl("from relationalai.early_access.dsl.orm.models import Concept")
        self._print_nl("from relationalai.early_access.dsl.orm.constraints import Range")
        self._print_nl("from relationalai.early_access.dsl.orm.constraints import RingType")
        self._nl()

        self._print_nl(f"model = Model(name='{model.name}', is_primary={model.is_primary}"
                       f"{self._print_if_true('dry_run', model.qb_model()._dry_run)})")
        # use class stubs for new Concepts which extend orm QB Concept
        if not pyi_enabled:
            self._print_nl("Concept = model.Concept")
        if enums_enabled:
            self._print_nl("Enum = model.Enum")
        self._print_nl("Relationship = model.Relationship")
        self._print_nl("Unique = model.unique")
        self._print_nl("Mandatory = model.mandatory")
        self._print_nl("RoleValueConstraint = model.role_value_constraint")
        self._print_nl("InclusiveSubtypeConstraint = model.inclusive_subtype_constraint")
        self._print_nl("ExclusiveSubtypeConstraint = model.exclusive_subtype_constraint")
        self._print_nl("RingConstraint = model.ring")

        self._nl()
        self._handle_concepts(model, enums_enabled, pyi_enabled)

        self._nl()
        self._handle_relationships(model)

        self._nl()
        self._handle_ref_schemes(model)

        self._nl()
        self._handle_uniqueness_constraints()

        self._nl()
        self._handle_mandatory_constraints()

        self._nl()
        self._handle_role_value_constraints()

        self._nl()
        self._handle_subtype_constraints()

        self._nl()
        self._handle_ring_constraints()

    def _handle_concepts(self, model: Model, enums_enabled: bool, pyi_enabled: bool) -> None:
        concepts_map = model.concepts_map()
        enums_map = model.enums_map()
        sorted_concepts = self._sort_dependency_graph(model.concepts())
        # print Concept classes
        if pyi_enabled:
            for name in sorted_concepts:
                c = concepts_map.get(name)
                if c is None:
                    raise ValueError(f"The concept '{name}' was not declared but used as concept domain.")
                # skip enums when they are enabled for printer
                if c._is_enum() and enums_enabled:
                    continue
                self._print_nl(f"class {name}Concept(Concept):")
                self._print_nl(f'\t"""Represents {name} concept."""')
            self._nl()
        for name in sorted_concepts:
            c = concepts_map.get(name)
            if c is None:
                raise ValueError(f"The concept '{name}' was not declared but used as concept domain.")
            if enums_enabled and c._is_enum():
                self._print_nl(f"{name} = Enum('{name}', {[e.name for e in enums_map[name]]})")
            else:
                extends = c._extends
                if not enums_enabled and c._is_enum():
                    # todo: derive enum type
                    # ORM adapter produces only string enums
                    extends_elements = ["qb.String"]
                else:
                    extends_elements = [f"{self._get_type(ext)}" for ext in extends]
                extends_str = f", extends=[{', '.join(extends_elements)}]" if extends else ""
                concept_class = name + 'Concept' if pyi_enabled else 'Concept'
                params = f"({'model, ' if pyi_enabled else ''}'{name}'{extends_str})"
                self._print_nl(f"{name} = {concept_class}{params}")

    def _handle_relationships(self, model: Model) -> None:
        for rel in model.relationships():
            # skip autogenerated 'name' Relationship for Enums
            if rel._name == 'name' and rel._parent is not None and rel._parent._is_enum():
                continue
            # print a root Relationship
            self._print_nl(f"{self._get_relationship_name(rel)} = "
                           f"Relationship('{rel._madlib}'{self._print_if_not_empty_and_not_equal('short_name', rel._passed_short_name, rel._name)})")
            # print remaining RelationshipReadings if any
            for r in rel._readings[1:]:
                self._print_nl(f"{self._get_relationship_name(r)} = "
                               f"{self._get_relationship_name(rel)}.alt('{r._madlib}'{self._print_if_not_empty_and_not_equal('short_name', r._passed_short_name, r._name)})")

    def _handle_ref_schemes(self, model: Model) -> None:
        for concept_name, concept in model.entity_types_map().items():
            if not isinstance(concept, Concept) or not concept._reference_schemes:
                continue

            for ref_scheme in concept._reference_schemes:
                rel_names = [self._get_relationship_name(rel) for rel in ref_scheme]
                self._print_nl(f"{concept_name}.identify_by({', '.join(rel_names)})")

    def _handle_uniqueness_constraints(self) -> None:
        for constraint in self._unique_constraints:
            if not constraint.is_preferred_identifier:
                elements = []
                for role in constraint.roles():
                    elements.append(self._get_role_name(role))
                self._print_nl(f"Unique({', '.join(elements)})")

    def _handle_mandatory_constraints(self):
        for constraint in self._mandatory_constraints:
            role = constraint.roles()[0]
            role_name = self._get_role_name(role)
            self._print_nl(f"Mandatory({role_name})")

    def _handle_role_value_constraints(self) -> None:
        for constraint in self._role_value_constraints:
            role = constraint.roles()[0]
            elements = []

            for value in constraint.values():
                if isinstance(value, Range):
                    if value._start is not None and value._end is not None:
                        elements.append(f"Range.between({value._start}, {value._end})")
                    elif value._end is not None:
                        elements.append(f"Range.to_value({value._end})")
                    else:
                        elements.append(f"Range.from_value({value._start})")
                elif isinstance(value, str):
                    elements.append(repr(value))
                else:
                    elements.append(str(value))

            role_name = self._get_role_name(role)
            self._print_nl(f"RoleValueConstraint({role_name}, [{', '.join(elements)}])")

    def _handle_subtype_constraints(self):
        self._emit_subtype_constraints(self._inclusive_subtype_constraints, "InclusiveSubtypeConstraint")
        self._nl()
        self._emit_subtype_constraints(self._exclusive_subtype_constraints, "ExclusiveSubtypeConstraint")

    def _handle_ring_constraints(self):
        for constraint in self._ring_constraints:
            elements = []
            for role in constraint.roles():
                elements.append(self._get_role_name(role))
            cst_types = []
            for tp in constraint.types:
                cst_types.append(f"RingType.{tp.name}")
            self._print_nl(f"RingConstraint([{', '.join(cst_types)}], {', '.join(elements)})")

    def _emit_subtype_constraints(self, constraints, constraint_type: str):
        for constraint in constraints:
            concepts = [self._get_type(c) for c in constraint.concepts()]
            self._print_nl(f"{constraint_type}({', '.join(concepts)})")

    def _process_constraints(self, model: Model) -> None:
        for c in model.constraints():
            if isinstance(c, Unique):
                self._unique_constraints.add(c)
            elif isinstance(c, Mandatory):
                self._mandatory_constraints.add(c)
            elif isinstance(c, RoleValueConstraint):
                self._role_value_constraints.add(c)
            elif isinstance(c, InclusiveSubtypeConstraint):
                self._inclusive_subtype_constraints.add(c)
            elif isinstance(c, ExclusiveSubtypeConstraint):
                self._exclusive_subtype_constraints.add(c)
            elif isinstance(c, RingConstraint):
                self._ring_constraints.add(c)


    def _get_role_name(self, role: Role) -> str:
        return f"{self._get_relationship_name(role._relationship)}['{role._field_ref._name}']"

    def _sort_dependency_graph(self, concepts: list[qb.Concept]):
        nodes = []
        edges = []
        for concept in concepts:
            name = concept._name
            nodes.append(name)
            for ext in concept._extends:
                if not self._is_builtin_type(ext):
                    edges.append((ext._name, name))
        return topological_sort(nodes, edges)

    @staticmethod
    def _get_relationship_name(rel: typing.Union[qb.Relationship, qb.RelationshipReading]) -> str:
        if rel._parent and rel._short_name:
            parent_name = rel._parent._name
            prefix = "qb." if parent_name in Concept.builtins else ""
            return f"{prefix}{parent_name}.{rel._short_name}"
        return rel._name

    def _get_type(self, t: Concept) -> str:
        name = self._get_concept_name(t)
        return f"qb.{name}" if name in Concept.builtins else name

    def _is_builtin_type(self, t: Concept) -> bool:
        return self._get_concept_name(t) in Concept.builtins

    @staticmethod
    def _get_concept_name(t: Concept) -> str:
        return 'Integer' if t._name == 'Int' else t._name

    @staticmethod
    def _print_if_not_empty(label: str, value: typing.Optional[str]) -> str:
        return f", {label}='{value}'" if value else ""

    @staticmethod
    def _print_if_not_empty_and_not_equal(label: str, value: typing.Optional[str], compare_to: typing.Optional[str]) -> str:
        return f", {label}='{value}'" if value and value != compare_to else ""

    @staticmethod
    def _print_if_true(label: str, value: bool) -> str:
        return f", {label}=True" if value else ""

EXCLUDED_RELATIONSHIPS = ["shape"]

class InterfacePrinter(BasePrinter):

    def __init__(self, io: typing.Optional[typing.IO[str]] = None):
        # Set the base class field (frozen)
        object.__setattr__(self, 'io', io)

    def to_string(self, model: Model, enums_enabled: bool=False) -> None:
        if enums_enabled:
            self._print_nl("from enum import Enum")
        self._print_nl("from typing import Union")
        self._print_nl("import relationalai.early_access.builder as qb")
        self._print_nl("from relationalai.early_access.dsl.orm.models import Concept")
        self._nl()
        if enums_enabled:
            self._handle_enums(model)
        self._handle_concepts(model, enums_enabled)

    def _handle_enums(self, model: Model) -> None:
        for e in model.enums():
            self._print_nl(f"class {e.__name__}(Enum):")
            for member in e:
                self._print_nl(f"\t{member.name}: {e.__name__}")
            self._print_nl("\n".join(self._get_declared_relationships(e._concept, EXCLUDED_RELATIONSHIPS))) #type: ignore
            self._nl()

    def _handle_concepts(self, model: Model, enums_enabled) -> None:
        for c in model.concepts():
            if enums_enabled and c._is_enum():
                continue
            exclude_list = list(EXCLUDED_RELATIONSHIPS)
            # when enum is declared but enums are disabled print it as a Concept without "name" Relationship
            if c._is_enum():
                exclude_list.append("name")
            self._print_nl(f"class {c._name}Concept(Concept):")
            rel_strs = self._get_declared_relationships(c, exclude_list)
            self._print_nl("\tpass" if len(rel_strs) == 0 else "\n".join(rel_strs))
            self._print_nl(f"{c._name}: {c._name}Concept")
            self._nl()

    def _get_declared_relationships(self, c: qb.Concept, exclude_list: list[str]) -> list[str]:
        rel_strs = []
        for rel in c._relationships.keys():
            if rel not in exclude_list:
                rel_strs.append(f"\t{rel}: Union[qb.Relationship, qb.RelationshipReading]")
        return rel_strs