"""MariaDB Dialect.

https://mariadb.com/kb/en/sql-statements-structure/
"""

from sqlfluff.core.dialects import load_raw_dialect
from sqlfluff.core.parser import (
    AnyNumberOf,
    AnySetOf,
    BaseSegment,
    Bracketed,
    Dedent,
    Delimited,
    Indent,
    Matchable,
    OneOf,
    OptionallyDelimited,
    ParseMode,
    Ref,
    Sequence,
)
from sqlfluff.dialects import dialect_mysql as mysql
from sqlfluff.dialects.dialect_mariadb_keywords import (
    mariadb_reserved_keywords,
    mariadb_unreserved_keywords,
)

# ansi_dialect = load_raw_dialect("ansi")
mysql_dialect = load_raw_dialect("mysql")
mariadb_dialect = mysql_dialect.copy_as(
    "mariadb",
    formatted_name="MariaDB",
    docstring="""**Default Casing**: ``lowercase``

**Quotes**: String Literals: ``''``, ``""`` or ``@``,
Identifiers: |back_quotes|.

The dialect for `MariaDB <https://www.mariadb.org/>`_.""",
)
mariadb_dialect.update_keywords_set_from_multiline_string(
    "unreserved_keywords", mariadb_unreserved_keywords
)
mariadb_dialect.sets("reserved_keywords").clear()
mariadb_dialect.update_keywords_set_from_multiline_string(
    "reserved_keywords", mariadb_reserved_keywords
)
mariadb_dialect.replace(
    AddDropSystemVersioningGrammar=Sequence(
        OneOf("ADD", "DROP"),
        "SYSTEM",
        "VERSIONING",
    ),
    TriggerOrReplaceGrammar=Sequence("OR", "REPLACE"),
)


class ColumnConstraintSegment(mysql.ColumnConstraintSegment):
    """A column option; each CREATE TABLE column can have 0 or more."""

    match_grammar: Matchable = OneOf(
        mysql.ColumnConstraintSegment.match_grammar,
        Sequence(
            Sequence("GENERATED", "ALWAYS", optional=True),
            "AS",
            Bracketed(Ref("ExpressionSegment")),
            OneOf("PERSISTENT", "STORED", "VIRTUAL", optional=True),
        ),
    )


class CreateUserStatementSegment(mysql.CreateUserStatementSegment):
    """`CREATE USER` statement.

    https://mariadb.com/kb/en/create-user/
    """

    match_grammar = mysql.CreateUserStatementSegment.match_grammar.copy(
        insert=[Ref("OrReplaceGrammar", optional=True)],
        before=Ref.keyword("USER"),
    )


class DeleteStatementSegment(BaseSegment):
    """A `DELETE` statement.

    https://mariadb.com/kb/en/delete/
    """

    type = "delete_statement"
    match_grammar = Sequence(
        "DELETE",
        Ref.keyword("LOW_PRIORITY", optional=True),
        Ref.keyword("QUICK", optional=True),
        Ref.keyword("IGNORE", optional=True),
        OneOf(
            Sequence(
                "FROM",
                Delimited(
                    Ref("DeleteTargetTableSegment"),
                    terminators=["USING"],
                ),
                Ref("DeleteUsingClauseSegment"),
                Ref("WhereClauseSegment", optional=True),
            ),
            Sequence(
                Delimited(
                    Ref("DeleteTargetTableSegment"),
                    terminators=["FROM"],
                ),
                Ref("FromClauseSegment"),
                Ref("WhereClauseSegment", optional=True),
            ),
            Sequence(
                Ref("FromClauseSegment"),
                Ref("SelectPartitionClauseSegment", optional=True),
                Ref("WhereClauseSegment", optional=True),
                Ref("OrderByClauseSegment", optional=True),
                Ref("LimitClauseSegment", optional=True),
                Ref("ReturningClauseSegment", optional=True),
            ),
        ),
    )


