# -*- coding: utf-8 -*-
#
# Copyright (C) 2018 Shinichi Takii, shinichi.takii@gmail.com
#
# This module is part of python-ddlparse and is released under
# the BSD License: https://opensource.org/licenses/BSD-3-Clause

"""Parse DDL statements"""

import json
import re
import textwrap
from collections import OrderedDict
from enum import IntEnum

from pyparsing import (
    CaselessKeyword,
    Forward,
    Group,
    OneOrMore,
    Optional,
    Regex,
    Suppress,
    Word,
    alphanums,
    delimitedList,
)


class DdlParseBase:

    NAME_CASE = IntEnum("NAME_CASE", "original lower upper")
    DATABASE = IntEnum("DATABASE", "mysql, postgresql, oracle, redshift")

    def __init__(self, source_database=None):
        self._source_database = source_database

    @property
    def source_database(self):
        """
        Source database option

        :param source_database: enum DdlParse.DATABASE
        """
        return self._source_database

    @source_database.setter
    def source_database(self, source_database):
        self._source_database = source_database


class DdlParseTableColumnBase(DdlParseBase):
    def __init__(self, source_database=None):
        super().__init__(source_database)
        self._name = ""

    @property
    def name(self):
        """name"""
        return self._name

    @name.setter
    def name(self, name):
        self._name = name

    def get_name(self, name_case=DdlParseBase.NAME_CASE.original):
        """
        Get Name converted case

        :param name_case: name case type
            * DdlParse.NAME_CASE.original : Return to no convert
            * DdlParse.NAME_CASE.lower : Return to lower
            * DdlParse.NAME_CASE.upper : Return to upper

        :return: name
        """
        if name_case == self.NAME_CASE.lower:
            return self._name.lower()
        elif name_case == self.NAME_CASE.upper:
            return self._name.upper()
        else:
            return self._name


