import typing
from abc import abstractmethod, ABC
from collections import OrderedDict
from typing import Union, Callable

import v0.relationalai.semantics as qb
from v0.relationalai.semantics import define, where
from v0.relationalai.semantics.internal.internal import ConceptConstruct, Ref, RelationshipRef, ConceptFilter
from v0.relationalai.early_access.dsl.bindings.common import Binding, BindableAttribute, IdentifierConceptBinding, \
    SubtypeConceptBinding, BindableTable, ReferentConceptBinding
from v0.relationalai.early_access.dsl.codegen.helpers import reference_entity, construct_entity
from v0.relationalai.early_access.dsl.orm.constraints import Unique
from v0.relationalai.early_access.dsl.orm.relationships import Role
from v0.relationalai.semantics.metamodel.util import OrderedSet


def _get_map_madlib_and_name(name_prefix: str, value_concept: qb.Concept) -> tuple[str, str]:
    """
    Generates a name for the map based on the value concept.
    """
    name = f'{name_prefix}_row_to_{value_concept}'
    madlib = f'{name} {{row_id:RowId}} {{val:{value_concept}}}'
    return madlib, name


class GeneratedRelation(qb.Relationship):
    """
    A class representing a relation generated from a model.
    """
    def __init__(self, madlib, model, name):
        super().__init__(madlib, model=model.qb_model(), short_name=name)

    def __repr__(self):
        return f'GeneratedRelation({self._name})'


class InternallyGeneratedRelation(GeneratedRelation):
    """
    A class representing a relation generated by analyzing the model.
    """
    def __init__(self, madlib, dsl_model, name):
        super().__init__(madlib, dsl_model, name)


class AbstractMap(GeneratedRelation):
    """
    A class representing an abstract map.
    """
    def __init__(self, madlib, dsl_model, name):
        super().__init__(madlib, dsl_model, name)

    def value_player(self) -> qb.Concept:
        """
        Returns the value player of the map.
        This is the concept that the map maps to.
        """
        raise NotImplementedError()

    def formula(self):
        source_ref = self.table()
        rez_val_ref = self.value_player().ref()
        return rez_val_ref, [self(source_ref, rez_val_ref)]


class RoleMap(AbstractMap):
    """
    A class representing a base role map relation.
    """

    def __init__(self, madlib, model, name, role):
        super().__init__(madlib, model, name)
        self._role = role

    def role(self):
        """
        Returns the role of the role map.
        """
        return self._role

    def value_player(self):
        """
        Returns the value player of the role map.
        """
        return self._role.player()

    @abstractmethod
    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def table(self) -> BindableTable:
        """
        Returns the table associated with the role map.
        """
        return self.column().table

    def __repr__(self):
        return f'RoleMap({self._name})'