class FlushStatementSegment(mysql.FlushStatementSegment):
    """A `Flush` statement.

    https://mariadb.com/kb/en/flush/
    """

    match_grammar: Matchable = Sequence(
        "FLUSH",
        OneOf(
            "NO_WRITE_TO_BINLOG",
            "LOCAL",
            optional=True,
        ),
        OneOf(
            Delimited(
                Sequence("BINARY", "LOGS"),
                Sequence("ENGINE", "LOGS"),
                Sequence("ERROR", "LOGS"),
                Sequence("GENERAL", "LOGS"),
                Sequence("QUERY", "CACHE"),
                Sequence("SLOW", "LOGS"),
                Sequence(Ref.keyword("RESET", optional=True), "MASTER"),
                Sequence(OneOf("GLOBAL", "SESSION", optional=True), "STATUS"),
                Sequence(
                    "RELAY",
                    "LOGS",
                    Sequence("FOR", "CHANNEL", optional=True),
                    Ref("ObjectReferenceSegment"),
                ),
                "HOSTS",
                "LOGS",
                "PRIVILEGES",
                "CHANGED_PAGE_BITMAPS",
                "CLIENT_STATISTICS",
                "DES_KEY_FILE",
                "INDEX_STATISTICS",
                "QUERY_RESPONSE_TIME",
                "SLAVE",
                "SSL",
                "TABLE_STATISTICS",
                "USER_STATISTICS",
                "USER_VARIABLES",
                "USER_RESOURCES",
            ),
            Sequence(
                "TABLES",
                Sequence(
                    Delimited(Ref("TableReferenceSegment"), terminators=["WITH"]),
                    optional=True,
                ),
                Sequence(
                    "WITH",
                    "READ",
                    "LOCK",
                    Sequence("AND", "DISABLE", "CHECKPOINT", optional=True),
                    optional=True,
                ),
            ),
            Sequence(
                "TABLES",
                Sequence(
                    Delimited(Ref("TableReferenceSegment"), terminators=["FOR"]),
                    optional=False,
                ),
                Sequence("FOR", "EXPORT", optional=True),
            ),
        ),
    )


class GroupByClauseSegment(BaseSegment):
    """A `GROUP BY` clause like in `SELECT`."""

    type = "groupby_clause"

    match_grammar: Matchable = Sequence(
        "GROUP",
        "BY",
        Indent,
        Sequence(
            Delimited(
                Sequence(
                    OneOf(
                        Ref("ColumnReferenceSegment"),
                        # Can `GROUP BY 1`
                        Ref("NumericLiteralSegment"),
                        # Can `GROUP BY coalesce(col, 1)`
                        Ref("ExpressionSegment"),
                    ),
                    OneOf("ASC", "DESC", optional=True),
                ),
                terminators=[Ref("GroupByClauseTerminatorGrammar")],
            ),
        ),
        Ref("WithRollupClauseSegment", optional=True),
        Dedent,
    )


class InsertStatementSegment(BaseSegment):
    """An `INSERT` statement.

    https://mariadb.com/kb/en/insert/
    """

    type = "insert_statement"
    match_grammar = Sequence(
        "INSERT",
        OneOf(
            "LOW_PRIORITY",
            "DELAYED",
            "HIGH_PRIORITY",
            optional=True,
        ),
        Ref.keyword("IGNORE", optional=True),
        Ref.keyword("INTO", optional=True),
        Ref("TableReferenceSegment"),
        Sequence(
            "PARTITION",
            Bracketed(
                Ref("SingleIdentifierListSegment"),
            ),
            optional=True,
        ),
        Ref("BracketedColumnReferenceListGrammar", optional=True),
        OneOf(
            Ref("ValuesClauseSegment"),
            Ref("SetClauseListSegment"),
            Ref(
                "SelectableGrammar",
                terminators=[Ref("ReturningClauseSegment")],
            ),
            optional=True,
        ),
        Ref("InsertRowAliasSegment", optional=True),
        Ref("UpsertClauseListSegment", optional=True),
        Ref("ReturningClauseSegment", optional=True),
    )


class ReplaceSegment(BaseSegment):
    """A `REPLACE` statement.

    As per https://mariadb.com/kb/en/replace/
    """

    type = "replace_statement"

    match_grammar = Sequence(
        "REPLACE",
        OneOf("LOW_PRIORITY", "DELAYED", optional=True),
        Sequence("INTO", optional=True),
        Ref("TableReferenceSegment"),
        Ref("SelectPartitionClauseSegment", optional=True),
        OneOf(
            Sequence(
                Ref("BracketedColumnReferenceListGrammar", optional=True),
                Ref("ValuesClauseSegment"),
            ),
            Ref("SetClauseListSegment"),
            Sequence(
                Ref("BracketedColumnReferenceListGrammar", optional=True),
                Ref("SelectStatementSegment"),
            ),
        ),
        Ref("ReturningClauseSegment", optional=True),
    )