class DdlParseColumn(DdlParseTableColumnBase):
    """Column define info"""

    def __init__(
        self,
        name,
        data_type_array,
        array_brackets=None,
        constraint=None,
        source_database=None,
    ):
        """
        :param data_type_array[]: Column data type ['data type name'] 
        or ['data type name', '(length)'] or ['data type name', '(precision, scale)']
        :param array_brackets: Column array brackets string '[]' or '[][]...'
        :param constraint: Column constraint string
        :param source_database: enum DdlParse.DATABASE
        """

        super().__init__(source_database)
        self._name = name
        self._set_data_type(data_type_array)
        self.constraint = constraint
        self._array_dimensional = (
            0 if array_brackets is None else array_brackets.count("[]")
        )

    @property
    def data_type(self):
        return self._data_type

    @property
    def length(self):
        return self._length

    @property
    def precision(self):
        return self._length

    @property
    def scale(self):
        return self._scale

    def _set_data_type(self, data_type_array):
        self._data_type = data_type_array[0].upper()
        self._length = None
        self._scale = None

        if len(data_type_array) < 2:
            return

        matches = re.findall(r"(\d+)\s*,*\s*(\d*)", data_type_array[-1])
        if len(matches) > 0:
            self._length = int(matches[0][0])
            self._scale = (
                None
                if len(matches[0]) < 2 or matches[0][1] == "" or int(matches[0][1]) == 0
                else int(matches[0][1])
            )

        if re.search(r"^\D+", data_type_array[1]):
            self._data_type += " {}".format(data_type_array[1])

    @property
    def constraint(self):
        """Constraint string"""
        constraint_arr = []
        if self._not_null:
            constraint_arr.append("PRIMARY KEY" if self._pk else "NOT NULL")
        if self._unique:
            constraint_arr.append("UNIQUE")

        return " ".join(constraint_arr)

    @constraint.setter
    def constraint(self, constraint):
        self._constraint = None if constraint is None else constraint.upper()

        self._not_null = (
            False
            if self._constraint is None
            or not re.search(r"(NOT NULL|PRIMARY KEY)", self._constraint)
            else True
        )
        self._pk = (
            False
            if self._constraint is None
            or not re.search("PRIMARY KEY", self._constraint)
            else True
        )
        self._unique = (
            False
            if self._constraint is None or not re.search("UNIQUE", self._constraint)
            else True
        )

        self._comment = None
        if constraint is not None:
            matches = re.findall(
                r"(?:\bCOMMENT\b\s+)(['\"])(.+)\1", constraint, re.IGNORECASE
            )
            if len(matches) > 0:
                self._comment = matches[0][1]

    @property
    def comment(self):
        return self._comment

    @comment.setter
    def comment(self, comment):
        self._comment = comment

    @property
    def description(self):
        return self.comment

    @description.setter
    def description(self, description):
        self.comment = description

    @property
    def array_dimensional(self):
        """array dimensional number"""
        return self._array_dimensional

    @property
    def not_null(self):
        return self._not_null

    @not_null.setter
    def not_null(self, flag):
        self._not_null = flag

    @property
    def primary_key(self):
        return self._pk

    @primary_key.setter
    def primary_key(self, flag):
        self._pk = flag

    @property
    def unique(self):
        return self._unique

    @unique.setter
    def unique(self, flag):
        self._unique = flag

    @property
    def snowflake_data_type(self):
        """Get SNOWFLAKE SQL data type"""

        # Snowflake data type = {source_database: [data type, ...], ...}
        SNOWFLAKE_DATA_TYPE_DIC = OrderedDict()
        SNOWFLAKE_DATA_TYPE_DIC["VARCHAR"] = {
            None: [re.compile(r"(CHAR|TEXT|CLOB|UUID)")]
        }
        SNOWFLAKE_DATA_TYPE_DIC["NUMBER"] = {None: [re.compile(r"INT|SERIAL|YEAR")]}
        SNOWFLAKE_DATA_TYPE_DIC["FLOAT"] = {
            None: [re.compile(r"(FLOAT|DOUBLE)"), "REAL", "MONEY"]
        }
        SNOWFLAKE_DATA_TYPE_DIC["TIMESTAMP_NTZ"] = {
            None: ["DATETIME", "TIMESTAMP", "TIMESTAMP WITHOUT TIME ZONE"],
            self.DATABASE.oracle: ["DATE"],
        }
        SNOWFLAKE_DATA_TYPE_DIC["TIMESTAMP_TZ"] = {
            None: ["TIMESTAMPTZ", "TIMESTAMP WITH TIME ZONE"]
        }
        SNOWFLAKE_DATA_TYPE_DIC["DATE"] = {None: ["DATE"]}
        SNOWFLAKE_DATA_TYPE_DIC["TIME"] = {None: ["TIME"]}
        SNOWFLAKE_DATA_TYPE_DIC["BOOLEAN"] = {None: [re.compile(r"BOOL")]}
        SNOWFLAKE_DATA_TYPE_DIC["VARIANT"] = {None: [re.compile(r"JSONB|JSON")]}

        for bq_type, conditions in SNOWFLAKE_DATA_TYPE_DIC.items():
            for source_db, source_datatypes in conditions.items():
                for source_datatype in source_datatypes:

                    if isinstance(source_datatype, str):
                        if self._data_type == source_datatype and (
                            self._source_database == source_db
                            or (self._source_database is not None and source_db is None)
                        ):
                            return bq_type

                    elif re.search(source_datatype, self._data_type) and (
                        self._source_database == source_db
                        or (self._source_database is not None and source_db is None)
                    ):
                        return bq_type

        if self._data_type in ["NUMERIC", "NUMBER", "DECIMAL"]:
            if self.scale is not None and (self._data_type in ["NUMERIC", "DECIMAL"]):
                return "NUMERIC(%d, %d)" % (self._length, self._scale)

            if self._scale is not None:
                print(self._scale)
                return "FLOAT"

            if (
                self._data_type == "NUMBER"
                and self._source_database == self.DATABASE.oracle
                and self._length is None
            ):
                return "FLOAT"

            return "INTEGER"

        raise ValueError("Unknown data type : '{}'".format(self._data_type))

    @property
    def snowflake_standard_data_type(self):
        """Get SNOWFLAKE Standard SQL data type"""
        data_type = self.snowflake_data_type
        return data_type

    def to_snowflake_field(self, name_case=DdlParseBase.NAME_CASE.original):
        """Generate SNOWFLAKE JSON field define"""

        col_name = self.get_name(name_case)

        # no or one dimensional array data type
        type = self.snowflake_standard_data_type

        col = OrderedDict()
        col["name"] = col_name
        col["type"] = type
        if self.description is not None:
            col["description"] = self.description

        return json.dumps(col, ensure_ascii=False)


