"""
CompositeType provides means to interact with
`PostgreSQL composite types`_. Currently this type features:

* Easy attribute access to composite type fields
* Supports SQLAlchemy TypeDecorator types
* Ability to include composite types as part of PostgreSQL arrays
* Type creation and dropping

Installation
^^^^^^^^^^^^

CompositeType automatically attaches `before_create` and `after_drop` DDL
listeners. These listeners create and drop the composite type in the
database. This means it works out of the box in your test environment where
you create the tables on each test run.

When you already have your database set up you should call
:func:`register_composites` after you've set up all models.

::

    register_composites(conn)



Usage
^^^^^

::

    from collections import OrderedDict

    import sqlalchemy as sa
    from sqlalchemy_utils import CompositeType, CurrencyType


    class Account(Base):
        __tablename__ = 'account'
        id = sa.Column(sa.Integer, primary_key=True)
        balance = sa.Column(
            CompositeType(
                'money_type',
                [
                    sa.Column('currency', CurrencyType),
                    sa.Column('amount', sa.Integer)
                ]
            )
        )


Creation
~~~~~~~~
When creating CompositeType, you can either pass in a tuple or a dictionary.

::
    account1 = Account()
    account1.balance = ('USD', 15)

    account2 = Account()
    account2.balance = {'currency': 'USD', 'amount': 15}

    session.add(account1)
    session.add(account2)
    session.commit()


Accessing fields
^^^^^^^^^^^^^^^^

CompositeType provides attribute access to underlying fields. In the following
example we find all accounts with balance amount more than 5000.


::

    session.query(Account).filter(Account.balance.amount > 5000)


Arrays of composites
^^^^^^^^^^^^^^^^^^^^

::

    from sqlalchemy.dialects.postgresql import ARRAY


    class Account(Base):
        __tablename__ = 'account'
        id = sa.Column(sa.Integer, primary_key=True)
        balances = sa.Column(
            ARRAY(
                CompositeType(
                    'money_type',
                    [
                        sa.Column('currency', CurrencyType),
                        sa.Column('amount', sa.Integer)
                    ]
                ),
                dimensions=1
            )
        )


.. _PostgreSQL composite types:
    https://www.postgresql.org/docs/current/rowtypes.html


Related links:

https://schinckel.net/2014/09/24/using-postgres-composite-types-in-django/
"""
from collections import namedtuple

import sqlalchemy as sa
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.schema import _CreateDropBase
from sqlalchemy.sql.expression import FunctionElement
from sqlalchemy.types import (
    SchemaType,
    to_instance,
    TypeDecorator,
    UserDefinedType
)

from .. import ImproperlyConfigured

psycopg2 = None
CompositeCaster = None
adapt = None
AsIs = None
register_adapter = None
try:
    import psycopg2
    from psycopg2.extensions import adapt, AsIs, register_adapter
    from psycopg2.extras import CompositeCaster
except ImportError:
    pass


class CompositeElement(FunctionElement):
    """
    Instances of this class wrap a Postgres composite type.
    """
    def __init__(self, base, field, type_):
        self.name = field
        self.type = to_instance(type_)

        super().__init__(base)


@compiles(CompositeElement)
def _compile_pgelem(expr, compiler, **kw):
    return f'({compiler.process(expr.clauses, **kw)}).{expr.name}'


# TODO: Make the registration work on connection level instead of global level
registered_composites = {}