class ReturningClauseSegment(BaseSegment):
    """This is a `RETURNING` clause.

    A RETURNING clause returns values modified by a
    INSERT, DELETE or REPLACE query.

    https://mariadb.com/kb/en/insert/
    https://mariadb.com/kb/en/delete/
    https://mariadb.com/kb/en/replace/
    """

    type = "returning_clause"

    match_grammar: Matchable = Sequence(
        "RETURNING",
        Indent,
        Delimited(
            Ref("SelectClauseElementSegment"),
            allow_trailing=True,
        ),
        Dedent,
        terminators=[Ref("SelectClauseTerminatorGrammar")],
        parse_mode=ParseMode.GREEDY_ONCE_STARTED,
    )


class SelectStatementSegment(mysql.SelectStatementSegment):
    """A `SELECT` statement.

    https://mariadb.com/kb/en/select/
    """

    # Inherit most of the parse grammar from the original.
    match_grammar = mysql.UnorderedSelectStatementSegment.match_grammar.copy(
        insert=[
            Ref("OrderByClauseSegment", optional=True),
            Ref("LimitClauseSegment", optional=True),
            Ref("NamedWindowSegment", optional=True),
            Ref("IntoClauseSegment", optional=True),
        ],
        terminators=[
            Ref("SetOperatorSegment"),
            Ref("UpsertClauseListSegment"),
            Ref("WithCheckOptionSegment"),
            Ref("ReturningClauseSegment"),
        ],
        # Overwrite the terminators, because we want to remove some from the
        # expression above.
        replace_terminators=True,
    )


class WithRollupClauseSegment(BaseSegment):
    """A `WITH ROLLUP` clause after the `GROUP BY` clause."""

    type = "with_rollup_clause"

    match_grammar = Sequence(
        "WITH",
        "ROLLUP",
    )


class IndexTypeSegment(BaseSegment):
    """Index type specification: USING {BTREE | HASH | RTREE}."""

    type = "index_type"

    match_grammar = Sequence(
        "USING",
        OneOf("BTREE", "HASH", "RTREE"),
    )


class IndexOptionSegment(BaseSegment):
    """Index option specification for MariaDB."""

    type = "index_option"

    match_grammar = AnyNumberOf(
        # KEY_BLOCK_SIZE [=] value
        Sequence(
            "KEY_BLOCK_SIZE",
            Ref("EqualsSegment", optional=True),
            Ref("NumericLiteralSegment"),
        ),
        # index_type
        Ref("IndexTypeSegment"),
        # WITH PARSER parser_name
        Sequence(
            "WITH",
            "PARSER",
            Ref("ObjectReferenceSegment"),
        ),
        # COMMENT 'string'
        Sequence(
            "COMMENT",
            Ref("QuotedLiteralSegment"),
        ),
        # CLUSTERING={YES| NO}
        Sequence(
            "CLUSTERING",
            Ref("EqualsSegment", optional=True),
            OneOf("YES", "NO"),
        ),
        # IGNORED | NOT IGNORED
        OneOf(
            "IGNORED",
            Sequence("NOT", "IGNORED"),
        ),
    )


class AlgorithmOptionSegment(BaseSegment):
    """Algorithm option: ALGORITHM [=] {DEFAULT|INPLACE|COPY|NOCOPY|INSTANT}."""

    type = "algorithm_option"

    match_grammar = Sequence(
        "ALGORITHM",
        Ref("EqualsSegment", optional=True),
        OneOf("DEFAULT", "INPLACE", "COPY", "NOCOPY", "INSTANT"),
    )


class LockOptionSegment(BaseSegment):
    """Lock option: LOCK [=] {DEFAULT|NONE|SHARED|EXCLUSIVE}."""

    type = "lock_option"

    match_grammar = Sequence(
        "LOCK",
        Ref("EqualsSegment", optional=True),
        OneOf("DEFAULT", "NONE", "SHARED", "EXCLUSIVE"),
    )