class InlineValueMap:
    """
    A non-QB relationship class that represents an inline value map.
    """
    def __init__(self, model, binding: Binding, role):
        self._binding = binding
        self._role = role
        self._value_player = role.player()
        self._initialize()

    def _initialize(self):
        role = self._role
        if not role.player()._is_primitive():
            raise TypeError(f'Cannot construct a value map for {role}, concept is not a value type')

        table_name = self._binding.column.table.physical_name()
        concept = role.player()
        _, self._name = _get_map_madlib_and_name(table_name, concept)
        self._has_transforms = bool(self._binding.transform_with)

    def _construct_body_atoms(self, filter_atoms):
        result_type = self.value_player()
        if self._has_transforms:
            transform_result = self._generate_transform_body()
            return result_type(transform_result), [
                where(
                    self._binding.column,  # needed when transform-to-a-literal is used to preserve row_id
                    filter_atoms
                )]
        else:
            return result_type(self._binding.column), [
                filter_atoms,
            ]

    def _generate_transform_body(self):
        transformations = self._binding.transform_with
        if transformations:
            transformations = (transformations,) if not isinstance(transformations, tuple) else transformations
            return self._chain_transformations(self._binding.column, transformations)
        raise ValueError("Incorrect state, cannot apply binding transformations with empty value converter(s)")

    def _chain_transformations(self, col_val_ref, transformations):
        result = col_val_ref
        for transformer in transformations:
            result = self._apply_transformations(result, transformer)
        return result

    @staticmethod
    def _apply_transformations(input_var, transformer):
        if isinstance(transformer, Callable):
            return transformer(input_var)
        else:
            raise TypeError(f'Expected a Relationship or Callable, got {type(transformer)}')

    def value_player(self):
        """
        Returns the value player of the inline value map.
        """
        return self._value_player

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this inline value map.
        """
        return self._binding.column

    def table(self) -> BindableTable:
        """
        Returns the table associated with the inline value map.
        """
        return self.column().table

    def binding(self) -> Binding:
        """
        Returns the binding of the value map.
        """
        return self._binding

    def formula(self):
        """
        Constructs the body formula for the value map.
        This is used to generate the head and the body of the relation.
        """
        filter_atoms = where(self._binding.filter_by)
        return self._construct_body_atoms(filter_atoms)

    def __repr__(self):
        return f'@inline ValueMap({self._name})'


class ValueMap(RoleMap):
    """
    A class representing a value map relation.
    """

    def __init__(self, model, binding: Binding, role):
        # =
        # Skip QB relationship initialization if *inline* is True, as __call__ will return a QB expression
        # =
        madlib, name = self._handle_params(binding, role)
        super().__init__(madlib, model, name, role)
        self._binding = binding
        self._inline_value_map = InlineValueMap(model, binding, role)
        self._generate_body()

    @staticmethod
    def _handle_params(binding, role):
        if not role.player()._is_primitive():
            raise TypeError(f'Cannot construct a value map for {role}, concept is not a value type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(table_name, concept)

    def _generate_body(self):
        source_ref = self.table()
        val_ref, body_atoms = self._inline_value_map.formula()

        define(self(source_ref, val_ref)).where(*body_atoms)

    def value_player(self):
        return self._inline_value_map.value_player()

    def column(self) -> BindableAttribute:
        return self._inline_value_map.column()

    def binding(self) -> Binding:
        return self._inline_value_map.binding()

    def formula(self):
        source_ref = self.table()
        rez_val_ref = self.value_player().ref()
        return rez_val_ref, [self(source_ref, rez_val_ref)]

    def __repr__(self):
        return f'ValueMap({self._name})'


class MaterializedEntityMap(AbstractMap):
    """
    A class representing a materialized entity map relation.
    This is used to materialize an entity map from a set of role maps.
    """

    def __init__(self, model, inline_entity_map: 'AbstractInlineEntityMap'):
        super().__init__(inline_entity_map.madlib(), model, inline_entity_map.short_name())
        self._inline_entity_map = inline_entity_map
        self._generate_body()

    def _generate_body(self):
        # materialize the population of the entity map as necessary
        self._inline_entity_map.materialize_population()

        source_ref = self.table()
        rez_val_ref, body_formula = self._inline_entity_map.formula()
        where(*body_formula).define(self(source_ref, rez_val_ref))

    def value_player(self):
        """
        Returns the result value player of the materialized entity map.
        This is the same as the value player of the inline entity map.
        """
        return self._inline_entity_map.value_player()

    def column(self) -> BindableAttribute:
        return self._inline_entity_map.column()

    def table(self) -> BindableTable:
        """
        Returns the table associated with the materialized entity map.
        """
        return self._inline_entity_map.table()

    def __repr__(self):
        # replace @inline with @materialized
        return self._inline_entity_map.__repr__().replace('@inline', '@materialized')


class AbstractInlineEntityMap(ABC):
    """
    A non-QB relationship class that represents an inline entity map.
    This is used to construct entities from a set of role maps.
    """

    def __init__(self, model, madlib, name):
        self._model = model
        self._madlib = madlib
        self._name = name

    def __call__(self, *args, **kwargs):
        if bool(args):
            raise ValueError(f'Expected no arguments to a Value Map, got {len(args)}: {args}')

        return self._construct_body_formula()

    def madlib(self):
        """
        Returns the madlib of the inline entity map.
        """
        return self._madlib

    def short_name(self):
        """
        Returns the name of the inline entity map.
        """
        return self._name

    @abstractmethod
    def value_player(self):
        """
        Returns the value player of the role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    @abstractmethod
    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def table(self) -> BindableTable:
        """
        Returns the table associated with the role map.
        """
        return self.column().table

    def formula(self):
        return self._construct_body_formula()

    @abstractmethod
    def _construct_body_formula(self) -> tuple[Union[ConceptConstruct, ConceptFilter, Ref, RelationshipRef], list[typing.Any]]:
        """
        Constructs the body formula for the inline entity map.
        This is used to generate the body of the relation.
        """
        pass

    def materialize(self):
        return MaterializedEntityMap(self._model, self)

    @abstractmethod
    def materialize_population(self):
        """
        Materializes the population of the entity map, using the body of the entity map.
        This has to check the conditions and only update the population if those are met.
        """
        pass


