from typing import Optional, List, Dict, Any

import agate
import dbt.exceptions
import dbt
from dbt.adapters.base.relation import SchemaSearchMap
from dbt.adapters.sql import SQLAdapter
from dbt.node_types import NodeType

from dbt.adapters.spark import SparkConnectionManager
from dbt.adapters.spark import SparkRelation
from dbt.adapters.spark import SparkColumn


from dbt.adapters.base import BaseRelation

from dbt.logger import GLOBAL_LOGGER as logger

GET_COLUMNS_IN_RELATION_MACRO_NAME = 'get_columns_in_relation'
LIST_SCHEMAS_MACRO_NAME = 'list_schemas'
LIST_RELATIONS_MACRO_NAME = 'list_relations_without_caching'
DROP_RELATION_MACRO_NAME = 'drop_relation'

KEY_TABLE_OWNER = 'Owner'
KEY_TABLE_STATISTICS = 'Statistics'


class SparkAdapter(SQLAdapter):
    COLUMN_NAMES = (
        'table_database',
        'table_schema',
        'table_name',
        'table_type',
        'table_comment',
        'table_owner',
        'column_name',
        'column_index',
        'column_type',
        'column_comment',

        'stats:bytes:label',
        'stats:bytes:value',
        'stats:bytes:description',
        'stats:bytes:include',

        'stats:rows:label',
        'stats:rows:value',
        'stats:rows:description',
        'stats:rows:include',
    )

    Relation = SparkRelation
    Column = SparkColumn
    ConnectionManager = SparkConnectionManager

    AdapterSpecificConfigs = frozenset({"file_format", "location_root",
                                        "partition_by", "clustered_by",
                                        "buckets"})

    @classmethod
    def date_function(cls) -> str:
        return 'current_timestamp()'

    @classmethod
    def convert_text_type(cls, agate_table, col_idx):
        return "string"

    @classmethod
    def convert_number_type(cls, agate_table, col_idx):
        decimals = agate_table.aggregate(agate.MaxPrecision(col_idx))
        return "double" if decimals else "bigint"

    @classmethod
    def convert_date_type(cls, agate_table, col_idx):
        return "date"

    @classmethod
    def convert_time_type(cls, agate_table, col_idx):
        return "time"

    @classmethod
    def convert_datetime_type(cls, agate_table, col_idx):
        return "timestamp"

    def quote(self, identifier):
        return '`{}`'.format(identifier)

    def add_schema_to_cache(self, schema) -> str:
        """Cache a new schema in dbt. It will show up in `list relations`."""
        if schema is None:
            name = self.nice_connection_name()
            dbt.exceptions.raise_compiler_error(
                'Attempted to cache a null schema for {}'.format(name)
            )
        if dbt.flags.USE_CACHE:
            self.cache.add_schema(None, schema)
        # so jinja doesn't render things
        return ''

    def list_relations_without_caching(
        self, information_schema, schema
    ) -> List[SparkRelation]:
        kwargs = {'information_schema': information_schema, 'schema': schema}
        try:
            results = self.execute_macro(
                LIST_RELATIONS_MACRO_NAME,
                kwargs=kwargs,
                release=True
            )
        except dbt.exceptions.RuntimeException as e:
            if hasattr(e, 'msg') and f"Database '{schema}' not found" in e.msg:
                return []
            else:
                description = "Error while retrieving information about"
                logger.debug(f"{description} {schema}: {e.msg}")
                return []

        relations = []
        for _schema, name, _, information in results:
            rel_type = ('view' if 'Type: VIEW' in information else 'table')
            relation = self.Relation.create(
                schema=_schema,
                identifier=name,
                type=rel_type
            )
            relations.append(relation)

        return relations

    def get_relation(
        self, database: str, schema: str, identifier: str
    ) -> Optional[BaseRelation]:
        if not self.Relation.include_policy.database:
            database = None

        return super().get_relation(database, schema, identifier)

    def parse_describe_extended(
            self,
            relation: Relation,
            raw_rows: List[agate.Row]
    ) -> List[SparkColumn]:
        # Convert the Row to a dict
        dict_rows = [dict(zip(row._keys, row._values)) for row in raw_rows]
        # Find the separator between the rows and the metadata provided
        # by the DESCRIBE TABLE EXTENDED statement
        pos = self.find_table_information_separator(dict_rows)

        # Remove rows that start with a hash, they are comments
        rows = [
            row for row in raw_rows[0:pos]
            if not row['col_name'].startswith('#')
        ]
        metadata = {
            col['col_name']: col['data_type'] for col in raw_rows[pos + 1:]
        }

        raw_table_stats = metadata.get(KEY_TABLE_STATISTICS)
        table_stats = SparkColumn.convert_table_stats(raw_table_stats)
        return [SparkColumn(
            table_database=relation.database,
            table_schema=relation.schema,
            table_name=relation.name,
            table_type=relation.type,
            table_owner=metadata.get(KEY_TABLE_OWNER),
            table_stats=table_stats,
            column=column['col_name'],
            column_index=idx,
            dtype=column['data_type'],
        ) for idx, column in enumerate(rows)]

    @staticmethod
    def find_table_information_separator(rows: List[dict]) -> int:
        pos = 0
        for row in rows:
            if not row['col_name'] or row['col_name'].startswith('#'):
                break
            pos += 1
        return pos

    def get_columns_in_relation(self, relation: Relation) -> List[SparkColumn]:
        rows: List[agate.Row] = super().get_columns_in_relation(relation)
        return self.parse_describe_extended(relation, rows)

    @staticmethod
    def _parse_relation(relation: Relation,
                        table_columns: List[Column],
                        rel_type: str,
                        properties: Dict[str, str] = None) -> List[dict]:
        properties = properties or {}
        statistics = {}
        table_owner_key = 'Owner'

        # First check if it is present in the properties
        table_owner = properties.get(table_owner_key)

        found_detailed_table_marker = False
        for column in table_columns:
            if column.name == '# Detailed Table Information':
                found_detailed_table_marker = True

            # In case there is another column with the name Owner
            if not found_detailed_table_marker:
                continue

            if not table_owner and column.name == table_owner_key:
                table_owner = column.data_type

            if column.name == 'Statistics':
                # format: 1109049927 bytes, 14093476 rows
                statistics = {stats.split(" ")[1]: int(stats.split(" ")[0]) for
                              stats in column.data_type.split(', ')}

        columns = []
        for column_index, column in enumerate(table_columns):
            # Fixes for pseudo-columns with no type
            if column.name in {
                '# Partition Information',
                '# col_name',
                ''
            }:
                continue
            elif column.name == '# Detailed Table Information':
                # Loop until the detailed table information
                break
            elif column.data_type is None:
                continue

            column_data = (
                relation.database,
                relation.schema,
                relation.name,
                rel_type,
                None,
                table_owner,
                column.name,
                column_index,
                column.data_type,
                None,

                # Table level stats
                'Table size',
                statistics.get("bytes"),
                "The size of the table in bytes",
                statistics.get("bytes") is not None,

                # Column level stats
                'Number of rows',
                statistics.get("rows"),
                "The number of rows in the table",
                statistics.get("rows") is not None
            )

            column_dict = dict(zip(SparkAdapter.COLUMN_NAMES, column_data))
            columns.append(column_dict)

        return columns

    def _massage_column_for_catalog(
        self, column: SparkColumn
    ) -> Dict[str, Any]:
        dct = column.to_dict()
        # different expectations here - Column.column is the name
        dct['column_name'] = dct.pop('column')
        dct['column_type'] = dct.pop('dtype')
        # table_database can't be None in core.
        if dct['table_database'] is None:
            dct['table_database'] = dct['table_schema']
        return dct

    def _get_catalog_for_relations(self, database: str, schema: str):
        with self.connection_named(f'{database}.{schema}'):
            columns = []
            for relation in self.list_relations(database, schema):
                logger.debug("Getting table schema for relation {}", relation)
                columns.extend(
                    self._massage_column_for_catalog(col)
                    for col in self.get_columns_in_relation(relation)
                )
        return agate.Table.from_object(columns)

    def _get_cache_schemas(self, manifest, exec_only=False):
        info_schema_name_map = SchemaSearchMap()
        for node in manifest.nodes.values():
            if exec_only and node.resource_type not in NodeType.executable():
                continue
            relation = self.Relation.create(
                database=node.database,
                schema=node.schema,
                identifier='information_schema',
                quote_policy=self.config.quoting,
            )
            key = relation.information_schema_only()
            info_schema_name_map[key] = {node.schema}
        return info_schema_name_map

    def _get_one_catalog(
        self, information_schema, schemas, manifest,
    ) -> agate.Table:
        name = f'{information_schema.database}.information_schema'

        if len(schemas) != 1:
            dbt.exceptions.raise_compiler_error(
                'Expected only one schema in spark _get_one_catalog'
            )

        database = information_schema.database
        schema = list(schemas)[0]

        with self.connection_named(name):
            columns = []
            for relation in self.list_relations(database, schema):
                logger.debug("Getting table schema for relation {}", relation)
                columns.extend(
                    self._massage_column_for_catalog(col)
                    for col in self.get_columns_in_relation(relation)
                )
            return agate.Table.from_object(columns)

    def check_schema_exists(self, database, schema):
        results = self.execute_macro(
            LIST_SCHEMAS_MACRO_NAME,
            kwargs={'database': database}
        )

        exists = True if schema in [row[0] for row in results] else False
        return exists