class WaitOptionSegment(BaseSegment):
    """Wait option: WAIT n | NOWAIT."""

    type = "wait_option"

    match_grammar = OneOf(
        Sequence("WAIT", Ref("NumericLiteralSegment")),
        "NOWAIT",
    )


class CreateIndexStatementSegment(mysql.CreateIndexStatementSegment):
    """A `CREATE INDEX` statement for MariaDB."""

    type = "create_index_statement"

    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        OneOf("UNIQUE", "FULLTEXT", "SPATIAL", optional=True),
        "INDEX",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("IndexReferenceSegment"),
        Ref("IndexTypeSegment", optional=True),
        "ON",
        Ref("TableReferenceSegment"),
        Ref("BracketedKeyPartListGrammar"),
        Ref("WaitOptionSegment", optional=True),
        Ref("IndexOptionSegment", optional=True),
        AnySetOf(
            Ref("AlgorithmOptionSegment"),
            Ref("LockOptionSegment"),
            optional=True,
        ),
    )


class CreateProcedureStatementSegment(mysql.CreateProcedureStatementSegment):
    """A `CREATE PROCEDURE` statement.

    https://mariadb.com/docs/server/server-usage/stored-routines/stored-procedures/create-procedure
    """

    match_grammar = mysql.CreateProcedureStatementSegment.match_grammar.copy(
        insert=[Ref("OrReplaceGrammar", optional=True)],
        before=Ref("ProcedureKeywordSegment"),
    )


class CreateFunctionStatementSegment(mysql.CreateFunctionStatementSegment):
    """A `CREATE FUNCTION` statement.

    https://mariadb.com/docs/server/reference/sql-statements/data-definition/create/create-function
    """

    match_grammar = mysql.CreateFunctionStatementSegment.match_grammar.copy(
        insert=[Ref("OrReplaceGrammar", optional=True)],
        before=Ref("FunctionKeywordSegment"),
    )


