# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida_core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""
The dummy model encodes the model defined by django in backends.djsite
using SQLAlchemy.
This is done to query the database with more performant ORM of SA.
"""

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from sqlalchemy.ext.declarative import declarative_base

from sqlalchemy import (
    Column, Table, ForeignKey, UniqueConstraint, create_engine,
    select, func, join, and_, or_, not_, except_, case, exists,
    text
)

from sqlalchemy.types import (
    Integer, String, DateTime, Float, Boolean, Text,
)
from sqlalchemy.orm import (
    relationship,
    backref,
    sessionmaker,
    foreign, mapper, aliased
)

from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm.attributes import InstrumentedAttribute
from sqlalchemy.sql.elements import Cast
from sqlalchemy.dialects.postgresql import UUID, JSONB, INTEGER, array
# TO COMPILE MY OWN FUNCTIONALITIES:
from sqlalchemy.sql.expression import FunctionElement, cast
from sqlalchemy.sql.base import ImmutableColumnCollection
from sqlalchemy.ext.compiler import compiles

# SETTINGS:
from aiida.common.setup import get_profile_config
from aiida.backends import settings

# EXCEPTIONS
from aiida.common.exceptions import DbContentError, MissingPluginError

# MISC
from aiida.backends.sqlalchemy.models.utils import uuid_func
from aiida.utils import timezone
from aiida.common.datastructures import (calc_states, _sorted_datastates,
                                         sort_states)

Base = declarative_base()


class DbLink(Base):
    __tablename__ = "db_dblink"
    id = Column(Integer, primary_key=True)
    input_id = Column(
        Integer,
        ForeignKey('db_dbnode.id', deferrable=True, initially="DEFERRED")
    )
    output_id = Column(
        Integer,
        ForeignKey(
            'db_dbnode.id',
            ondelete="CASCADE",
            deferrable=True,
            initially="DEFERRED"
        )
    )
    type = Column(String(255))
    input = relationship("DbNode", primaryjoin="DbLink.input_id == DbNode.id")
    output = relationship("DbNode", primaryjoin="DbLink.output_id == DbNode.id")
    label = Column(String(255), index=True, nullable=False)


class DbCalcState(Base):
    __tablename__ = "db_dbcalcstate"
    id = Column(Integer, primary_key=True)
    dbnode = relationship(
        'DbNode',
        backref=backref('dbstates', passive_deletes=True)
    )
    state = Column(String(255))
    time = Column(DateTime(timezone=True), default=timezone.now)
    dbnode_id = Column(
        Integer,
        ForeignKey(
            'db_dbnode.id', ondelete="CASCADE",
            deferrable=True, initially="DEFERRED"
        )
    )
    __table_args__ = (
        UniqueConstraint('dbnode_id', 'state'),
    )


class DbAttribute(Base):
    __tablename__ = "db_dbattribute"
    id = Column(Integer, primary_key=True)
    dbnode_id = Column(Integer, ForeignKey('db_dbnode.id'))
    key = Column(String(255))
    datatype = Column(String(10))
    tval = Column(String, default='')
    fval = Column(Float, default=None, nullable=True)
    ival = Column(Integer, default=None, nullable=True)
    bval = Column(Boolean, default=None, nullable=True)
    dval = Column(DateTime, default=None, nullable=True)


class DbExtra(Base):
    __tablename__ = "db_dbextra"
    id = Column(Integer, primary_key=True)
    dbnode_id = Column(Integer, ForeignKey('db_dbnode.id'))
    key = Column(String(255))
    datatype = Column(String(10))
    tval = Column(String, default='')
    fval = Column(Float, default=None, nullable=True)
    ival = Column(Integer, default=None, nullable=True)
    bval = Column(Boolean, default=None, nullable=True)
    dval = Column(DateTime, default=None, nullable=True)


class DbComputer(Base):
    __tablename__ = "db_dbcomputer"

    id = Column(Integer, primary_key=True)

    uuid = Column(UUID(as_uuid=True), default=uuid_func)
    name = Column(String(255), unique=True, nullable=False)
    hostname = Column(String(255))

    description = Column(Text, nullable=True)
    enabled = Column(Boolean)

    transport_type = Column(String(255))
    scheduler_type = Column(String(255))

    transport_params = Column(String(255))
    _metadata = Column('metadata', String(255), default="{}")

    def get_aiida_class(self):
        from aiida.backends.djsite.db.models import DbComputer as DjangoSchemaDbComputer
        djcomputer = DjangoSchemaDbComputer(
            id=self.id, uuid=self.uuid, name=self.name,
            hostname=self.hostname, description=self.description,
            enabled=self.enabled, transport_type=self.transport_type,
            scheduler_type=self.scheduler_type, transport_params=self.transport_params,
            metadata=self._metadata
        )
        return djcomputer.get_aiida_class()


class DbUser(Base):
    __tablename__ = "db_dbuser"

    id = Column(Integer, primary_key=True)
    email = Column(String(254), unique=True, index=True)
    password = Column(String(128))  # Clear text password ?
    first_name = Column(String(254), nullable=True)
    last_name = Column(String(254), nullable=True)
    institution = Column(String(254), nullable=True)

    is_staff = Column(Boolean, default=False)
    is_active = Column(Boolean, default=False)

    last_login = Column(DateTime(timezone=True), default=timezone.now)
    date_joined = Column(DateTime(timezone=True), default=timezone.now)

    def get_aiida_class(self):
        from aiida.backends.djsite.db.models import DbUser as DjangoSchemaDbUser
        djuser = DjangoSchemaDbUser(
            id=self.id, email=self.email, password=self.password,
            first_name=self.first_name, last_name=self.last_name,
            institution=self.institution, is_staff=self.is_staff,
            is_active=self.is_active, last_login=self.last_login,
            date_joined=self.date_joined
        )
        return djuser.get_aiida_class()


table_groups_nodes = Table(
    'db_dbgroup_dbnodes',
    Base.metadata,
    Column('id', Integer, primary_key=True),
    Column(
        'dbnode_id', Integer,
        ForeignKey('db_dbnode.id', deferrable=True, initially="DEFERRED")
    ),
    Column(
        'dbgroup_id', Integer,
        ForeignKey('db_dbgroup.id', deferrable=True, initially="DEFERRED")
    )
)


class DbGroup(Base):
    __tablename__ = "db_dbgroup"

    id = Column(Integer, primary_key=True)

    uuid = Column(UUID(as_uuid=True), default=uuid_func)
    name = Column(String(255), index=True)

    type = Column(String(255), default="", index=True)

    time = Column(DateTime(timezone=True), default=timezone.now)
    description = Column(Text, nullable=True)

    user_id = Column(
        Integer,
        ForeignKey(
            'db_dbuser.id', ondelete='CASCADE',
            deferrable=True, initially="DEFERRED")
    )
    user = relationship('DbUser', backref=backref('dbgroups', cascade='merge'))

    dbnodes = relationship('DbNode', secondary=table_groups_nodes,
                           backref="dbgroups", lazy='dynamic')

    __table_args__ = (
        UniqueConstraint('name', 'type'),
    )

    def __str__(self):
        if self.type:
            return '<DbGroup [type: {}] "{}">'.format(self.type, self.name)
        else:
            return '<DbGroup [user-defined] "{}">'.format(self.name)

    def get_aiida_class(self):
        from aiida.orm.implementation.django.group import Group as DjangoAiidaGroup
        from aiida.backends.djsite.db.models import DbGroup as DjangoSchemaDbGroup
        dbgroup = DjangoSchemaDbGroup(
            id=self.id,
            type=self.type,
            uuid=self.uuid,
            name=self.name,
            time=self.time,
            description=self.description,
            user_id=self.user_id,
        )

        return DjangoAiidaGroup(dbgroup=dbgroup)


class DbNode(Base):
    __tablename__ = "db_dbnode"
    id = Column(Integer, primary_key=True)
    uuid = Column(UUID(as_uuid=True), default=uuid_func)
    type = Column(String(255), index=True)
    process_type = Column(String(255), index=True)
    label = Column(String(255), index=True, nullable=True)
    description = Column(Text(), nullable=True)
    ctime = Column(DateTime(timezone=True), default=timezone.now)
    mtime = Column(DateTime(timezone=True), default=timezone.now)
    dbcomputer_id = Column(
        Integer,
        ForeignKey('db_dbcomputer.id', deferrable=True, initially="DEFERRED"),
        nullable=True
    )
    dbcomputer = relationship(
        'DbComputer',
        backref=backref('dbnodes', passive_deletes=True)
    )
    user_id = Column(
        Integer,
        ForeignKey('db_dbuser.id', deferrable=True, initially="DEFERRED"),
        nullable=False
    )
    user = relationship('DbUser', backref='dbnodes')

    public = Column(Boolean, default=False)

    nodeversion = Column(Integer, default=1)

    attributes = relationship('DbAttribute', uselist=True, backref='dbnode')
    extras = relationship('DbExtra', uselist=True, backref='dbnode')

    outputs = relationship(
        "DbNode",
        secondary="db_dblink",
        primaryjoin="DbNode.id == DbLink.input_id",
        secondaryjoin="DbNode.id == DbLink.output_id",
        backref=backref("inputs", passive_deletes=True),
        passive_deletes=True
    )

    def get_aiida_class(self):
        """
        Return the corresponding instance of
        :func:`~aiida.orm.implementation.django.node.Node`
        or a subclass return by the plugin loader.

        .. todo::
            The behavior is quite pathetic, creating a django DbNode instance
            to instantiate the aiida instance.
            These means that every time you load Aiida instances with
            the QueryBuilder when using Django as a backend, three instances
            are instantiated for every Aiida instance you load!
            Could be fixed by allowing DbNode from the dummy nodel to be passed
            to AiidaNode's __init__.

        :returns: An instance of the plugin class
        """
        # I need to import the DbNode in the Django model,
        # and instantiate an object that has the same attributes as self.
        from aiida.backends.djsite.db.models import DbNode as DjangoSchemaDbNode
        dbnode = DjangoSchemaDbNode(
            id=self.id, type=self.type, process_type=self.process_type, uuid=self.uuid, ctime=self.ctime,
            mtime=self.mtime, label=self.label, description=self.description, dbcomputer_id=self.dbcomputer_id,
            user_id=self.user_id, public=self.public, nodeversion=self.nodeversion
        )
        return dbnode.get_aiida_class()

    @hybrid_property
    def user_email(self):
        """
        Returns: the email of the user
        """
        return self.user.email

    @user_email.expression
    def user_email(cls):
        """
        Returns: the email of the user at a class level (i.e. in the database)
        """
        return select([DbUser.email]).where(DbUser.id == cls.user_id).label(
            'user_email')

    # Computer name
    @hybrid_property
    def computer_name(self):
        """
        Returns: the of the computer
        """
        return self.dbcomputer.name

    @computer_name.expression
    def computer_name(cls):
        """
        Returns: the name of the computer at a class level (i.e. in the 
        database)
        """
        return select([DbComputer.name]).where(DbComputer.id ==
                                               cls.dbcomputer_id).label(
            'computer_name')

    # State
    @hybrid_property
    def state(self):
        """
        Return the most recent state from DbCalcState
        """
        if not self.id:
            return None
        all_states = DbCalcState.query.filter(DbCalcState.dbnode_id == self.id).all()
        if all_states:
            # return max((st.time, st.state) for st in all_states)[1]
            return sort_states(((dbcalcstate.state, dbcalcstate.state.value)
                                for dbcalcstate in all_states),
                               use_key=True)[0]
        else:
            return None

    @state.expression
    def state(cls):
        """
        Return the expression to get the 'latest' state from DbCalcState,
        to be used in queries, where 'latest' is defined using the state order
        defined in _sorted_datastates.
        """
        # Sort first the latest states
        whens = {
            v: idx for idx, v
            in enumerate(_sorted_datastates[::-1], start=1)}
        custom_sort_order = case(value=DbCalcState.state,
                                 whens=whens,
                                 else_=100)  # else: high value to put it at the bottom

        # Add numerical state to string, to allow to sort them
        states_with_num = select([
            DbCalcState.id.label('id'),
            DbCalcState.dbnode_id.label('dbnode_id'),
            DbCalcState.state.label('state_string'),
            custom_sort_order.label('num_state')
        ]).select_from(DbCalcState).alias()

        # Get the most 'recent' state (using the state ordering, and the min function) for
        # each calc
        calc_state_num = select([
            states_with_num.c.dbnode_id.label('dbnode_id'),
            func.min(states_with_num.c.num_state).label('recent_state')
        ]).group_by(states_with_num.c.dbnode_id).alias()

        # Join the most-recent-state table with the DbCalcState table
        all_states_q = select([
            DbCalcState.dbnode_id.label('dbnode_id'),
            DbCalcState.state.label('state_string'),
            calc_state_num.c.recent_state.label('recent_state'),
            custom_sort_order.label('num_state'),
        ]).select_from(  # DbCalcState).alias().join(
            join(DbCalcState, calc_state_num, DbCalcState.dbnode_id == calc_state_num.c.dbnode_id)).alias()

        # Get the association between each calc and only its corresponding most-recent-state row
        subq = select([
            all_states_q.c.dbnode_id.label('dbnode_id'),
            all_states_q.c.state_string.label('state')
        ]).select_from(all_states_q).where(all_states_q.c.num_state == all_states_q.c.recent_state).alias()

        # Final filtering for the actual query
        return select([subq.c.state]). \
            where(
            subq.c.dbnode_id == cls.id,
        ). \
            label('laststate')


profile = get_profile_config(settings.AIIDADB_PROFILE)


def get_aldjemy_session():
    """
    Use aldjemy to make a session

    .. note:
        Use only in this case. In normal production mode
        it is safer make session explictly because it is more robust
    """
    from aldjemy.core import get_engine
    engine = get_engine()
    _Session = sessionmaker(bind=engine)
    return _Session()


session = get_aldjemy_session()