class SimpleConstructorEntityMap(AbstractInlineEntityMap):
    """
    A class representing an entity map relation.
    """

    def __init__(self, model, binding: Binding, role: Role, identifier_to_role_map: OrderedDict[Unique, 'RoleMap']):
        madlib, name = self._handle_params(binding, role)
        super().__init__(model, madlib, name)
        self._binding = binding
        self._role = role
        self._identifier_to_role_map = identifier_to_role_map
        self._reference_role_map = list(identifier_to_role_map.items())[-1][1]  # last role map is the reference one

    @staticmethod
    def _handle_params(binding, role):
        if role.player()._is_primitive():
            raise TypeError(f'Cannot construct an entity map for {role}, concept is not an entity type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(table_name, concept)

    def _construct_body_formula(self):
        concept = self._role.player()
        value_ref, ref_body_atoms = self._reference_role_map.formula()
        entity_ref = reference_entity(concept, value_ref, check_population=False)
        body_atoms = [
            *ref_body_atoms
        ]
        return entity_ref, body_atoms

    def _concept_population_atom(self, entity):
        """
        Generates the *optional* population atom for the entity map.
        Only used for bindings that have a FK.
        """
        if self._should_reference():
            return where(
                self._role.player()(entity)
            )
        else:
            return where()

    def _should_reference(self) -> bool:
        return self._reference_role_map.column().references_column is not None

    def value_player(self):
        """
        Returns the value player of the entity map.
        """
        return self._role.player()

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity map.
        """
        return self._binding.column

    def materialize_population(self):
        concept = self._role.player()
        if not self._should_reference():
            # then create entities
            role_maps = self._identifier_to_role_map.values()
            map_value_refs, subformulas = [], []
            for role_map in role_maps:
                rez_val, subformula = role_map.formula()
                map_value_refs.append(rez_val)
                subformulas.extend(subformula)
            where(
                *subformulas
            ).define(construct_entity(concept, *map_value_refs))

    def __repr__(self):
        return f'@inline CtorEntityMap({self._name})'


class ReferentEntityMap(AbstractInlineEntityMap):
    """
    A class representing a referent entity map relation.
    """

    def __init__(self, model, binding: Binding, role: Role, constructing_role_map: 'Map'):
        madlib, name = self._handle_params(binding, role)
        super().__init__(model, madlib, name)
        self._binding = binding
        self._role = role
        self._constructing_role_map = constructing_role_map

    @staticmethod
    def _handle_params(binding, role):
        if role.player()._is_primitive():
            raise TypeError(f'Cannot construct an entity map for {role}, concept is not an entity type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(f'{table_name}_ref', concept)

    def _construct_body_formula(self):
        concept = self._role.player()
        orig_entity_ref, subformula_atoms = self._constructing_role_map.formula()

        entity_ref = reference_entity(concept, orig_entity_ref)
        return entity_ref, [
            *subformula_atoms,
        ]

    def value_player(self):
        return self._role.player()

    def binding(self) -> Binding:
        """
        Returns the binding of the entity map.
        """
        return self._binding

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity map.
        """
        return self._binding.column

    def materialize_population(self):
        """
        Referent entity maps do not need to materialize population, as they always reference existing entities.
        """
        pass

    def __repr__(self):
        return f'@inline ReferentEntityMap({self._name})'


class CompositeEntityMap(AbstractInlineEntityMap):
    """
    A class representing a composite entity map relation.

    Takes a sequence of EntityMaps and constructs a composite entity type.
    """

    def __init__(self, model, entity_concept: qb.Concept, *role_maps: 'Map'):
        madlib, name = self._handle_params(entity_concept, *role_maps)
        super().__init__(model, madlib, name)
        self._entity_concept = entity_concept
        self._role_maps = list(role_maps)

    def _handle_params(self, entity_concept: qb.Concept, *role_maps: 'Map'):
        if entity_concept._is_primitive():
            raise TypeError(f'Cannot construct a composite entity map for {entity_concept},'
                            f' concept is not an entity type')
        if len(role_maps) < 2:
            raise ValueError('CompositeEntityMap requires at least two EntityMaps')

        role_map = role_maps[0]
        table = role_map.table()
        self._table = table
        return _get_map_madlib_and_name(table.physical_name(), entity_concept)

    def _construct_body_formula(self):
        var_refs, body_atoms = self._body_formula()

        # the rule to populate the entity map
        entity_ref = reference_entity(self._entity_concept, *var_refs, check_population=False)
        return entity_ref, body_atoms

    @staticmethod
    def _should_reference() -> bool:
        """
        Determines if the composite entity map should reference a column.
        Currently not supported fully.
        """
        return False

    def _body_formula(self):
        var_refs, subformulas = zip(*(role_map.formula() for role_map in self._role_maps))
        return var_refs, [
            *(atom for subformula in subformulas for atom in subformula)
        ]

    def value_player(self):
        """
        Returns the value player of the composite entity map.
        """
        return self._entity_concept

    def table(self) -> BindableTable:
        """
        Returns the table associated with the composite entity map.
        """
        return self._table

    def column(self) -> BindableAttribute:
        raise ValueError('CompositeEntityMap does not index a single column')

    def materialize_population(self):
        var_refs, body_atoms = self._body_formula()
        if not self._should_reference():
            # construct entities
            entity_ref = construct_entity(self._entity_concept, *var_refs)
            where(
                *body_atoms,
            ).define(entity_ref)

    def __repr__(self):
        return f'@inline CompositeEntityMap({self._name})'


AbstractEntityMap = Union['SimpleConstructorEntityMap', 'ReferentEntityMap', 'CompositeEntityMap', 'UnionEntityMap']
Map = Union['ValueMap', 'InlineValueMap', 'UnionEntityMap', 'AbstractInlineEntityMap', 'MaterializedEntityMap']

class EntitySubtypeMap(AbstractInlineEntityMap):
    """
    A class representing an entity subtype map relation.
    """

    def __init__(self, model, binding: Union[IdentifierConceptBinding, SubtypeConceptBinding, ReferentConceptBinding],
                 ctor_entity_map: Union[AbstractInlineEntityMap, MaterializedEntityMap]):
        # type of the entity subtype map is the parent type coming from the ctor_entity_map
        madlib, name = self._handle_params(binding, ctor_entity_map.value_player())
        super().__init__(model, madlib, name)
        self._binding = binding
        self._ctor_entity_map = ctor_entity_map

    @staticmethod
    def _handle_params(binding: Union[IdentifierConceptBinding, SubtypeConceptBinding, ReferentConceptBinding],
                       result_concept: qb.Concept):
        table = binding.column.table
        entity_concept = binding.entity_type
        # check that the entity concept is a subtype of the result concept
        assert entity_concept._isa(result_concept), f"Entity concept {entity_concept} must be a subtype of result concept {result_concept}"
        name = f'{table.physical_name()}_row_to_{entity_concept}'
        madlib = f'{name} {{row_id:RowId}} {{val:{result_concept}}}'
        return madlib, name

    def _construct_body_formula(self):
        parent_type_ref, body_atoms = self._generate_body_atoms()

        # note: have to return the parent type here, as otherwise it gets a join on the subtype population
        #       but doesn't cast to the type of the subtype...
        return parent_type_ref, body_atoms

    def _generate_body_atoms(self):
        filtering_atom = where(self._binding.filter_by)
        parent_type_ref, subformula_atoms = self._ctor_entity_map.formula()
        return parent_type_ref, [
            *subformula_atoms,
            filtering_atom
        ]

    def binding(self):
        """
        Returns the binding of the entity subtype map.
        """
        return self._binding

    def value_player(self):
        """
        Returns the subtype of the entity subtype map.
        """
        return self._binding.entity_type

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity subtype map.
        """
        return self._binding.column

    def table(self) -> BindableTable:
        """
        Returns the table associated with the entity subtype map.
        """
        return self._binding.column.table

    def materialize_population(self):
        subtype = self._binding.entity_type
        parent_type_ref, body_atoms = self._generate_body_atoms()
        # TODO instead of Person.new maybe try Employee.new to construct these?
        # derive subtype population
        where(*body_atoms).define(subtype(parent_type_ref))

    def __repr__(self):
        return f'@inline EntitySubtypeMap({self._name})'


class UnionEntityMap(AbstractMap):
    """
    A class representing a union entity map relation.
    """

    def __init__(self, model, entity_concept: qb.Concept, *entity_maps: AbstractEntityMap, generate_population: bool=False):
        madlib, name = self._handle_params(entity_concept, *entity_maps)
        super().__init__(madlib, model, name)
        self._entity_type = entity_concept
        self._entity_maps = OrderedSet().update(entity_maps)
        self._generate_population = generate_population
        self._generate_body()

    def _handle_params(self, entity_concept: qb.Concept, *entity_maps: AbstractEntityMap):
        if len(entity_maps) == 0:
            raise ValueError('UnionEntityMap requires at least one EntityMap')
        # pick an arbitrary entity map to get the table, as they must all have the same
        table = entity_maps[0].table()
        self._table = table
        return _get_map_madlib_and_name(table.physical_name(), entity_concept)

    def _generate_body(self):
        for entity_map in self._entity_maps:
            self._generate_body_rule(entity_map)

    def _generate_body_rule(self, entity_map: AbstractEntityMap):
        source_ref = self.table()
        rez_val_ref, subformula_atoms = entity_map.formula()

        # derive union entity map
        where(
            *subformula_atoms
        ).define(self(source_ref, rez_val_ref))

    def value_player(self):
        """
        Returns the type of the entity map.
        """
        return self._entity_type

    def table(self):
        """
        Returns the table associated with the entity map.
        """
        return self._table

    def update(self, entity_map: AbstractEntityMap):
        """
        Updates the union entity map with a new entity map.
        """
        if entity_map in self._entity_maps:
            return
        self._entity_maps.add(entity_map)
        self._generate_body_rule(entity_map)

    def materialize_population(self):
        if self._generate_population:
            # derive type population
            for entity_map in self._entity_maps:
                rez_val_ref, subformula_atoms = entity_map.formula()
                where(
                    *subformula_atoms
                ).define(self._entity_type(rez_val_ref))

    def __repr__(self):
        return f'UnionEntityMap({self._name})'