class DdlParseColumnDict(OrderedDict, DdlParseBase):
    """
    Columns dictionary collection

    * Orderd dictionary
    * Dict with case insensitive keys
      (SQL is case insensitive)
    """

    def __init__(self, source_database=None):
        super().__init__()
        self.source_database = source_database

    def __getitem__(self, key):
        return super().__getitem__(key.lower())

    def __setitem__(self, key, value):
        super().__setitem__(key.lower(), value)

    def append(
        self,
        column_name,
        data_type_array=None,
        array_brackets=None,
        constraint=None,
        source_database=None,
    ):
        if source_database is None:
            source_database = self.source_database

        column = DdlParseColumn(
            column_name, data_type_array, array_brackets, constraint, source_database
        )
        self.__setitem__(column_name, column)
        return column

    def to_snowflake_fields(self, name_case=DdlParseBase.NAME_CASE.original):
        """
        Generate SnowFlake JSON fields define

        :param name_case: name case type
            * DdlParse.NAME_CASE.original : Return to no convert
            * DdlParse.NAME_CASE.lower : Return to lower
            * DdlParse.NAME_CASE.upper : Return to upper

        :return: SnowFlake JSON fields define
        """

        sf_fields = []

        for col in self.values():
            sf_fields.append(col.to_snowflake_field(name_case))

        return "[{}]".format(",".join(sf_fields))


class DdlParseTable(DdlParseTableColumnBase):
    """Table define info"""

    def __init__(self, source_database=None):
        super().__init__(source_database)
        self._schema = None
        self._columns = DdlParseColumnDict(source_database)

    @property
    def source_database(self):
        """
        Source database option

        :param source_database: enum DdlParse.DATABASE
        """
        return super().source_database

    @source_database.setter
    def source_database(self, source_database):
        super(self.__class__, self.__class__).source_database.__set__(
            self, source_database
        )
        self._columns.source_database = source_database

    @property
    def is_temp(self):
        """Temporary Table Flag"""
        return self._is_temp

    @is_temp.setter
    def is_temp(self, flag):
        self._is_temp = flag

    @property
    def schema(self):
        """Schema name"""
        return self._schema

    @schema.setter
    def schema(self, schema):
        self._schema = schema

    @property
    def columns(self):
        """DdlParseColumn dictionary collection"""
        return self._columns

    def to_snowflake_fields(self, name_case=DdlParseBase.NAME_CASE.original):
        """
        Generate SnowFlake JSON fields define

        :param name_case: name case type
            * DdlParse.NAME_CASE.original : Return to no convert
            * DdlParse.NAME_CASE.lower : Return to lower
            * DdlParse.NAME_CASE.upper : Return to upper

        :return: Snowflake JSON fields define
        """

        return self._columns.to_snowflake_fields(name_case)

    def to_snowflake_ddl(self, name_case=DdlParseBase.NAME_CASE.original):
        """
        Generate Snowflake CREATE TABLE statements

        :param name_case: name case type
            * DdlParse.NAME_CASE.original : Return to no convert
            * DdlParse.NAME_CASE.lower : Return to lower
            * DdlParse.NAME_CASE.upper : Return to upper

        :return: SNOWFLAKE CREATE TABLE statements
        """

        if self.schema is None:
            dataset = "dataset"
        elif name_case == self.NAME_CASE.lower:
            dataset = self.schema.lower()
        elif name_case == self.NAME_CASE.upper:
            dataset = self.schema.upper()
        else:
            dataset = self.schema

        cols_defs = []
        for col in self.columns.values():
            col_name = col.get_name(name_case)

            # no array data type
            type = col.snowflake_standard_data_type
            not_null = " NOT NULL" if col.not_null else ""

            # cols_defs.append("{name} {type}{not_null}".format(
            cols_defs.append(
                "{name} {type}{not_null}{description}".format(
                    name=col_name,
                    type=type,
                    not_null=not_null,
                    description=' OPTIONS (description = "{}")'.format(
                        col.description.replace('"', '\\"')
                    )
                    if col.description is not None
                    else "",
                )
            )

        return textwrap.dedent(
            """\
            #standardSQL
            CREATE TABLE {dataset}.{table}
            (
              {colmns_define}
            )"""
        ).format(
            dataset=dataset,
            table=self.get_name(name_case),
            colmns_define=",\n  ".join(cols_defs),
        )


