# -*- coding: utf-8 -*-
"""
    pyrseas.dbobject.schema
    ~~~~~~~~~~~~~~~~~~~~~~~

    This defines two classes, Schema and SchemaDict, derived from
    DbObject and DbObjectDict, respectively.
"""
import os

from pyrseas.yamlutil import yamldump
from . import DbObjectDict, DbObject
from . import quote_id, commentable, ownable, grantable
from .dbtype import BaseType, Composite, Domain, Enum, Range
from .table import Table, Sequence
from .view import View, MaterializedView


class Schema(DbObject):
    """A database schema definition, i.e., a named collection of tables,
    views, triggers and other schema objects."""

    keylist = ['name']
    catalog = 'pg_namespace'

    @property
    def allprivs(self):
        return 'UC'

    def __init__(self, name, description=None, owner=None, privileges=[],
                 oldname=None, oid=None):
        """Initialize the schema

        :param name: schema name (from nspname)
        :param description: comment text (from obj_description())
        :param owner: owner name (from rolname via nspowner)
        :param privileges: access privileges (from nspacl)
        :param oldname: previous name of schema
        """
        super(Schema, self).__init__(name, description)
        self._init_own_privs(owner, privileges)
        self.oldname = None

    @staticmethod
    def query(dbversion=None):
        return """
            SELECT nspname AS name, rolname AS owner,
                   array_to_string(nspacl, ',') AS privileges,
                   obj_description(n.oid, 'pg_namespace') AS description, n.oid
            FROM pg_namespace n
                 JOIN pg_roles r ON (r.oid = nspowner)
            WHERE nspname NOT IN ('information_schema', 'pg_toast')
                  AND nspname NOT LIKE 'pg_temp\_%'
                  AND nspname NOT LIKE 'pg_toast_temp\_%'
            ORDER BY nspname"""

    @staticmethod
    def from_map(name, inobj):
        """Initialize a schema instance from a YAML map

        :param name: schema name
        :param inobj: YAML map of the schema
        :return: schema instance
        """
        obj = Schema(
            name, inobj.pop('description', None), inobj.pop('owner', None),
            inobj.pop('privileges', []))
        obj.fix_privileges()
        obj.set_oldname(inobj)
        return obj

    def extern_dir(self, root='.'):
        """Return the path to a directory to hold the schema objects.

        :return: directory path
        """
        (dir, ext) = os.path.splitext(os.path.join(root,
                                                   self.extern_filename()))
        return dir

    def to_map(self, db, dbschemas, opts):
        """Convert tables, etc., dictionaries to a YAML-suitable format

        :param dbschemas: dictionary of schemas
        :param opts: options to include/exclude schemas/tables, etc.
        :return: dictionary
        """
        if self.name == 'pyrseas':
            return {}
        no_owner = opts.no_owner
        no_privs = opts.no_privs
        schbase = {} if no_owner else {'owner': self.owner}
        if not no_privs:
            schbase.update({'privileges': self.map_privs()})
        if self.description is not None:
            schbase.update(description=self.description)

        schobjs = []
        seltbls = getattr(opts, 'tables', [])
        if hasattr(self, 'tables'):
            for objkey in self.tables:
                if not seltbls or objkey in seltbls:
                    obj = self.tables[objkey]
                    schobjs.append((obj, obj.to_map(db, dbschemas, opts)))

        def mapper(objtypes):
            if hasattr(self, objtypes):
                schemadict = getattr(self, objtypes)
                for objkey in schemadict:
                    if objtypes == 'sequences' or (
                            not seltbls or objkey in seltbls):
                        obj = schemadict[objkey]
                        schobjs.append((obj, obj.to_map(db, opts)))

        for objtypes in ['ftables', 'sequences', 'views', 'matviews']:
            mapper(objtypes)

        def mapper2(objtypes, privs=False):
            if hasattr(self, objtypes):
                schemadict = getattr(self, objtypes)
                for objkey in schemadict:
                    obj = schemadict[objkey]
                    if privs:
                        dct = obj.to_map(db, no_owner, no_privs)
                    else:
                        dct = obj.to_map(db, no_owner)
                    schobjs.append((obj, dct))

        if hasattr(opts, 'tables') and not opts.tables or \
                not hasattr(opts, 'tables'):
            for objtypes in ('conversions', 'collations', 'operators',
                             'operclasses', 'operfams', 'tsconfigs',
                             'tsdicts', 'tsparsers', 'tstempls'):
                mapper2(objtypes)
            for objtypes in ('types', 'domains'):
                mapper2(objtypes, True)
            if hasattr(self, 'functions'):
                for objkey in self.functions:
                    obj = self.functions[objkey]
                    schobjs.append((obj, obj.to_map(db, no_owner, no_privs)))

        # special case for pg_catalog schema
        if self.name == 'pg_catalog' and not schobjs:
            return {}

        if hasattr(self, 'datacopy') and self.datacopy:
            dir = self.extern_dir(opts.data_dir)
            if not os.path.exists(dir):
                os.mkdir(dir)
            for tbl in self.datacopy:
                self.tables[tbl].data_export(dbschemas.dbconn, dir)

        if opts.multiple_files:
            dir = self.extern_dir(opts.metadata_dir)
            if not os.path.exists(dir):
                os.mkdir(dir)
            filemap = {}
            for obj, objmap in schobjs:
                if objmap is not None:
                    extkey = obj.extern_key()
                    filepath = os.path.join(dir, obj.extern_filename())
                    with open(filepath, 'a') as f:
                        f.write(yamldump({extkey: objmap}))
                    outobj = {extkey:
                              os.path.relpath(filepath, opts.metadata_dir)}
                    filemap.update(outobj)
            # always write the schema YAML file
            filepath = self.extern_filename()
            extkey = self.extern_key()
            with open(os.path.join(opts.metadata_dir, filepath), 'a') as f:
                f.write(yamldump({extkey: schbase}))
            filemap.update(schema=filepath)
            return {extkey: filemap}

        schmap = dict((obj.extern_key(), objmap) for obj, objmap in schobjs
                      if objmap is not None)
        schmap.update(schbase)
        return {self.extern_key(): schmap}

    @commentable
    @grantable
    @ownable
    def create(self, dbversion=None):
        """Return SQL statements to CREATE the schema

        :return: SQL statements
        """
        # special case for --revert
        if self.name == 'pg_catalog':
            return []
        return ["CREATE SCHEMA %s" % quote_id(self.name)]

    def data_import(self, opts):
        """Generate SQL to import data from the tables in this schema

        :param opts: options to include/exclude schemas/tables, etc.
        :return: list of SQL statements
        """
        stmts = []
        if hasattr(self, 'datacopy') and self.datacopy:
            dir = self.extern_dir(opts.data_dir)
            for tbl in self.datacopy:
                stmts.append(self.tables[tbl].data_import(dir))
        return stmts

    def drop(self):
        if self.name not in ('public', 'pg_catalog'):
            return super(Schema, self).drop()
        else:
            return []


