import json
import re

import xmltodict
from collections import defaultdict
from relationalai.early_access.dsl.adapters.orm.model import ORMEntityType, ORMValueType, ORMRole, ORMSubtypeFact, \
    ORMUniquenessConstraint, ORMExclusionConstraint, ORMMandatoryConstraint, ORMReadingOrderRole, ORMReadingOrder, \
    ExclusiveInclusiveSubtypeFact, SubtypeFact, ExclusiveSubtypeFact, ORMInclusionConstraint, RoleConstraint, \
    ORMRoleValueConstraint, ORMRingConstraint, ORMRingType
from relationalai.util.graph import topological_sort


class ORMParser:

    def __init__(self, orm_file_path):
        with open(orm_file_path) as orm_file:
            data = xmltodict.parse(orm_file.read())

        self._ontology = json.loads(json.dumps(data))

        self._model_name = self._parse_model_name()
        self._entity_types = self._parse_entity_types()
        self._value_types = self._parse_value_types()
        self._object_types = {**self._value_types, **self._entity_types}
        self._roles = self._parse_roles()
        self._fact_type_to_roles = self._parse_fact_type_to_roles()
        self._role_value_constraints = self._parse_role_value_constraints()
        self._internal_uniqueness_constraints, self._external_uniqueness_constraints = self._parse_uniqueness_constraints()
        (self._unique_roles, self._fact_type_to_internal_ucs, self._fact_type_to_complex_ucs,
         self._identifier_fact_type_to_entity_type) = self._process_internal_uniqueness_constraints()
        self._exclusion_constraints = self._parse_exclusion_constraints()
        self._inclusion_constraints = self._parse_inclusion_constraints()
        self._role_constraints = self._parse_role_constraints()
        self._mandatory_constraints = self._parse_mandatory_constraints()
        self._ring_constraints = self._parse_ring_constraints()
        self._subtype_facts = self._parse_subtype_facts()
        self._sorted_subtype_facts = self._sort_subtype_facts()
        self._fact_type_readings = self._parse_fact_types_reading_orders()
        self._mandatory_roles = self._parse_mandatory_roles()

    def model_name(self):
        return self._model_name

    def entity_types(self) -> dict[str, ORMEntityType]:
        return self._entity_types

    def object_types(self):
        return self._object_types

    def value_types(self):
        return self._value_types

    def roles(self):
        return self._roles

    def role_value_constraints(self):
        return self._role_value_constraints

    def subtype_facts(self):
        return self._subtype_facts

    def external_uniqueness_constraints(self):
        return self._external_uniqueness_constraints

    def inclusion_constraints(self):
        return self._inclusion_constraints

    def exclusion_constraints(self):
        return self._exclusion_constraints

    def role_constraints(self):
        return self._role_constraints

    def ring_constraints(self):
        return self._ring_constraints

    def unique_roles(self):
        return self._unique_roles

    def mandatory_roles(self):
        return self._mandatory_roles

    def fact_type_readings(self):
        return self._fact_type_readings

    def fact_type_to_internal_ucs(self):
        return self._fact_type_to_internal_ucs

    def fact_type_to_complex_ucs(self):
        return self._fact_type_to_complex_ucs

    def fact_type_to_roles(self):
        return self._fact_type_to_roles

    def identifier_fact_type_to_entity_type(self):
        return self._identifier_fact_type_to_entity_type

    def sorted_subtype_facts(self):
        return self._sorted_subtype_facts

    def _parse_fact_type_to_roles(self):
        fact_type_data = defaultdict(list)
        for role in self._roles.values():
            relationship_name = role.relationship_name
            fact_type_data[relationship_name].append(role)
        return fact_type_data

    def _parse_model_name(self):
        model_name = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "@Name")
        return model_name if model_name else "ORMModel"

    def _parse_entity_types(self):
        entity_types = {}
        orm_entity_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                            "orm:Objects", "orm:EntityType")
        if orm_entity_types:
            for et in orm_entity_types:
                id = et["@id"]
                name = et['@Name']
                ref_mode = et.get("@_ReferenceMode") or None
                entity_types[id] = ORMEntityType(id, name, ref_mode)
        return entity_types

    def _parse_value_types(self):
        value_types = {}
        data_types = self._parse_data_types()
        orm_value_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                           "orm:Objects", "orm:ValueType")
        if orm_value_types:
            for vt in orm_value_types:
                if not vt.get("@IsImplicitBooleanValue"):
                    id = vt["@id"]
                    name = vt["@Name"]
                    data_type = data_types[self._get_nested(vt, "orm:ConceptualDataType", "@ref")]
                    value_types[id] = ORMValueType(id, name, data_type)
        return value_types

    def _parse_data_types(self):
        data_types = {}
        orm_data_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:DataTypes")
        if orm_data_types:
            for k, v in orm_data_types.items():
                if v is not None:
                    data_types[v["@id"]] = k[4:]
        return data_types

    def _parse_roles(self):
        roles = {}
        role_to_player = self._parse_role_player()
        orm_fact_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:Facts", "orm:Fact")
        if orm_fact_types:
            for ft in orm_fact_types:
                relationship_name = (ft['@_Name'])
                orm_roles = self._get_nested(ft, "orm:FactRoles", "orm:Role")
                if orm_roles:
                    roles_list = orm_roles if isinstance(orm_roles, list) else [orm_roles]
                    for ro in roles_list:
                        role_id = ro["@id"]
                        role_name = ro["@Name"]
                        role_player = self._get_nested(ro, "orm:RolePlayer", "@ref")
                        if role_player in self._object_types:
                            roles[role_id] = ORMRole(role_id, role_name, relationship_name, role_to_player[role_id])
        return roles

    def _parse_role_player(self):
        role_to_player = {}
        object_types = ((self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:Objects",
                                          "orm:EntityType") or []) +
                        (self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:Objects",
                                          "orm:ValueType") or []))
        for ot in object_types:
            roles = self._get_nested(ot, "orm:PlayedRoles", "orm:Role")
            if roles:
                roles_list = roles if isinstance(roles, list) else [roles]
                for role in roles_list:
                    role_to_player[role["@ref"]] = ot["@id"]
        return role_to_player

    def _parse_role_value_constraints(self):
        role_value_constraints = {}
        fact_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:Facts", "orm:Fact")
        if fact_types:
            for ft in fact_types:
                roles = self._get_nested(ft, "orm:FactRoles", "orm:Role")
                if roles and isinstance(roles, list):
                    for ro in roles:
                        value_ranges = self._get_nested(ro, "orm:ValueRestriction", "orm:RoleValueConstraint",
                                                        "orm:ValueRanges", "orm:ValueRange")
                        if value_ranges:
                            values = []
                            for rvc in value_ranges:
                                if rvc.get("@MinValue") != rvc.get("@MaxValue"):
                                    raise Exception("Unsupported value range constraint.")
                                values.append(rvc["@MinValue"])
                            role_value_constraints[ro["@id"]] = ORMRoleValueConstraint(self._roles[ro["@id"]], values)
        return role_value_constraints

    def _parse_role_constraints(self):
        constraints = []
        for ec in self._inclusion_constraints.values():
            if set(ec.roles).issubset(self._roles):
                constraints.append(RoleConstraint(ec.roles, ec.exclusive, True))
        for ec in self._exclusion_constraints.values():
            if set(ec.roles).issubset(self._roles):
                constraints.append(RoleConstraint(ec.roles, True, ec.inclusive))
        return constraints

    def _parse_subtype_facts(self):
        subtype_of = defaultdict(list)
        for subtype_arrow in self._parse_subtype_arrows():
            subtype_object = self._build_subtype_object(subtype_arrow)
            subtype_of[subtype_object.supertype_name].append(subtype_object)
        return subtype_of

    def _parse_subtype_arrows(self):
        subtype_of = []
        orm_subtype_facts = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                             "orm:Facts", "orm:SubtypeFact")
        if orm_subtype_facts:
            subtype_facts = orm_subtype_facts if isinstance(orm_subtype_facts, list) else [orm_subtype_facts]
            for sft in subtype_facts:
                subtype_of.append(self._parse_subtype_arrow(sft))
        return subtype_of

    def _parse_subtype_arrow(self, subtype_fact):
        fact_roles = subtype_fact["orm:FactRoles"]
        subtype_id = self._get_nested(fact_roles, "orm:SubtypeMetaRole", "@id")
        subtype = self._get_nested(fact_roles, "orm:SubtypeMetaRole", "orm:RolePlayer", "@ref")
        supertype_id = self._get_nested(fact_roles, "orm:SupertypeMetaRole", "@id")
        supertype = self._get_nested(fact_roles, "orm:SupertypeMetaRole", "orm:RolePlayer", "@ref")
        return ORMSubtypeFact(subtype_id, subtype, supertype_id, supertype)

    def _parse_uniqueness_constraints(self):
        internal = {}
        external = {}
        orm_ucs = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                   "orm:Constraints", "orm:UniquenessConstraint")
        if orm_ucs:
            for uc in orm_ucs:
                uc_id = uc["@id"]
                pid = uc.get("orm:PreferredIdentifierFor", None)
                identifies = None
                if pid is not None and pid["@ref"] in self._entity_types:
                    identifies = pid["@ref"]
                uc_roles = []
                roles = self._get_nested(uc, "orm:RoleSequence", "orm:Role")
                if roles:
                    roles_list = roles if isinstance(roles, list) else [roles]
                    for ro in roles_list:
                        role_id = ro["@ref"]
                        uc_roles.append(role_id)
                constraint = ORMUniquenessConstraint(uc_id, uc_roles, identifies)
                target = internal if uc.get("@IsInternal") == "true" else external
                target[uc_id] = constraint
        return internal, external

    def _process_internal_uniqueness_constraints(self):
        unique_roles = []
        fact_type_to_internal_uc = defaultdict(list)
        fact_type_to_complex_uc = defaultdict(list)
        identifier_fact_type_to_entity_type = defaultdict()

        for uc in self._internal_uniqueness_constraints.values():
            if uc.identifies is None and len(uc.roles) == 1:
                unique_roles.extend(uc.roles)

            role_id = uc.roles[0]
            role = self._roles.get(role_id)
            if not role or not role.relationship_name:
                continue

            fact_type = role.relationship_name

            fact_type_to_internal_uc[fact_type].append(uc)

            if uc.identifies is None:
                if len(uc.roles) > 1:
                    fact_type_to_complex_uc[fact_type].append(uc)
            else:
                identifier_fact_type_to_entity_type[fact_type] = self._entity_types[uc.identifies]

        return unique_roles, fact_type_to_internal_uc, fact_type_to_complex_uc, identifier_fact_type_to_entity_type

    def _parse_mandatory_constraints(self):
        mandatory_constraints = {}
        orm_mcs = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                   "orm:Constraints", "orm:MandatoryConstraint")
        if orm_mcs:
            for mc in orm_mcs:
                mc_id = mc["@id"]
                if mc.get("@IsSimple", None) is not None:
                    mc_roles = []
                    roles = self._get_nested(mc, "orm:RoleSequence", "orm:Role")
                    if roles:
                        roles_list = roles if isinstance(roles, list) else [roles]
                        for ro in roles_list:
                            role_id = ro["@ref"]
                            mc_roles.append(role_id)
                    constraint = ORMMandatoryConstraint(mc_id, mc_roles)
                    mandatory_constraints[mc_id] = constraint
        return mandatory_constraints

    def _parse_inclusion_constraints(self):
        inclusion_constraints = {}
        orm_mcs = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                   "orm:Constraints", "orm:MandatoryConstraint")
        if orm_mcs:
            for mc in orm_mcs:
                mc_id = mc["@id"]
                mc_name = mc["@Name"]
                if mc_name.startswith("InclusiveOrConstraint"):
                    mc_roles = []
                    roles = self._get_nested(mc, "orm:RoleSequence", "orm:Role")
                    if roles:
                        roles_list = roles if isinstance(roles, list) else [roles]
                        for ro in roles_list:
                            role_id = ro["@ref"]
                            mc_roles.append(role_id)
                    exclusive = mc.get("orm:ExclusiveOrExclusionConstraint", None) is not None
                    constraint = ORMInclusionConstraint(mc_id, mc_roles, exclusive)
                    inclusion_constraints[mc_id] = constraint
        return inclusion_constraints

    def _parse_exclusion_constraints(self):
        exclusion_constraints = {}
        orm_mcs = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                   "orm:Constraints", "orm:ExclusionConstraint")
        if orm_mcs:
            for ec in orm_mcs:
                ec_id = ec["@id"]
                ec_roles = []
                role_sequences = self._get_nested(ec, "orm:RoleSequences", "orm:RoleSequence")
                role_sequences_list = role_sequences if isinstance(role_sequences, list) else [role_sequences]
                for rs in role_sequences_list:
                    roles = self._get_nested(rs, "orm:Role")
                    if roles:
                        roles_list = roles if isinstance(roles, list) else [roles]
                        for ro in roles_list:
                            role_id = ro["@ref"]
                            ec_roles.append(role_id)
                inclusive = ec.get("orm:ExclusiveOrMandatoryConstraint", None) is not None
                constraint = ORMExclusionConstraint(ec_id, ec_roles, inclusive)
                exclusion_constraints[ec_id] = constraint
        return exclusion_constraints

    def _parse_ring_constraints(self):
        ring_constraints = {}
        orm_rcs = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel",
                                   "orm:Constraints", "orm:RingConstraint")
        if orm_rcs:
            rcs_list = orm_rcs if isinstance(orm_rcs, list) else [orm_rcs]
            for rc in rcs_list:
                rc_id = rc["@id"]
                rc_type = rc["@Type"]
                rc_roles = []
                roles = self._get_nested(rc, "orm:RoleSequence", "orm:Role")
                if roles:
                    roles_list = roles if isinstance(roles, list) else [roles]
                    for ro in roles_list:
                        role_id = ro["@ref"]
                        rc_roles.append(role_id)
                constraint = ORMRingConstraint(rc_id, rc_roles, self.process_ring_constraint(rc_type))
                ring_constraints[rc_id] = constraint
        return ring_constraints

    def _parse_fact_types_reading_orders(self):
        fact_types_readings = defaultdict(list)
        fact_types = self._get_nested(self._ontology, "ormRoot:ORM2", "orm:ORMModel", "orm:Facts", "orm:Fact")
        if fact_types:
            for ft in fact_types:
                relationship_name = (ft['@_Name'])
                reading_orders = self._get_nested(ft, "orm:ReadingOrders", "orm:ReadingOrder")
                if reading_orders:
                    ros = reading_orders if isinstance(reading_orders, list) else [reading_orders]
                    for ro in ros:
                        reading_order = self._parse_fact_types_reading_order(ro)
                        if reading_order:
                            fact_types_readings[relationship_name].append(reading_order)
        return fact_types_readings

    def _parse_fact_types_reading_order(self, reading_order):
        role_texts = self._get_nested(reading_order, "orm:Readings", "orm:Reading", "orm:ExpandedData", "orm:RoleText")
        if not role_texts:
            return None

        role_text_list = role_texts if isinstance(role_texts, list) else [role_texts]

        # Step 1: Build roles with whatever RoleText exists
        index_to_role = {}
        for rt in role_text_list:
            index = int(rt["@RoleIndex"])
            index_to_role[index] = ORMReadingOrderRole(
                index=index,
                prefix=rt.get("@PreBoundText", "").strip() or None,
                postfix=rt.get("@PostBoundText", "").strip() or None,
                text=rt.get("@FollowingText", "").strip() or None,
                role=None
            )

        # Step 2: Ensure all roles from RoleSequence are represented, even if no RoleText exists
        roles = self._get_nested(reading_order, "orm:RoleSequence", "orm:Role")
        if roles:
            roles_list = roles if isinstance(roles, list) else [roles]
            for i, role in enumerate(roles_list):
                role_id = role.get("@ref")
                orm_role = self._roles.get(role_id, None)
                if i in index_to_role:
                    index_to_role[i].role = orm_role
                else:
                    # Fill in with None/defaults if RoleText is missing
                    index_to_role[i] = ORMReadingOrderRole(
                        index=i,
                        prefix=None,
                        postfix=None,
                        text=None,
                        role=orm_role
                    )

        return ORMReadingOrder([index_to_role[i] for i in sorted(index_to_role)])

    @staticmethod
    def _get_nested(d, *keys):
        for key in keys:
            d = d.get(key)
            if d is None:
                return None
        return d

    def _build_subtype_object(self, subtype_arrow):
        sub_name = self._object_types[subtype_arrow.subtype].name
        sup_name = self._object_types[subtype_arrow.supertype].name
        for ec, ec_metadata in self._exclusion_constraints.items():
            if ec_metadata.inclusive and subtype_arrow.supertype_role_id in ec_metadata.roles:
                return ExclusiveInclusiveSubtypeFact(sub_name, sup_name)
            elif subtype_arrow.supertype_role_id in ec_metadata.roles:
                return ExclusiveSubtypeFact(sub_name, sup_name)
        return SubtypeFact(sub_name, sup_name)

    def _parse_mandatory_roles(self):
        mandatory_roles = []
        for mc in self._mandatory_constraints.values():
            mandatory_roles.extend(mc.roles)
        return mandatory_roles

    def _sort_subtype_facts(self):
        # Build a dependency graph to sort subtype facts so that child entity are declared
        # after the parent entity that they extend
        subtype_facts = self._subtype_facts
        nodes = []
        edges = []
        sorted_subtype_facts = []
        # Parent depends on child, i.e., parent entity must be declared before child entity
        for parent in subtype_facts:
            nodes.append(parent)
            for child in subtype_facts[parent]:
                edges.append((parent, child.subtype_name))
        for parent in topological_sort(nodes, edges):
            sorted_subtype_facts.extend(subtype_facts[parent])
        return sorted_subtype_facts

    @staticmethod
    def process_ring_constraint(rc_type):
        pattern = '|'.join(r_type.value for r_type in ORMRingType)
        matches = re.finditer(pattern, rc_type)

        ring_types = []
        for match in matches:
            ring_types.append(ORMRingType(match.group(0)))
        return ring_types