class DdlParse(DdlParseBase):
    """DDL parser"""

    _LPAR, _RPAR, _COMMA, _SEMICOLON, _DOT, _DOUBLEQUOTE, _BACKQUOTE, _SPACE = map(
        Suppress, '(),;."` '
    )
    _CREATE, _TABLE, _TEMP, _CONSTRAINT, _NOT_NULL, _PRIMARY_KEY, _UNIQUE, _UNIQUE_KEY, _FOREIGN_KEY, _REFERENCES, _KEY, _CHAR_SEMANTICS, _BYTE_SEMANTICS = map(
        CaselessKeyword,
        "CREATE, TABLE, TEMP, CONSTRAINT, NOT NULL, PRIMARY KEY, UNIQUE, UNIQUE KEY, FOREIGN KEY, REFERENCES, KEY, CHAR, BYTE".replace(
            ", ", ","
        ).split(
            ","
        ),
    )
    _SUPPRESS_QUOTE = _BACKQUOTE | _DOUBLEQUOTE

    _COMMENT = Suppress("--" + Regex(r".+"))

    _CREATE_TABLE_STATEMENT = (
        Suppress(_CREATE)
        + Optional(_TEMP)("temp")
        + Suppress(_TABLE)
        + Optional(Suppress(CaselessKeyword("IF NOT EXISTS")))
        + Optional(_SUPPRESS_QUOTE)
        + Optional(
            Word(alphanums + "_")("schema")
            + Optional(_SUPPRESS_QUOTE)
            + _DOT
            + Optional(_SUPPRESS_QUOTE)
        )
        + Word(alphanums + "_<>")("table")
        + Optional(_SUPPRESS_QUOTE)
        + _LPAR
        + delimitedList(
            OneOrMore(
                _COMMENT
                |
                # Ignore Index
                Suppress(_KEY + Word(alphanums + "_'`() "))
                | Group(
                    Optional(
                        Suppress(_CONSTRAINT)
                        + Optional(_SUPPRESS_QUOTE)
                        + Word(alphanums + "_")("name")
                        + Optional(_SUPPRESS_QUOTE)
                    )
                    + (
                        (
                            (_PRIMARY_KEY ^ _UNIQUE ^ _UNIQUE_KEY ^ _NOT_NULL)("type")
                            + Optional(_SUPPRESS_QUOTE)
                            + Optional(Word(alphanums + "_"))("name")
                            + Optional(_SUPPRESS_QUOTE)
                            + _LPAR
                            + Group(
                                delimitedList(
                                    Optional(_SUPPRESS_QUOTE)
                                    + Word(alphanums + "_")
                                    + Optional(_SUPPRESS_QUOTE)
                                )
                            )("constraint_columns")
                            + _RPAR
                        )
                        | (
                            (_FOREIGN_KEY)("type")
                            + _LPAR
                            + Group(
                                delimitedList(
                                    Optional(_SUPPRESS_QUOTE)
                                    + Word(alphanums + "_")
                                    + Optional(_SUPPRESS_QUOTE)
                                )
                            )("constraint_columns")
                            + _RPAR
                            + Optional(
                                Suppress(_REFERENCES)
                                + Optional(_SUPPRESS_QUOTE)
                                + Word(alphanums + "_")("references_table")
                                + Optional(_SUPPRESS_QUOTE)
                                + _LPAR
                                + Group(
                                    delimitedList(
                                        Optional(_SUPPRESS_QUOTE)
                                        + Word(alphanums + "_")
                                        + Optional(_SUPPRESS_QUOTE)
                                    )
                                )("references_columns")
                                + _RPAR
                            )
                        )
                    )
                )("constraint")
                | Group(
                    Optional(_SUPPRESS_QUOTE)
                    + Word(alphanums + "_")("name")
                    + Optional(_SUPPRESS_QUOTE)
                    + Group(
                        Word(alphanums + "_")
                        + Optional(
                            CaselessKeyword("WITHOUT TIME ZONE")
                            ^ CaselessKeyword("WITH TIME ZONE")
                            ^ CaselessKeyword("PRECISION")
                            ^ CaselessKeyword("VARYING")
                        )
                        + Optional(
                            _LPAR
                            + Regex(r"\d+\s*,*\s*\d*")
                            + Optional(Suppress(_CHAR_SEMANTICS | _BYTE_SEMANTICS))
                            + _RPAR
                        )
                    )("type")
                    + Optional(Word(r"\[\]"))("array_brackets")
                    + Optional(
                        Regex(
                            r"(?!--)(\b(COMMENT|DEFAULT)\b\s+[^,]+|([A-Za-z0-9_'\": -]|[^\x01-\x7E])*)",
                            re.IGNORECASE,
                        )
                    )("constraint")
                )("column")
                | _COMMENT
            )
        )("columns")
    )

    _DDL_PARSE_EXPR = Forward()
    _DDL_PARSE_EXPR << OneOrMore(_COMMENT | _CREATE_TABLE_STATEMENT)

    def __init__(self, ddl=None, source_database=None):
        super().__init__(source_database)
        self._ddl = ddl
        self._table = DdlParseTable(source_database)

    @property
    def source_database(self):
        """
        Source database option

        :param source_database: enum DdlParse.DATABASE
        """
        return super().source_database

    @source_database.setter
    def source_database(self, source_database):
        super(self.__class__, self.__class__).source_database.__set__(
            self, source_database
        )
        self._table.source_database = source_database

    @property
    def ddl(self):
        """DDL script"""
        return self._ddl

    @ddl.setter
    def ddl(self, ddl):
        self._ddl = ddl

    def parse(self, ddl=None, source_database=None):
        """
        Parse DDL script.

        :param ddl: DDL script
        :return: DdlParseTable, Parsed table define info.
        """

        if ddl is not None:
            self._ddl = ddl

        if source_database is not None:
            self.source_database = source_database

        if self._ddl is None:
            raise ValueError("DDL is not specified")

        ret = self._DDL_PARSE_EXPR.parseString(self._ddl)
        # print(ret.dump())

        if "schema" in ret:
            self._table.schema = ret["schema"]

        self._table.name = ret["table"]
        self._table.is_temp = True if "temp" in ret else False

        for ret_col in ret["columns"]:

            if ret_col.getName() == "column":
                # add column
                col = self._table.columns.append(
                    column_name=ret_col["name"],
                    data_type_array=ret_col["type"],
                    array_brackets=ret_col["array_brackets"]
                    if "array_brackets" in ret_col
                    else None,
                    constraint=ret_col["constraint"]
                    if "constraint" in ret_col
                    else None,
                )

            elif ret_col.getName() == "constraint":
                # set column constraint
                for col_name in ret_col["constraint_columns"]:
                    col = self._table.columns[col_name]

                    if ret_col["type"] == "PRIMARY KEY":
                        col.not_null = True
                        col.primary_key = True
                    elif ret_col["type"] in ["UNIQUE", "UNIQUE KEY"]:
                        col.unique = True
                    elif ret_col["type"] == "NOT NULL":
                        col.not_null = True

        return self._table