PREFIXES = {'domain ': 'types', 'type': 'types', 'table ': 'tables',
            'view ': 'tables', 'sequence ': 'tables',
            'materialized view ': 'tables',
            'function ': 'functions', 'aggregate ': 'functions',
            'operator family ': 'operfams', 'operator class ': 'operclasses',
            'conversion ': 'conversions', 'text search dictionary ': 'tsdicts',
            'text search template ': 'tstempls',
            'text search parser ': 'tsparsers',
            'text search configuration ': 'tsconfigs',
            'foreign table ': 'ftables', 'collation ': 'collations'}
SCHOBJS1 = ['types', 'tables', 'ftables']
SCHOBJS2 = ['collations', 'conversions', 'functions', 'operators',
            'operclasses', 'operfams', 'tsconfigs', 'tsdicts', 'tsparsers',
            'tstempls']


class SchemaDict(DbObjectDict):
    "The collection of schemas in a database.  Minimally, the 'public' schema."

    cls = Schema

    def from_map(self, inmap, newdb):
        """Initialize the dictionary of schemas by converting the input map

        :param inmap: the input YAML map defining the schemas
        :param newdb: collection of dictionaries defining the database

        Starts the recursive analysis of the input map and
        construction of the internal collection of dictionaries
        describing the database objects.
        """
        for key in inmap:
            (objtype, spc, sch) = key.partition(' ')
            if spc != ' ' or objtype != 'schema':
                raise KeyError("Unrecognized object type: %s" % key)
            inobj = inmap[key]
            schema = self[sch] = Schema.from_map(sch, inobj)
            objdict = {}
            for key in sorted(inobj.keys()):
                mapped = False
                for prefix in PREFIXES:
                    if key.startswith(prefix):
                        otype = PREFIXES[prefix]
                        if otype not in objdict:
                            objdict[otype] = {}
                        objdict[otype].update({key: inobj[key]})
                        mapped = True
                        break
                # Needs separate processing because it overlaps
                # operator classes and operator families
                if not mapped and key.startswith('operator '):
                    otype = 'operators'
                    if otype not in objdict:
                        objdict[otype] = {}
                    objdict[otype].update({key: inobj[key]})
                    mapped = True
                elif key == 'oldname':
                    inobj.pop(key)
                    mapped = True
                if not mapped and key != 'schema':
                    raise KeyError("Expected typed object, found '%s'" % key)

            for objtype in SCHOBJS1:
                if objtype in objdict:
                    subobjs = getattr(newdb, objtype)
                    subobjs.from_map(schema, objdict[objtype], newdb)
            for objtype in SCHOBJS2:
                if objtype in objdict:
                    subobjs = getattr(newdb, objtype)
                    subobjs.from_map(schema, objdict[objtype])

    def link_refs(self, db, datacopy):
        """Connect various schema objects to their respective schemas

        :param db: dictionary of dictionaries of all objects
        :param datacopy: dictionary of data copying info
        """
        def link_one(targdict, objtype, objkeys, subtype=None):
            schema = self[objkeys[0]]
            if subtype is None:
                subtype = objtype
            if not hasattr(schema, subtype):
                setattr(schema, subtype, {})
            objdict = getattr(schema, subtype)
            key = objkeys[1] if len(objkeys) == 2 else objkeys[1:]
            objdict.update({key: targdict[objkeys]})

        targ = db.types
        for keys in targ:
            dbtype = targ[keys]
            if isinstance(dbtype, Domain):
                link_one(targ, 'types', keys, 'domains')
            elif isinstance(dbtype, Enum) or isinstance(dbtype, Composite) \
                 or isinstance(dbtype, BaseType) or isinstance(dbtype, Range):
                link_one(targ, 'types', keys)
        targ = db.tables
        for keys in targ:
            table = targ[keys]
            type_ = 'tables'
            if isinstance(table, Table):
                link_one(targ, type_, keys)
            elif isinstance(table, Sequence):
                link_one(targ, type_, keys, 'sequences')
            elif isinstance(table, MaterializedView):
                link_one(targ, type_, keys, 'matviews')
            elif isinstance(table, View):
                link_one(targ, type_, keys, 'views')
        targ = db.functions
        for keys in targ:
            link_one(targ, 'functions', keys)
        for objtype in ['operators', 'operclasses', 'operfams', 'conversions',
                        'tsconfigs', 'tsdicts', 'tsparsers', 'tstempls',
                        'ftables', 'collations']:
            targ = getattr(db, objtype)
            for keys in targ:
                link_one(targ, objtype, keys)
        for key in datacopy:
            if not key.startswith('schema '):
                raise KeyError("Unrecognized object type: %s" % key)
            schema = self[key[7:]]
            if not hasattr(schema, 'datacopy'):
                schema.datacopy = []
            for tbl in datacopy[key]:
                if hasattr(schema, 'tables') and tbl in schema.tables:
                    schema.datacopy.append(tbl)

    def to_map(self, db, opts):
        """Convert the schema dictionary to a regular dictionary

        :param opts: options to include/exclude schemas/tables, etc.
        :return: dictionary

        Invokes the `to_map` method of each schema to construct a
        dictionary of schemas.
        """
        schemas = {}
        selschs = getattr(opts, 'schemas', [])
        for sch in self:
            if not selschs or sch in selschs:
                if hasattr(opts, 'excl_schemas') and opts.excl_schemas \
                        and sch in opts.excl_schemas:
                    continue
                schemas.update(self[sch].to_map(db, self, opts))

        return schemas

    def data_import(self, opts):
        """Iterate over schemas with tables to be imported

        :param opts: options to include/exclude schemas/tables, etc.
        :return: list of SQL statements
        """
        return [self[sch].data_import(opts) for sch in self]