class TableOptionsSegment(mysql.TableOptionsSegment):
    """MariaDB table options for CREATE/ALTER TABLE statements.

    https://mariadb.com/kb/en/create-table/
    """

    type = "table_options"

    match_grammar = OptionallyDelimited(
        OneOf(
            # [STORAGE] ENGINE [=] engine_name
            Sequence(
                Ref.keyword("STORAGE", optional=True),
                "ENGINE",
                Ref("EqualsSegment", optional=True),
                OneOf(Ref("QuotedLiteralSegment"), Ref("NakedIdentifierSegment")),
            ),
            # AUTO_INCREMENT [=] number
            Sequence(
                "AUTO_INCREMENT",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # AVG_ROW_LENGTH [=] number
            Sequence(
                "AVG_ROW_LENGTH",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # [DEFAULT] CHARACTER SET [=] charset_name
            Sequence(
                Ref.keyword("DEFAULT", optional=True),
                Ref("CharsetGrammar"),
                Ref("EqualsSegment", optional=True),
                OneOf(
                    Ref("QuotedLiteralSegment"),
                    Ref("NakedIdentifierSegment"),
                    "DEFAULT",
                ),
            ),
            # CHECKSUM [=] {0 | 1}
            Sequence(
                "CHECKSUM",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # [DEFAULT] COLLATE [=] collation_name
            Sequence(
                Ref.keyword("DEFAULT", optional=True),
                "COLLATE",
                Ref("EqualsSegment", optional=True),
                OneOf(
                    Ref("QuotedLiteralSegment"),
                    Ref("NakedIdentifierSegment"),
                    "DEFAULT",
                ),
            ),
            # COMMENT [=] 'string'
            Sequence(
                "COMMENT",
                Ref("EqualsSegment", optional=True),
                Ref("QuotedLiteralSegment"),
            ),
            # CONNECTION [=] 'connect_string'
            Sequence(
                "CONNECTION",
                Ref("EqualsSegment", optional=True),
                Ref("QuotedLiteralSegment"),
            ),
            # DATA DIRECTORY [=] 'absolute path to directory'
            Sequence(
                "DATA",
                "DIRECTORY",
                Ref("EqualsSegment", optional=True),
                Ref("QuotedLiteralSegment"),
            ),
            # DELAY_KEY_WRITE [=] {0 | 1}
            Sequence(
                "DELAY_KEY_WRITE",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # ENCRYPTED [=] {YES | NO}
            Sequence(
                "ENCRYPTED",
                Ref("EqualsSegment", optional=True),
                OneOf("YES", "NO"),
            ),
            # ENCRYPTION_KEY_ID [=] number
            Sequence(
                "ENCRYPTION_KEY_ID",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # IETF_QUOTES [=] {YES | NO}
            Sequence(
                "IETF_QUOTES",
                Ref("EqualsSegment", optional=True),
                OneOf("YES", "NO"),
            ),
            # INDEX DIRECTORY [=] 'absolute path to directory'
            Sequence(
                "INDEX",
                "DIRECTORY",
                Ref("EqualsSegment", optional=True),
                Ref("QuotedLiteralSegment"),
            ),
            # INSERT_METHOD [=] { NO | FIRST | LAST }
            Sequence(
                "INSERT_METHOD",
                Ref("EqualsSegment", optional=True),
                OneOf("NO", "FIRST", "LAST"),
            ),
            # KEY_BLOCK_SIZE [=] number
            Sequence(
                "KEY_BLOCK_SIZE",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # MAX_ROWS [=] number
            Sequence(
                "MAX_ROWS",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # MIN_ROWS [=] number
            Sequence(
                "MIN_ROWS",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # PACK_KEYS [=] {0 | 1 | DEFAULT}
            Sequence(
                "PACK_KEYS",
                Ref("EqualsSegment", optional=True),
                OneOf(Ref("NumericLiteralSegment"), "DEFAULT"),
            ),
            # PAGE_CHECKSUM [=] {0 | 1}
            Sequence(
                "PAGE_CHECKSUM",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # PAGE_COMPRESSED [=] {0 | 1}
            Sequence(
                "PAGE_COMPRESSED",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # PAGE_COMPRESSION_LEVEL [=] {0 .. 9}
            Sequence(
                "PAGE_COMPRESSION_LEVEL",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # PASSWORD [=] 'string'
            Sequence(
                "PASSWORD",
                Ref("EqualsSegment", optional=True),
                Ref("QuotedLiteralSegment"),
            ),
            # ROW_FORMAT [=] {DEFAULT|DYNAMIC|FIXED|COMPRESSED|REDUNDANT|COMPACT|PAGE}
            Sequence(
                "ROW_FORMAT",
                Ref("EqualsSegment", optional=True),
                OneOf(
                    "DEFAULT",
                    "DYNAMIC",
                    "FIXED",
                    "COMPRESSED",
                    "REDUNDANT",
                    "COMPACT",
                    "PAGE",
                ),
            ),
            # SEQUENCE [=] {0|1}
            Sequence(
                "SEQUENCE",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # STATS_AUTO_RECALC [=] {DEFAULT|0|1}
            Sequence(
                "STATS_AUTO_RECALC",
                Ref("EqualsSegment", optional=True),
                OneOf(
                    "DEFAULT",
                    Ref("NumericLiteralSegment"),
                ),
            ),
            # STATS_PERSISTENT [=] {DEFAULT|0|1}
            Sequence(
                "STATS_PERSISTENT",
                Ref("EqualsSegment", optional=True),
                OneOf(
                    "DEFAULT",
                    Ref("NumericLiteralSegment"),
                ),
            ),
            # STATS_SAMPLE_PAGES [=] {DEFAULT|number}
            Sequence(
                "STATS_SAMPLE_PAGES",
                Ref("EqualsSegment", optional=True),
                OneOf(
                    "DEFAULT",
                    Ref("NumericLiteralSegment"),
                ),
            ),
            # TABLESPACE tablespace_name
            Sequence(
                "TABLESPACE",
                Ref("NakedIdentifierSegment"),
            ),
            # TRANSACTIONAL [=] {0 | 1}
            Sequence(
                "TRANSACTIONAL",
                Ref("EqualsSegment", optional=True),
                Ref("NumericLiteralSegment"),
            ),
            # UNION [=] (tbl_name[,tbl_name]...)
            Sequence(
                "UNION",
                Ref("EqualsSegment", optional=True),
                Bracketed(Delimited(Ref("TableReferenceSegment"))),
            ),
            # WITH SYSTEM VERSIONING
            Sequence(
                "WITH",
                "SYSTEM",
                "VERSIONING",
            ),
        ),
    )