class CompositeType(UserDefinedType, SchemaType):
    """
    Represents a PostgreSQL composite type.

    :param name:
        Name of the composite type.
    :param columns:
        List of columns that this composite type consists of
    """
    python_type = tuple

    class comparator_factory(UserDefinedType.Comparator):
        def __getattr__(self, key):
            try:
                type_ = self.type.typemap[key]
            except KeyError:
                raise KeyError(
                    "Type '{}' doesn't have an attribute: '{}'".format(
                        self.name, key
                    )
                )

            return CompositeElement(self.expr, key, type_)

    def __init__(self, name, columns, quote=None, **kwargs):
        if psycopg2 is None:
            raise ImproperlyConfigured(
                "'psycopg2' package is required in order to use CompositeType."
            )
        SchemaType.__init__(
            self,
            name=name,
            quote=quote
        )
        self.columns = columns
        if name in registered_composites:
            self.type_cls = registered_composites[name].type_cls
        else:
            self.type_cls = namedtuple(
                self.name, [c.name for c in columns]
            )
        registered_composites[name] = self

        class Caster(CompositeCaster):
            def make(obj, values):
                return self.type_cls(*values)

        self.caster = Caster
        attach_composite_listeners()

    def get_col_spec(self):
        return self.name

    def bind_processor(self, dialect):
        def process(value):
            if value is None:
                return None

            processed_value = []
            for i, column in enumerate(self.columns):
                current_value = (
                    value.get(column.name)
                    if isinstance(value, dict)
                    else value[i]
                )

                if isinstance(column.type, TypeDecorator):
                    processed_value.append(
                        column.type.process_bind_param(
                            current_value, dialect
                        )
                    )
                else:
                    processed_value.append(current_value)
            return self.type_cls(*processed_value)
        return process

    def result_processor(self, dialect, coltype):
        def process(value):
            if value is None:
                return None
            cls = value.__class__
            kwargs = {}
            for column in self.columns:
                if isinstance(column.type, TypeDecorator):
                    kwargs[column.name] = column.type.process_result_value(
                        getattr(value, column.name), dialect
                    )
                else:
                    kwargs[column.name] = getattr(value, column.name)
            return cls(**kwargs)
        return process

    def create(self, bind=None, checkfirst=None):
        if (
            not checkfirst or
            not bind.dialect.has_type(bind, self.name, schema=self.schema)
        ):
            bind.execute(CreateCompositeType(self))

    def drop(self, bind=None, checkfirst=True):
        if (
            checkfirst and
            bind.dialect.has_type(bind, self.name, schema=self.schema)
        ):
            bind.execute(DropCompositeType(self))


def register_psycopg2_composite(dbapi_connection, composite):
    psycopg2.extras.register_composite(
        composite.name,
        dbapi_connection,
        globally=True,
        factory=composite.caster
    )

    def adapt_composite(value):
        dialect = PGDialect_psycopg2()
        adapted = [
            adapt(
                getattr(value, column.name)
                if not isinstance(column.type, TypeDecorator)
                else column.type.process_bind_param(
                    getattr(value, column.name),
                    dialect
                )
            )
            for column in
            composite.columns
        ]
        for value in adapted:
            if hasattr(value, 'prepare'):
                value.prepare(dbapi_connection)
        values = [
            value.getquoted().decode(dbapi_connection.encoding)
            for value in adapted
        ]
        return AsIs(
            '({})::{}'.format(
                ', '.join(values),
                dialect.identifier_preparer.quote(composite.name)
            )
        )

    register_adapter(composite.type_cls, adapt_composite)


def get_driver_connection(connection):
    try:
        # SQLAlchemy 2.0
        return connection.connection.driver_connection
    except AttributeError:
        return connection.connection.connection


def before_create(target, connection, **kw):
    for name, composite in registered_composites.items():
        composite.create(connection, checkfirst=True)
        register_psycopg2_composite(
            get_driver_connection(connection),
            composite
        )


def after_drop(target, connection, **kw):
    for name, composite in registered_composites.items():
        composite.drop(connection, checkfirst=True)


def register_composites(connection):
    for name, composite in registered_composites.items():
        register_psycopg2_composite(
            get_driver_connection(connection),
            composite
        )


def attach_composite_listeners():
    listeners = [
        (sa.MetaData, 'before_create', before_create),
        (sa.MetaData, 'after_drop', after_drop),
    ]
    for listener in listeners:
        if not sa.event.contains(*listener):
            sa.event.listen(*listener)


def remove_composite_listeners():
    listeners = [
        (sa.MetaData, 'before_create', before_create),
        (sa.MetaData, 'after_drop', after_drop),
    ]
    for listener in listeners:
        if sa.event.contains(*listener):
            sa.event.remove(*listener)


class CreateCompositeType(_CreateDropBase):
    pass


@compiles(CreateCompositeType)
def _visit_create_composite_type(create, compiler, **kw):
    type_ = create.element
    fields = ', '.join(
        '{name} {type}'.format(
            name=column.name,
            type=compiler.dialect.type_compiler.process(
                to_instance(column.type)
            )
        )
        for column in type_.columns
    )

    return 'CREATE TYPE {name} AS ({fields})'.format(
        name=compiler.preparer.format_type(type_),
        fields=fields
    )


class DropCompositeType(_CreateDropBase):
    pass


@compiles(DropCompositeType)
def _visit_drop_composite_type(drop, compiler, **kw):
    type_ = drop.element

    return f'DROP TYPE {compiler.preparer.format_type(type_)}'