"""The Databricks Dialect.

Functionally, it is quite similar to SparkSQL,
however it's much less strict on keywords.
It also has some extensions.
"""

from sqlfluff.core.dialects import load_raw_dialect
from sqlfluff.core.parser import (
    AnyNumberOf,
    AnySetOf,
    Anything,
    BaseSegment,
    Bracketed,
    CodeSegment,
    CommentSegment,
    Dedent,
    Delimited,
    IdentifierSegment,
    Indent,
    Matchable,
    NewlineSegment,
    OneOf,
    OptionallyBracketed,
    Ref,
    RegexLexer,
    RegexParser,
    Sequence,
    StringLexer,
    StringParser,
    SymbolSegment,
    TypedParser,
    WordSegment,
)
from sqlfluff.dialects import dialect_ansi as ansi
from sqlfluff.dialects import dialect_sparksql as sparksql
from sqlfluff.dialects.dialect_databricks_keywords import (
    RESERVED_KEYWORDS,
    UNRESERVED_KEYWORDS,
)

sparksql_dialect = load_raw_dialect("sparksql")
databricks_dialect = sparksql_dialect.copy_as(
    "databricks",
    formatted_name="Databricks",
    docstring="The dialect for `Databricks <https://databricks.com/>`_.",
)

databricks_dialect.sets("unreserved_keywords").update(UNRESERVED_KEYWORDS)
databricks_dialect.sets("unreserved_keywords").update(
    sparksql_dialect.sets("reserved_keywords")
)
databricks_dialect.sets("unreserved_keywords").difference_update(RESERVED_KEYWORDS)
databricks_dialect.sets("reserved_keywords").clear()
databricks_dialect.sets("reserved_keywords").update(RESERVED_KEYWORDS)

databricks_dialect.sets("date_part_function_name").update(["TIMEDIFF"])


databricks_dialect.insert_lexer_matchers(
    # Named Function Parameters:
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-function-invocation.html#named-parameter-invocation
    [
        StringLexer("right_arrow", "=>", CodeSegment),
    ],
    before="equals",
)


databricks_dialect.insert_lexer_matchers(
    # ?:: is the try_cast shorthand (an error-tolerating cast). It must be matched
    # before the "?" (question) matcher so that "?::" lexes as a single token
    # rather than as "?" followed by "::".
    # https://docs.databricks.com/aws/en/sql/language-manual/functions/questiondoublecolonsign
    [
        StringLexer("try_casting_operator", "?::", CodeSegment),
    ],
    before="question",
)


databricks_dialect.insert_lexer_matchers(
    # Databricks Pipeline Parameters:
    # https://docs.databricks.com/en/delta-live-tables/parameters.html
    # Must come before dollar_quote since both start with $
    [
        RegexLexer(
            "pipeline_parameter",
            r"\$\{[A-Za-z_][A-Za-z0-9_]*\}",
            CodeSegment,
        ),
    ],
    before="dollar_quote",
)


databricks_dialect.insert_lexer_matchers(
    # Notebook Cell Delimiter:
    # https://learn.microsoft.com/en-us/azure/databricks/notebooks/notebook-export-import#sql-1
    [
        RegexLexer("command", r"(\r?\n){2}-- COMMAND ----------(\r?\n)", CodeSegment),
    ],
    before="newline",
)

databricks_dialect.insert_lexer_matchers(
    # Bare notebook magic in the first cell must consume the notebook header,
    # because the lexer only matches against the remaining string slice.
    [
        RegexLexer(
            "notebook_start_bare_magic_sql",
            r"-- Databricks notebook source(?:\r?\n)%sql\b[^\r\n]*",
            CommentSegment,
            subdivider=RegexLexer("newline", r"\r\n|\n", NewlineSegment),
            trim_post_subdivide=RegexLexer(
                "notebook_start", r"-- Databricks notebook source", CommentSegment
            ),
        ),
        RegexLexer(
            "notebook_start_bare_magic_cell",
            r"(?s)-- Databricks notebook source(?:\r?\n)"
            r"%(?:python|scala|r|sh|md|run|fs|pip|conda)\b[^\r\n]*"
            r"(?:(?!(?:\r?\n){2}-- COMMAND ----------(?:\r?\n)).)*"
            r"(?=(?:\r?\n){2}-- COMMAND ----------(?:\r?\n)|\Z)",
            CodeSegment,
            subdivider=RegexLexer("newline", r"\r\n|\n", NewlineSegment),
            trim_post_subdivide=RegexLexer(
                "notebook_start", r"-- Databricks notebook source", CommentSegment
            ),
        ),
    ],
    before="inline_comment",
)


databricks_dialect.insert_lexer_matchers(
    # Databricks Notebook Start:
    # needed to insert "so early" to avoid magic + notebook
    # start to be interpreted as inline comments
    # https://learn.microsoft.com/en-us/azure/databricks/notebooks/notebooks-code#language-magic
    [
        RegexLexer(
            "notebook_start", r"-- Databricks notebook source(\r?\n){1}", CommentSegment
        ),
        RegexLexer(
            "magic_single_line",
            r"(-- MAGIC %)([^\n]{2,})( [^%]{1})([^\n]*)",
            CodeSegment,
        ),
        RegexLexer("magic_line", r"(-- MAGIC)( [^%]{1})([^\n]*)", CodeSegment),
        RegexLexer("magic_start", r"(-- MAGIC %)([^\n]{2,})(\r?\n)", CodeSegment),
        RegexLexer(
            "bare_magic_sql",
            r"(\r?\n)+-- COMMAND ----------(\r?\n)+%sql\b[^\r\n]*",
            CommentSegment,
            subdivider=RegexLexer("newline", r"\r\n|\n", NewlineSegment),
            trim_post_subdivide=RegexLexer(
                "command", r"-- COMMAND ----------", CodeSegment
            ),
        ),
        RegexLexer(
            "bare_magic_cell",
            r"(?s)(\r?\n)+-- COMMAND ----------(\r?\n)+"
            r"%(?:python|scala|r|sh|md|run|fs|pip|conda)\b[^\r\n]*"
            r"(?:(?!(?:\r?\n){2}-- COMMAND ----------(?:\r?\n)).)*"
            r"(?=(?:\r?\n){2}-- COMMAND ----------(?:\r?\n)|\Z)",
            CodeSegment,
            subdivider=RegexLexer("newline", r"\r\n|\n", NewlineSegment),
            trim_post_subdivide=RegexLexer(
                "command", r"-- COMMAND ----------", CodeSegment
            ),
        ),
        RegexLexer(
            "command",
            r"(\r?\n){2}-- COMMAND ----------(\r?\n)",
            CodeSegment,
            subdivider=RegexLexer("newline", r"\r\n|\n", NewlineSegment),
        ),
    ],
    before="inline_comment",
)


databricks_dialect.add(
    CommandCellSegment=TypedParser("command", CodeSegment, type="command_cell"),
    DoubleQuotedUDFBody=TypedParser(
        "double_quote",
        CodeSegment,
        type="udf_body",
        trim_chars=('"',),
    ),
    SingleQuotedUDFBody=TypedParser(
        "single_quote",
        CodeSegment,
        type="udf_body",
        trim_chars=("'",),
    ),
    DollarQuotedUDFBody=TypedParser(
        "dollar_quote",
        CodeSegment,
        type="udf_body",
        trim_chars=("$",),
    ),
    PipelineParameterSegment=TypedParser(
        "pipeline_parameter",
        CodeSegment,
        type="pipeline_parameter",
    ),
    RightArrowSegment=StringParser("=>", SymbolSegment, type="right_arrow"),
    TryCastOperatorSegment=StringParser(
        "?::", SymbolSegment, type="try_casting_operator"
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-principal.html
    PrincipalIdentifierSegment=OneOf(
        Ref("NakedIdentifierSegment"),
        Ref("BackQuotedIdentifierSegment"),
    ),
)

# Override SingleIdentifierGrammar to include parameterized segments
databricks_dialect.replace(
    SingleIdentifierGrammar=sparksql_dialect.get_grammar(
        "SingleIdentifierGrammar"
    ).copy(
        insert=[
            Ref("ParameterizedSegment"),
        ]
    ),
    CollateGrammar=Sequence("COLLATE", Ref("CollationReferenceSegment")),
)

databricks_dialect.add(
    PredictiveOptimizationGrammar=Sequence(
        OneOf("ENABLE", "DISABLE", "INHERIT"),
        "PREDICTIVE",
        "OPTIMIZATION",
    ),
    SetOwnerGrammar=Sequence(
        Ref.keyword("SET", optional=True),
        "OWNER",
        "TO",
        Ref("PrincipalIdentifierSegment"),
    ),
    SetTagOnGrammar=Sequence("SET", "TAG", "ON"),
    UnsetTagOnGrammar=Sequence("UNSET", "TAG", "ON"),
    SetTagsGrammar=Sequence(
        "SET",
        "TAGS",
        Ref("BracketedPropertyListGrammar"),
    ),
    UnsetTagsGrammar=Sequence(
        "UNSET",
        "TAGS",
        Ref("BracketedPropertyNameListGrammar"),
    ),
    ColumnDefaultGrammar=Sequence(
        "DEFAULT",
        OneOf(
            Ref("LiteralGrammar"),
            Ref("FunctionSegment"),
        ),
    ),
    ConstraintOptionGrammar=Sequence(
        Sequence("ENABLE", "NOVALIDATE", optional=True),
        Sequence("NOT", "ENFORCED", optional=True),
        Sequence("DEFERRABLE", optional=True),
        Sequence("INITIALLY", "DEFERRED", optional=True),
        OneOf("NORELY", "RELY", optional=True),
    ),
    ForeignKeyOptionGrammar=Sequence(
        Sequence("MATCH", "FULL", optional=True),
        AnySetOf(
            Sequence("ON", "UPDATE", "NO", "ACTION"),
            Sequence("ON", "DELETE", "NO", "ACTION"),
        ),
    ),
    DropConstraintGrammar=Sequence(
        "DROP",
        OneOf(
            Sequence(
                Ref("PrimaryKeyGrammar"),
                Ref("IfExistsGrammar", optional=True),
                OneOf(
                    "RESTRICT",
                    "CASCADE",
                    optional=True,
                ),
            ),
            Sequence(
                Ref("ForeignKeyGrammar"),
                Ref("IfExistsGrammar", optional=True),
                Bracketed(
                    Delimited(
                        Ref("ColumnReferenceSegment"),
                    )
                ),
            ),
            Sequence(
                "CONSTRAINT",
                Ref("IfExistsGrammar", optional=True),
                Ref("ObjectReferenceSegment"),
                OneOf(
                    "RESTRICT",
                    "CASCADE",
                    optional=True,
                ),
            ),
        ),
    ),
    AlterPartitionGrammar=Sequence(
        "PARTITION",
        Bracketed(
            Delimited(
                AnyNumberOf(
                    OneOf(
                        Ref("ColumnReferenceSegment"),
                        Ref("SetClauseSegment"),
                    ),
                    min_times=1,
                ),
            ),
        ),
    ),
    RowFilterClauseGrammar=Sequence(
        "ROW",
        "FILTER",
        Ref("ObjectReferenceSegment"),
        "ON",
        Bracketed(
            Delimited(
                OneOf(
                    Ref("ColumnReferenceSegment"),
                    Ref("LiteralGrammar"),
                ),
                optional=True,
            ),
        ),
    ),
    PropertiesBackTickedIdentifierSegment=RegexParser(
        r"`.+`",
        IdentifierSegment,
        type="properties_naked_identifier",
    ),
    LocationWithCredentialGrammar=Sequence(
        "LOCATION",
        Ref("QuotedLiteralSegment"),
        Sequence(
            "WITH",
            Bracketed(
                "CREDENTIAL",
                Ref("PrincipalIdentifierSegment"),
            ),
            optional=True,
        ),
    ),
    NotebookStart=TypedParser("notebook_start", CommentSegment, type="notebook_start"),
    MagicSingleLineGrammar=TypedParser(
        "magic_single_line", CodeSegment, type="magic_single_line"
    ),
    MagicLineGrammar=TypedParser("magic_line", CodeSegment, type="magic_line"),
    MagicStartGrammar=TypedParser("magic_start", CodeSegment, type="magic_start"),
    BareMagicCellGrammar=TypedParser(
        "bare_magic_cell", CodeSegment, type="bare_magic_cell"
    ),
    NotebookStartBareMagicCellGrammar=TypedParser(
        "notebook_start_bare_magic_cell", CodeSegment, type="bare_magic_cell"
    ),
    VariableNameIdentifierSegment=OneOf(
        Ref("NakedIdentifierSegment"),
        Ref("BackQuotedIdentifierSegment"),
    ),
)

databricks_dialect.replace(
    DelimiterGrammar=OneOf(Ref("SemicolonSegment"), Ref("CommandCellSegment")),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-describe-volume.html
    DescribeObjectGrammar=sparksql_dialect.get_grammar("DescribeObjectGrammar").copy(
        insert=[
            Sequence(
                "VOLUME",
                Ref("VolumeReferenceSegment"),
            ),
        ],
        at=0,
    ),
    # Add ParameterizedSegment to the LiteralGrammar to support named parameters
    LiteralGrammar=sparksql_dialect.get_grammar("LiteralGrammar").copy(
        insert=[
            Ref("ParameterizedSegment"),
        ]
    ),
    FunctionContentsExpressionGrammar=OneOf(
        Ref("ExpressionSegment"),
        Ref("NamedArgumentSegment"),
    ),
    PropertiesNakedIdentifierSegment=RegexParser(
        r"[A-Z_][A-Z0-9_]*",
        IdentifierSegment,
        type="properties_naked_identifier",
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-show-schemas.html
    # Differences between this and the SparkSQL version:
    # - Support for `FROM`|`IN` at the catalog level
    # - `LIKE` keyword is optional
    ShowDatabasesSchemasGrammar=Sequence(
        # SHOW { DATABASES | SCHEMAS }
        OneOf("DATABASES", "SCHEMAS"),
        Sequence(
            OneOf("FROM", "IN"),
            Ref("DatabaseReferenceSegment"),
            optional=True,
        ),
        Sequence(
            Ref.keyword("LIKE", optional=True),
            Ref("QuotedLiteralSegment"),
            optional=True,
        ),
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-show-functions.html
    # Differences between this and the SparkSQL version:
    # - Support for `FROM`|`IN` at the schema level
    # - `LIKE` keyword is optional
    ShowFunctionsGrammar=Sequence(
        # SHOW FUNCTIONS
        OneOf("USER", "SYSTEM", "ALL", optional=True),
        "FUNCTIONS",
        Sequence(
            Sequence(
                OneOf("FROM", "IN"),
                Ref("DatabaseReferenceSegment"),
                optional=True,
            ),
            Sequence(
                Ref.keyword("LIKE", optional=True),
                OneOf(
                    # qualified function from a database
                    Sequence(
                        Ref("DatabaseReferenceSegment"),
                        Ref("DotSegment"),
                        Ref("FunctionNameSegment"),
                        allow_gaps=False,
                    ),
                    # non-qualified function
                    Ref("FunctionNameSegment"),
                    # Regex/like string
                    Ref("QuotedLiteralSegment"),
                ),
                optional=True,
            ),
            optional=True,
        ),
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-show-tables.html
    # Differences between this and the SparkSQL version:
    # - `LIKE` keyword is optional
    ShowTablesGrammar=Sequence(
        # SHOW TABLES
        "TABLES",
        Sequence(
            OneOf("FROM", "IN"),
            Ref("DatabaseReferenceSegment"),
            optional=True,
        ),
        Sequence(
            Ref.keyword("LIKE", optional=True),
            Ref("QuotedLiteralSegment"),
            optional=True,
        ),
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-show-views.html
    # Only difference between this and the SparkSQL version:
    # - `LIKE` keyword is optional
    ShowViewsGrammar=Sequence(
        # SHOW VIEWS
        "VIEWS",
        Sequence(
            OneOf("FROM", "IN"),
            Ref("DatabaseReferenceSegment"),
            optional=True,
        ),
        Sequence(
            Ref.keyword("LIKE", optional=True),
            Ref("QuotedLiteralSegment"),
            optional=True,
        ),
    ),
    # https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-show-volumes.html
    ShowObjectGrammar=sparksql_dialect.get_grammar("ShowObjectGrammar").copy(
        insert=[
            Sequence(
                "VOLUMES",
                Sequence(
                    OneOf("FROM", "IN"),
                    Ref("DatabaseReferenceSegment"),
                    optional=True,
                ),
                Sequence(
                    Ref.keyword("LIKE", optional=True),
                    Ref("QuotedLiteralSegment"),
                    optional=True,
                ),
            )
        ],
    ),
    NotNullGrammar=Sequence(
        "NOT",
        "NULL",
    ),
    FunctionNameIdentifierSegment=OneOf(
        TypedParser("word", WordSegment, type="function_name_identifier"),
        Ref("BackQuotedIdentifierSegment"),
    ),
    PreTableFunctionKeywordsGrammar=OneOf("STREAM"),
    ColumnGeneratedGrammar=OneOf(
        Sequence(
            "GENERATED",
            "ALWAYS",
            "AS",
            Bracketed(
                OneOf(
                    Ref("FunctionSegment"),
                    Ref("BareFunctionSegment"),
                    Ref("ExpressionSegment"),
                ),
            ),
        ),
        Sequence(
            "GENERATED",
            OneOf(
                "ALWAYS",
                Sequence("BY", "DEFAULT"),
            ),
            "AS",
            "IDENTITY",
            Bracketed(
                Sequence(
                    Sequence(
                        "START",
                        "WITH",
                        Ref("NumericLiteralSegment"),
                        optional=True,
                    ),
                    Sequence(
                        "INCREMENT",
                        "BY",
                        Ref("NumericLiteralSegment"),
                        optional=True,
                    ),
                ),
                optional=True,
            ),
        ),
    ),
)

databricks_dialect.replace(
    PostFunctionGrammar=sparksql_dialect.get_grammar("PostFunctionGrammar").copy(
        insert=[
            Ref("WithinGroupClauseSegment"),
        ],
    ),
)


class WithinGroupClauseSegment(BaseSegment):
    """An WITHIN GROUP clause for ordered-set aggregate functions.

    https://docs.databricks.com/en/sql/language-manual/functions/percentile_cont.html
    """

    type = "withingroup_clause"
    match_grammar = Sequence(
        "WITHIN",
        "GROUP",
        Bracketed(Ref("OrderByClauseSegment")),
    )


class ShorthandCastSegment(ansi.ShorthandCastSegment):
    """A casting operation using '::' or '?::'.

    '?::' is shorthand for try_cast (an error-tolerating cast).
    https://docs.databricks.com/aws/en/sql/language-manual/functions/questiondoublecolonsign
    """

    match_grammar: Matchable = Sequence(
        OneOf(
            Ref("Expression_D_Grammar"),
            Ref("CaseExpressionSegment"),
        ),
        AnyNumberOf(
            Sequence(
                OneOf(Ref("CastOperatorSegment"), Ref("TryCastOperatorSegment")),
                Ref("DatatypeSegment"),
                Ref("TimeZoneGrammar", optional=True),
                allow_gaps=True,
            ),
            min_times=1,
        ),
    )


class IdentifierClauseSegment(BaseSegment):
    """An `IDENTIFIER` clause segment.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-names-identifier-clause.html
    """

    type = "identifier_clause_segment"
    match_grammar = Sequence(
        "IDENTIFIER",
        Bracketed(Ref("ExpressionSegment")),
    )


class ObjectReferenceSegment(ansi.ObjectReferenceSegment):
    """A reference to an object."""

    # Allow whitespace
    match_grammar: Matchable = Delimited(
        OneOf(Ref("SingleIdentifierGrammar"), Ref("IdentifierClauseSegment")),
        delimiter=Ref("ObjectReferenceDelimiterGrammar"),
        terminators=[Ref("ObjectReferenceTerminatorGrammar")],
        allow_gaps=False,
    )


class DatabaseReferenceSegment(ObjectReferenceSegment):
    """A reference to a database."""

    type = "database_reference"


class TableReferenceSegment(ObjectReferenceSegment):
    """A reference to an table, CTE, subquery or alias."""

    type = "table_reference"


class SchemaReferenceSegment(ObjectReferenceSegment):
    """A reference to a schema."""

    type = "schema_reference"


class TableExpressionSegment(sparksql.TableExpressionSegment):
    """The main table expression e.g. within a FROM clause.

    Enhance to allow for additional clauses allowed in Spark and Delta Lake.
    """

    match_grammar = sparksql.TableExpressionSegment.match_grammar.copy(
        insert=[
            Ref("IdentifierClauseSegment"),
        ],
        before=Ref("ValuesClauseSegment"),
    )


class CatalogReferenceSegment(ansi.ObjectReferenceSegment):
    """A reference to a catalog.

    https://docs.databricks.com/data-governance/unity-catalog/create-catalogs.html
    """

    type = "catalog_reference"

    # Allow catalog names to be identifiers or parameters
    match_grammar: Matchable = OneOf(
        Delimited(
            OneOf(Ref("SingleIdentifierGrammar"), Ref("IdentifierClauseSegment")),
            delimiter=Ref("ObjectReferenceDelimiterGrammar"),
            terminators=[Ref("ObjectReferenceTerminatorGrammar")],
            allow_gaps=False,
        ),
        Ref("ParameterizedSegment"),
    )


class VolumeReferenceSegment(ansi.ObjectReferenceSegment):
    """Volume reference."""

    type = "volume_reference"


class AlterCatalogStatementSegment(BaseSegment):
    """An `ALTER CATALOG` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-alter-catalog.html
    """

    type = "alter_catalog_statement"
    match_grammar = Sequence(
        "ALTER",
        "CATALOG",
        Ref("CatalogReferenceSegment"),
        OneOf(
            Ref("SetOwnerGrammar"),
            Ref("SetTagsGrammar"),
            Ref("UnsetTagsGrammar"),
            Ref("PredictiveOptimizationGrammar"),
        ),
    )


class CreateCatalogStatementSegment(BaseSegment):
    """A `CREATE CATALOG` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-create-catalog.html
    """

    type = "create_catalog_statement"
    match_grammar = Sequence(
        "CREATE",
        "CATALOG",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("CatalogReferenceSegment"),
        Ref("CommentGrammar", optional=True),
    )


class DropCatalogStatementSegment(BaseSegment):
    """A `DROP CATALOG` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-drop-catalog.html
    """

    type = "drop_catalog_statement"
    match_grammar = Sequence(
        "DROP",
        "CATALOG",
        Ref("IfExistsGrammar", optional=True),
        Ref("CatalogReferenceSegment"),
        Ref("DropBehaviorGrammar", optional=True),
    )


class UseCatalogStatementSegment(BaseSegment):
    """A `USE CATALOG` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-use-catalog.html
    """

    type = "use_catalog_statement"
    match_grammar = Sequence(
        "USE",
        "CATALOG",
        Ref("CatalogReferenceSegment"),
    )


class UseDatabaseStatementSegment(sparksql.UseDatabaseStatementSegment):
    """A `USE DATABASE` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-ddl-usedb.html
    """

    type = "use_database_statement"
    match_grammar = Sequence(
        "USE",
        OneOf("DATABASE", "SCHEMA", optional=True),
        Ref("DatabaseReferenceSegment"),
    )


class AlterDatabaseStatementSegment(sparksql.AlterDatabaseStatementSegment):
    """An `ALTER DATABASE/SCHEMA` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-alter-schema.html
    """

    match_grammar = Sequence(
        "ALTER",
        OneOf("DATABASE", "SCHEMA"),
        Ref("DatabaseReferenceSegment"),
        OneOf(
            Sequence(
                "SET",
                Ref("DatabasePropertiesGrammar"),
            ),
            Ref("SetOwnerGrammar"),
            Ref("SetTagsGrammar"),
            Ref("UnsetTagsGrammar"),
            Ref("PredictiveOptimizationGrammar"),
        ),
    )


class AlterVolumeStatementSegment(BaseSegment):
    """Alter Volume Statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-alter-volume.html
    """

    type = "alter_volume_statement"

    match_grammar = Sequence(
        "ALTER",
        "VOLUME",
        Ref("VolumeReferenceSegment"),
        OneOf(
            Sequence(
                "RENAME",
                "TO",
                Ref("VolumeReferenceSegment"),
            ),
            Ref("SetOwnerGrammar"),
            Ref("SetTagsGrammar"),
            Ref("UnsetTagsGrammar"),
        ),
    )


class CreateVolumeStatementSegment(BaseSegment):
    """Create Volume Statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-volume.html
    """

    type = "create_volume_statement"

    match_grammar = OneOf(
        # You can create a non-external volume without a location
        Sequence(
            "CREATE",
            "VOLUME",
            Ref("IfNotExistsGrammar", optional=True),
            Ref("VolumeReferenceSegment"),
            Ref("CommentGrammar", optional=True),
        ),
        # Or you can create an external volume that must have a location
        Sequence(
            "CREATE",
            "EXTERNAL",
            "VOLUME",
            Ref("IfNotExistsGrammar", optional=True),
            Ref("VolumeReferenceSegment"),
            Ref("LocationGrammar"),
            Ref("CommentGrammar", optional=True),
        ),
    )


class DropVolumeStatementSegment(BaseSegment):
    """Drop Volume Statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-drop-volume.html
    """

    type = "drop_volume_statement"

    match_grammar = Sequence(
        "DROP",
        "VOLUME",
        Ref("IfExistsGrammar", optional=True),
        Ref("VolumeReferenceSegment"),
    )


class CreateDatabaseStatementSegment(sparksql.CreateDatabaseStatementSegment):
    """A `CREATE DATABASE` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-schema.html
    """

    match_grammar = sparksql.CreateDatabaseStatementSegment.match_grammar.copy(
        insert=[
            Sequence(
                Ref.keyword("MANAGED", optional=True),
                "LOCATION",
                Ref("QuotedLiteralSegment"),
                optional=True,
            ),
        ],
        at=5,
        remove=[
            Ref("LocationGrammar", optional=True),
        ],
    )


class CreateViewStatementSegment(BaseSegment):
    """A `CREATE VIEW` statement.

    https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-create-view

    Does not inherit from SparkSQL to properly support Databricks-specific syntax
    and to be distinct from `CREATE MATERIALIZED VIEW`.
    """

    type = "create_view_statement"

    _schema_binding_clause = Sequence(
        "WITH",
        OneOf(
            "METRICS",
            Sequence(
                "SCHEMA",
                OneOf(
                    "BINDING",
                    "COMPENSATION",
                    Sequence(
                        Ref.keyword("TYPE", optional=True),
                        "EVOLUTION",
                    ),
                ),
            ),
        ),
    )

    _view_clauses = AnyNumberOf(
        Ref("CommentClauseSegment"),
        Sequence(
            "DEFAULT",
            "COLLATION",
            Ref("ObjectReferenceSegment"),  # collation_name
        ),
        Ref("TablePropertiesGrammar"),
        Sequence("LANGUAGE", "YAML"),
        _schema_binding_clause,
    )

    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        Ref("TemporaryGrammar", optional=True),
        "VIEW",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("TableReferenceSegment"),
        Sequence(
            Bracketed(
                Delimited(
                    Sequence(
                        Ref("ColumnReferenceSegment"),
                        Ref("CommentClauseSegment", optional=True),
                    ),
                ),
            ),
            optional=True,
        ),
        _view_clauses,
        "AS",
        OneOf(
            OptionallyBracketed(Ref("SelectableGrammar")),
            # YAML metric view definition: $$ yaml_string $$
            Ref("DollarQuotedUDFBody"),
        ),
    )


class MaterializedViewExpectationConstraintSegment(BaseSegment):
    """A data quality expectation for a materialized view."""

    type = "constraint_statement"

    match_grammar = Sequence(
        "CONSTRAINT",
        Ref("ObjectReferenceSegment"),
        "EXPECT",
        Bracketed(Ref("ExpressionSegment")),
        Sequence(
            "ON",
            "VIOLATION",
            OneOf(
                Sequence("FAIL", "UPDATE"),
                Sequence("DROP", "ROW"),
            ),
            optional=True,
        ),
    )


class CreateMaterializedViewStatementSegment(BaseSegment):
    """A `CREATE MATERIALIZED VIEW` Statement.

    Covers both standard Databricks SQL and Lakeflow/DLT syntax:
    https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-create-materialized-view
    https://docs.databricks.com/aws/en/ldp/developer/ldp-sql-ref-create-materialized-view
    """

    type = "create_materialized_view_statement"

    _schedule = OneOf(
        Sequence(
            "SCHEDULE",
            Ref.keyword("REFRESH", optional=True),
            OneOf(
                Sequence(
                    "EVERY",
                    Ref("NumericLiteralSegment"),
                    OneOf("HOUR", "HOURS", "DAY", "DAYS", "WEEK", "WEEKS"),
                ),
                Sequence(
                    "CRON",
                    Ref("QuotedLiteralSegment"),
                    Sequence(
                        "AT",
                        "TIME",
                        "ZONE",
                        Ref("QuotedLiteralSegment"),
                        optional=True,
                    ),
                ),
            ),
        ),
        Sequence(
            "TRIGGER",
            "ON",
            "UPDATE",
            Sequence(
                "AT",
                "MOST",
                "EVERY",
                Ref("IntervalExpressionSegment"),
                optional=True,
            ),
        ),
        optional=True,
    )

    _view_clauses = AnyNumberOf(
        Ref("PartitionSpecGrammar"),
        Ref("TableClusterByClauseSegment"),
        Ref("CommentGrammar"),
        Sequence(
            "DEFAULT",
            "COLLATION",
            Ref("ObjectReferenceSegment"),
        ),
        Ref("TablePropertiesGrammar"),
        _schedule,
        Sequence(
            "WITH",
            Ref("RowFilterClauseGrammar"),
        ),
    )

    match_grammar = Sequence(
        "CREATE",
        # OR REPLACE for standard SQL, OR REFRESH for DLT
        OneOf(Ref("OrReplaceGrammar"), Ref("OrRefreshGrammar"), optional=True),
        Ref.keyword("PRIVATE", optional=True),  # DLT-specific
        "MATERIALIZED",
        "VIEW",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("TableReferenceSegment"),
        Bracketed(
            Sequence(
                Ref("ColumnFieldDefinitionSegment"),
                AnyNumberOf(
                    Sequence(
                        Ref("CommaSegment"),
                        Ref("ColumnFieldDefinitionSegment"),
                    ),
                ),
                AnyNumberOf(
                    Sequence(
                        Ref("CommaSegment"),
                        Ref("MaterializedViewExpectationConstraintSegment"),
                    ),
                ),
                AnyNumberOf(
                    Sequence(
                        Ref("CommaSegment"),
                        Ref("TableConstraintSegment"),
                    ),
                ),
            ),
            optional=True,
        ),
        _view_clauses,
        "AS",
        OptionallyBracketed(Ref("SelectableGrammar")),
    )


class MaskStatementSegment(BaseSegment):
    """A `MASK` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-column-mask.html
    """

    type = "mask_statement"
    match_grammar = Sequence(
        "MASK",
        Ref("FunctionNameSegment"),
        Sequence(
            "USING",
            "COLUMNS",
            Bracketed(
                AnyNumberOf(
                    OneOf(
                        Ref("ColumnReferenceSegment"),
                        Ref("ExpressionSegment"),
                    ),
                ),
            ),
            optional=True,
        ),
    )


class ColumnFieldDefinitionSegment(ansi.ColumnDefinitionSegment):
    """A column field definition, e.g. for CREATE TABLE or ALTER TABLE.

    This supports the iceberg syntax and allows for iceberg syntax such
    as ADD COLUMN a.b.
    """

    match_grammar: Matchable = Sequence(
        Ref("ColumnReferenceSegment"),  # Column name
        Ref("DatatypeSegment"),  # Column type
        Bracketed(Anything(), optional=True),  # For types like VARCHAR(100)
        AnyNumberOf(
            Ref("ColumnPropertiesSegment"),
            Ref("ColumnConstraintSegment"),
            Ref("ColumnDefaultGrammar"),  # For default values
        ),
    )


class PropertyNameSegment(sparksql.PropertyNameSegment):
    """A property name segment. Databricks allows for back quoted segments."""

    match_grammar = Sequence(
        OneOf(
            Delimited(
                OneOf(
                    Ref("PropertiesNakedIdentifierSegment"),
                    Ref("PropertiesBackTickedIdentifierSegment"),
                ),
                delimiter=Ref("DotSegment"),
                allow_gaps=False,
            ),
            Ref("SingleIdentifierGrammar"),
        ),
    )


class TableConstraintSegment(ansi.TableConstraintSegment):
    """A table constraint, e.g. for CREATE TABLE or ALTER TABLE.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-constraint.html
    """

    match_grammar = Sequence(
        "CONSTRAINT",
        OneOf(
            Sequence(
                Ref("ObjectReferenceSegment", optional=True),
                Ref("PrimaryKeyGrammar"),
                Bracketed(
                    Delimited(
                        Ref("ColumnReferenceSegment"),
                        Ref.keyword("TIMESERIES", optional=True),
                    ),
                ),
                Ref("ConstraintOptionGrammar", optional=True),
            ),
            Sequence(
                Ref("ObjectReferenceSegment", optional=True),
                Indent,
                Ref("ForeignKeyGrammar"),
                Bracketed(
                    Delimited(
                        Ref("ColumnReferenceSegment"),
                    ),
                ),
                "REFERENCES",
                Ref("TableReferenceSegment"),
                Ref("BracketedColumnReferenceListGrammar", optional=True),
                OneOf(
                    Ref("ForeignKeyOptionGrammar"),
                    Ref("ConstraintOptionGrammar"),
                    optional=True,
                ),
                Dedent,
            ),
            Sequence(
                Ref("ObjectReferenceSegment"),
                "CHECK",
                Bracketed(Ref("ExpressionSegment")),
                Ref.keyword("ENFORCED", optional=True),
            ),
        ),
    )


class UnsetTagStatementSegment(BaseSegment):
    """An `UNSET TAG ON` statement.

    https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-set-tag
    """

    type = "tag_statement"
    match_grammar = Sequence(
        Ref("UnsetTagOnGrammar"),
        OneOf(
            Sequence(
                "CATALOG",
                Ref("CatalogReferenceSegment"),
            ),
            Sequence(
                OneOf("DATABASE", "SCHEMA"),
                Ref("DatabaseReferenceSegment"),
            ),
            Sequence(
                OneOf("TABLE", "VIEW"),
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "VOLUME",
                Ref("VolumeReferenceSegment"),
            ),
            Sequence(
                "COLUMN",
                Ref("ColumnReferenceSegment"),
            ),
        ),
        OneOf(Ref("BackQuotedIdentifierSegment"), Ref("NakedIdentifierSegment")),
    )


class TagStatementSegment(BaseSegment):
    """An `SET TAG ON` statement.

    https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-set-tag
    """

    type = "tag_statement"
    match_grammar = Sequence(
        Ref("SetTagOnGrammar"),
        OneOf(
            Sequence(
                "CATALOG",
                Ref("CatalogReferenceSegment"),
            ),
            Sequence(
                OneOf("DATABASE", "SCHEMA"),
                Ref("DatabaseReferenceSegment"),
            ),
            Sequence(
                OneOf("TABLE", "VIEW"),
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "VOLUME",
                Ref("VolumeReferenceSegment"),
            ),
            Sequence(
                "COLUMN",
                Ref("ColumnReferenceSegment"),
            ),
        ),
        OneOf(Ref("BackQuotedIdentifierSegment"), Ref("NakedIdentifierSegment")),
        Ref("EqualsSegment"),
        OneOf(Ref("BackQuotedIdentifierSegment"), Ref("NakedIdentifierSegment")),
    )


class AlterTableStatementSegment(sparksql.AlterTableStatementSegment):
    """An `ALTER TABLE` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-alter-table.html
    """

    match_grammar = Sequence(
        "ALTER",
        "TABLE",
        Ref("TableReferenceSegment"),
        Indent,
        OneOf(
            Sequence(
                "RENAME",
                "TO",
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "ADD",
                OneOf("COLUMNS", "COLUMN"),
                Indent,
                OptionallyBracketed(
                    Delimited(
                        Sequence(
                            Ref("ColumnFieldDefinitionSegment"),
                            Ref("ColumnDefaultGrammar", optional=True),
                            Ref("CommentGrammar", optional=True),
                            Ref("FirstOrAfterGrammar", optional=True),
                            Ref("MaskStatementSegment", optional=True),
                        ),
                    ),
                ),
                Dedent,
            ),
            Sequence(
                OneOf("ALTER", "CHANGE"),
                Ref.keyword("COLUMN", optional=True),
                Ref("ColumnReferenceSegment"),
                OneOf(
                    Ref("CommentGrammar"),
                    Ref("FirstOrAfterGrammar"),
                    Sequence(
                        OneOf("SET", "DROP"),
                        "NOT",
                        "NULL",
                    ),
                    Sequence(
                        "TYPE",
                        Ref("DatatypeSegment"),
                    ),
                    Sequence(
                        "SET",
                        Ref("ColumnDefaultGrammar"),
                    ),
                    Sequence(
                        "DROP",
                        "DEFAULT",
                    ),
                    Sequence(
                        "SYNC",
                        "IDENTITY",
                    ),
                    Sequence(
                        "SET",
                        Ref("MaskStatementSegment"),
                    ),
                    Sequence(
                        "DROP",
                        "MASK",
                    ),
                    Ref("SetTagsGrammar"),
                    Ref("UnsetTagsGrammar"),
                ),
            ),
            Sequence(
                "DROP",
                OneOf("COLUMN", "COLUMNS", optional=True),
                Ref("IfExistsGrammar", optional=True),
                OptionallyBracketed(
                    Delimited(
                        Ref("ColumnReferenceSegment"),
                    ),
                ),
            ),
            Sequence(
                "RENAME",
                "COLUMN",
                Ref("ColumnReferenceSegment"),
                "TO",
                Ref("ColumnReferenceSegment"),
            ),
            Sequence(
                "ADD",
                Ref("TableConstraintSegment"),
            ),
            Ref("DropConstraintGrammar"),
            Sequence(
                "DROP",
                "FEATURE",
                Ref("ObjectReferenceSegment"),
                Sequence(
                    "TRUNCATE",
                    "HISTORY",
                    optional=True,
                ),
            ),
            Sequence(
                "ADD",
                Ref("IfNotExistsGrammar", optional=True),
                AnyNumberOf(Ref("AlterPartitionGrammar")),
            ),
            Sequence(
                "DROP",
                Ref("IfExistsGrammar", optional=True),
                AnyNumberOf(Ref("AlterPartitionGrammar")),
            ),
            Sequence(
                Ref("AlterPartitionGrammar"),
                "SET",
                Ref("LocationGrammar"),
            ),
            Sequence(
                Ref("AlterPartitionGrammar"),
                "RENAME",
                "TO",
                Ref("AlterPartitionGrammar"),
            ),
            Sequence(
                "RECOVER",
                "PARTITIONS",
            ),
            Sequence(
                "SET",
                Ref("RowFilterClauseGrammar"),
            ),
            Sequence(
                "DROP",
                "ROW",
                "FILTER",
            ),
            Sequence(
                "SET",
                Ref("TablePropertiesGrammar"),
            ),
            Ref("UnsetTablePropertiesGrammar"),
            Sequence(
                "SET",
                "SERDE",
                Ref("QuotedLiteralSegment"),
                Sequence(
                    "WITH",
                    "SERDEPROPERTIES",
                    Ref("BracketedPropertyListGrammar"),
                    optional=True,
                ),
            ),
            Sequence(
                "SET",
                Ref("LocationGrammar"),
            ),
            Ref("SetOwnerGrammar"),
            Sequence(
                Sequence(
                    "ALTER",
                    "COLUMN",
                    Ref("ColumnReferenceSegment"),
                    optional=True,
                ),
                Ref("SetTagsGrammar"),
            ),
            Sequence(
                Sequence(
                    "ALTER",
                    "COLUMN",
                    Ref("ColumnReferenceSegment"),
                    optional=True,
                ),
                Ref("UnsetTagsGrammar"),
            ),
            Ref("TableClusterByClauseSegment"),
            Ref("PredictiveOptimizationGrammar"),
        ),
        Dedent,
    )


class AlterViewStatementSegment(sparksql.AlterViewStatementSegment):
    """An `ALTER VIEW` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-alter-view.html
    """

    match_grammar = Sequence(
        "ALTER",
        Ref.keyword("MATERIALIZED", optional=True),
        "VIEW",
        Ref("TableReferenceSegment"),
        OneOf(
            Sequence(
                "RENAME",
                "TO",
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "SET",
                Ref("TablePropertiesGrammar"),
            ),
            Ref("UnsetTablePropertiesGrammar"),
            Sequence(
                "AS",
                Ref("SelectStatementSegment"),
            ),
            Sequence(
                "WITH",
                "SCHEMA",
                OneOf(
                    "BINDING",
                    "COMPENSATION",
                    Sequence(
                        Ref.keyword("TYPE", optional=True),
                        "EVOLUTION",
                    ),
                ),
            ),
            Ref("SetOwnerGrammar"),
            Ref("SetTagsGrammar"),
            Ref("UnsetTagsGrammar"),
            Sequence(
                Indent,
                OneOf(
                    Sequence(
                        OneOf("ADD", "ALTER"),
                        "SCHEDULE",
                        Ref.keyword("REFRESH", optional=True),
                        "CRON",
                        Ref("QuotedLiteralSegment"),
                        Sequence(
                            "AT",
                            "TIME",
                            "ZONE",
                            Ref("QuotedLiteralSegment"),
                            optional=True,
                        ),
                    ),
                    Sequence(
                        "DROP",
                        "SCHEDULE",
                    ),
                ),
                Dedent,
            ),
        ),
    )


class SetTimeZoneStatementSegment(BaseSegment):
    """A `SET TIME ZONE` statement.

    https://docs.databricks.com/sql/language-manual/sql-ref-syntax-aux-conf-mgmt-set-timezone.html
    """

    type = "set_timezone_statement"
    match_grammar = Sequence(
        "SET",
        "TIME",
        "ZONE",
        OneOf("LOCAL", Ref("QuotedLiteralSegment"), Ref("IntervalExpressionSegment")),
    )


class TableClusterByClauseSegment(sparksql.TableClusterByClauseSegment):
    """A `CLUSTER BY` clause in table definitions.

    https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-ddl-cluster-by
    """

    match_grammar = Sequence(
        "CLUSTER",
        "BY",
        Indent,
        OneOf(
            Ref("BracketedColumnReferenceListGrammar"),
            "AUTO",
            "NONE",
        ),
        Dedent,
    )


class OptimizeTableStatementSegment(BaseSegment):
    """An `OPTIMIZE` statement.

    https://docs.databricks.com/en/sql/language-manual/delta-optimize.html
    """

    type = "optimize_table_statement"
    match_grammar = Sequence(
        "OPTIMIZE",
        Ref("TableReferenceSegment"),
        Sequence(
            "WHERE",
            Ref("ExpressionSegment"),
            optional=True,
        ),
        Sequence(
            "ZORDER",
            "BY",
            Bracketed(Delimited(Ref("ColumnReferenceSegment"))),
            optional=True,
        ),
    )


class LimitClauseSegment(sparksql.LimitClauseSegment):
    """A `LIMIT` clause like in `SELECT`.

    Enhanced from SparkSQL to support parameterized values.
    """

    match_grammar = Sequence(
        "LIMIT",
        Indent,
        OneOf(
            Ref("NumericLiteralSegment"),
            "ALL",
            Ref("FunctionSegment"),
            Ref("ParameterizedSegment"),  # Add support for parameters
        ),
        Dedent,
    )


class MergeInsertClauseSegment(sparksql.MergeInsertClauseSegment):
    """`INSERT` clause within the `MERGE` statement.

    Databricks MERGE should not treat trailing `WHEN` as an alias for the
    inserted values clause, so we match the values list explicitly.
    """

    match_grammar: Matchable = Sequence(
        "INSERT",
        OneOf(
            Ref("WildcardExpressionSegment"),
            Sequence(
                Indent,
                Ref("BracketedColumnReferenceListGrammar"),
                Dedent,
                "VALUES",
                Bracketed(
                    Delimited(
                        Ref("ExpressionSegment"),
                    )
                ),
            ),
        ),
    )


class StatementSegment(sparksql.StatementSegment):
    """Overriding StatementSegment to allow for additional segment parsing."""

    match_grammar = sparksql.StatementSegment.match_grammar.copy(
        # Segments defined in Databricks SQL dialect
        insert=[
            # Unity Catalog
            Ref("AlterCatalogStatementSegment"),
            Ref("CreateCatalogStatementSegment"),
            Ref("DropCatalogStatementSegment"),
            Ref("UseCatalogStatementSegment"),
            Ref("AlterVolumeStatementSegment"),
            Ref("CreateVolumeStatementSegment"),
            Ref("DropVolumeStatementSegment"),
            Ref("CreateDatabaseStatementSegment"),
            Ref("SetTimeZoneStatementSegment"),
            Ref("OptimizeTableStatementSegment"),
            Ref("CreateDatabricksFunctionStatementSegment"),
            Ref("CreateTableCloneStatementSegment"),
            Ref("FunctionParameterListGrammarWithComments"),
            Ref("DeclareOrReplaceVariableStatementSegment"),
            Ref("CommentOnStatementSegment"),
            Ref("TagStatementSegment"),
            Ref("UnsetTagStatementSegment"),
            # Notebook grammar
            Ref("MagicCellStatementSegment"),
            # Databricks - Delta Live Tables
            Ref("ApplyChangesIntoStatementSegment"),
            Ref("CreateFlowStatementSegment"),
            Ref("CreateMaterializedViewStatementSegment"),
        ]
    )


class FunctionParameterListGrammarWithComments(BaseSegment):
    """The parameters for a function ie. `(column type COMMENT 'comment')`."""

    type = "function_parameter_list_with_comments"

    match_grammar: Matchable = Bracketed(
        Delimited(
            Sequence(
                Ref("FunctionParameterGrammar"),
                AnyNumberOf(
                    Sequence(
                        "DEFAULT",
                        OneOf(Ref("LiteralGrammar"), Ref("FunctionSegment")),
                        optional=True,
                    ),
                    Ref("CommentClauseSegment", optional=True),
                ),
            ),
            optional=True,
        ),
    )


class FunctionDefinitionGrammar(ansi.FunctionDefinitionGrammar):
    """This is the body of a `CREATE FUNCTION AS` statement."""

    match_grammar = Sequence(
        AnyNumberOf(
            Sequence(
                "LANGUAGE",
                OneOf(Ref.keyword("SQL"), Ref.keyword("PYTHON")),
                optional=True,
            ),
            Sequence(
                OneOf("DETERMINISTIC", Sequence("NOT", "DETERMINISTIC")),
                optional=True,
            ),
            Ref("CommentClauseSegment", optional=True),
            Sequence(
                OneOf(Sequence("CONTAINS", "SQL"), Sequence("READS", "SQL", "DATA")),
                optional=True,
            ),
            Sequence(
                OneOf(
                    Sequence(
                        "AS",
                        OneOf(
                            Ref("DoubleQuotedUDFBody"),
                            Ref("SingleQuotedUDFBody"),
                            Ref("DollarQuotedUDFBody"),
                            Bracketed(
                                OneOf(
                                    Ref("ExpressionSegment"),
                                    Ref("SelectStatementSegment"),
                                )
                            ),
                        ),
                    ),
                    Sequence(
                        "RETURN",
                        OneOf(
                            Ref("SetExpressionSegment"),
                            Ref("ExpressionSegment"),
                            Ref("SelectStatementSegment"),
                            Ref("WithCompoundStatementSegment"),
                        ),
                    ),
                )
            ),
        )
    )


class CreateDatabricksFunctionStatementSegment(BaseSegment):
    """A `CREATE FUNCTION` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html
    """

    type = "create_sql_function_statement"

    match_grammar: Matchable = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        Ref("TemporaryGrammar", optional=True),
        "FUNCTION",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("FunctionNameSegment"),
        Ref("FunctionParameterListGrammarWithComments"),
        Sequence(
            "RETURNS",
            OneOf(
                Ref("DatatypeSegment"),
                Sequence(
                    "TABLE",
                    Sequence(
                        Bracketed(
                            Delimited(
                                Sequence(
                                    Ref("ColumnReferenceSegment"),
                                    Ref("DatatypeSegment"),
                                    Ref("CommentGrammar", optional=True),
                                ),
                            ),
                        ),
                        optional=True,
                    ),
                ),
            ),
            optional=True,
        ),
        Ref("FunctionDefinitionGrammar"),
    )


class NamedArgumentSegment(BaseSegment):
    """Named argument to a function.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-function-invocation.html#named-parameter-invocation
    """

    type = "named_argument"
    match_grammar = Sequence(
        Ref("VariableNameIdentifierSegment"),
        Ref("RightArrowSegment"),
        Ref("ExpressionSegment"),
    )


class AliasExpressionSegment(sparksql.AliasExpressionSegment):
    """A reference to an object with an `AS` clause.

    The optional AS keyword allows both implicit and explicit aliasing.
    Note also that it's possible to specify just column aliases without aliasing the
    table as well:
    .. code-block:: sql

        SELECT * FROM VALUES (1,2) as t (a, b);
        SELECT * FROM VALUES (1,2) as (a, b);
        SELECT * FROM VALUES (1,2) as t;

    Note that in Spark SQL, identifiers are quoted using backticks (`my_table`) rather
    than double quotes ("my_table"). Quoted identifiers are allowed in aliases, but
    unlike ANSI which allows single quoted identifiers ('my_table') in aliases, this is
    not allowed in Spark and so the definition of this segment must depart from ANSI.
    """

    match_grammar = Sequence(
        Indent,
        Ref("AsAliasOperatorSegment", optional=True),
        OneOf(
            # maybe table alias and column aliases
            Sequence(
                Ref("SingleIdentifierGrammar", optional=True),
                Bracketed(Ref("SingleIdentifierListSegment")),
            ),
            # just a table alias
            Ref("SingleIdentifierGrammar"),
            exclude=OneOf(
                "LATERAL",
                Ref("JoinTypeKeywords"),
                "WINDOW",
                "PIVOT",
                "KEYS",
                "FROM",
                "FOR",
            ),
        ),
        Dedent,
    )


class GroupByClauseSegment(sparksql.GroupByClauseSegment):
    """Enhance `GROUP BY` clause like in `SELECT` for `CUBE`, `ROLLUP`, and `ALL`.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-qry-select-groupby.html
    """

    match_grammar = Sequence(
        "GROUP",
        "BY",
        Indent,
        OneOf(
            "ALL",
            Delimited(
                Ref("CubeRollupClauseSegment"),
                Ref("GroupingSetsClauseSegment"),
                Ref("ColumnReferenceSegment"),
                # Can `GROUP BY 1`
                Ref("NumericLiteralSegment"),
                # Can `GROUP BY coalesce(col, 1)`
                Ref("ExpressionSegment"),
            ),
            Sequence(
                Delimited(
                    Ref("ColumnReferenceSegment"),
                    # Can `GROUP BY 1`
                    Ref("NumericLiteralSegment"),
                    # Can `GROUP BY coalesce(col, 1)`
                    Ref("ExpressionSegment"),
                ),
                OneOf(
                    Ref("WithCubeRollupClauseSegment"),
                    Ref("GroupingSetsClauseSegment"),
                ),
            ),
        ),
        Dedent,
    )


class ColumnConstraintSegment(ansi.ColumnConstraintSegment):
    """A column constraint, e.g. for CREATE TABLE or ALTER TABLE.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-constraint.html
    """

    match_grammar = Sequence(
        Sequence(
            "CONSTRAINT",
            Ref("ObjectReferenceSegment"),
            optional=True,
        ),
        OneOf(
            Sequence(
                Ref("PrimaryKeyGrammar"),
                Ref("ConstraintOptionGrammar", optional=True),
            ),
            Sequence(
                Ref("ForeignKeyGrammar", optional=True),
                "REFERENCES",
                Ref("TableReferenceSegment"),
                Ref("BracketedColumnReferenceListGrammar", optional=True),
                OneOf(
                    Ref("ForeignKeyOptionGrammar"),
                    Ref("ConstraintOptionGrammar"),
                    optional=True,
                ),
            ),
        ),
    )


class CreateTableUsingStatementSegment(sparksql.CreateTableStatementSegment):
    """A `CREATE TABLE [USING]` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-using.html
    """

    type = "create_table_using_statement"

    match_grammar = Sequence(
        OneOf(
            Sequence(
                Sequence(
                    "CREATE",
                    "OR",
                    optional=True,
                ),
                "REPLACE",
                "TABLE",
            ),
            Sequence(
                "CREATE",
                Ref.keyword("EXTERNAL", optional=True),
                "TABLE",
                Ref("IfNotExistsGrammar", optional=True),
            ),
        ),
        Ref("TableReferenceSegment"),
        Ref("TableSpecificationSegment", optional=True),
        Sequence(
            "USING",
            Ref("DataSourceSegment"),
            optional=True,
        ),
        AnyNumberOf(Ref("TableClausesSegment")),
        Sequence(
            "AS",
            OneOf(
                Ref("SelectStatementSegment"),
                Ref("ValuesClauseSegment"),
            ),
            optional=True,
        ),
    )


class CreateTableCloneStatementSegment(BaseSegment):
    """A create table clone statement.

    https://docs.databricks.com/aws/en/sql/language-manual/delta-clone
    """

    type = "create_table_clone_statement"

    match_grammar = Sequence(
        OneOf(
            Sequence(
                Sequence(
                    "CREATE",
                    "OR",
                    optional=True,
                ),
                "REPLACE",
                "TABLE",
            ),
            Sequence(
                "CREATE",
                "TABLE",
                Ref("IfNotExistsGrammar", optional=True),
            ),
        ),
        Ref("TableReferenceSegment"),
        OneOf("SHALLOW", "DEEP", optional=True),
        "CLONE",
        Ref("TableReferenceSegment"),
        Ref("TablePropertiesGrammar", optional=True),
        Ref("LocationGrammar", optional=True),
    )


class TableSpecificationSegment(BaseSegment):
    """A table specification, e.g. for CREATE TABLE or ALTER TABLE.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-spec.html
    """

    type = "table_specification_segment"

    match_grammar = Bracketed(
        Delimited(
            Sequence(
                Ref("ColumnReferenceSegment"),
                Ref("DatatypeSegment"),
                AnyNumberOf(
                    Ref("ColumnPropertiesSegment"),
                ),
            ),
        ),
    )


class ColumnPropertiesSegment(BaseSegment):
    """Properties for a column in a table specification.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-spec.html
    """

    type = "column_properties_segment"

    match_grammar = OneOf(
        Ref("NotNullGrammar"),
        Ref("ColumnGeneratedGrammar"),
        Sequence(
            "DEFAULT",
            Ref("ColumnConstraintDefaultGrammar"),
        ),
        Ref("CollateGrammar"),
        Ref("CommentGrammar"),
        Ref("ColumnConstraintSegment"),
        Ref("MaskStatementSegment"),
    )


class TableClausesSegment(BaseSegment):
    """Clauses for a table specification.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-table-spec.html
    """

    type = "table_clauses_segment"

    match_grammar = OneOf(
        Ref("PartitionClauseSegment"),
        Ref("TableClusterByClauseSegment"),
        Ref("LocationWithCredentialGrammar"),
        Ref("OptionsGrammar"),
        Ref("CommentGrammar"),
        Ref("TablePropertiesGrammar"),
        Sequence(
            "WITH",
            Ref("RowFilterClauseGrammar"),
        ),
    )


class DeclareOrReplaceVariableStatementSegment(BaseSegment):
    """A `DECLARE [OR REPLACE] VARIABLE` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-declare-variable.html
    """

    type = "declare_or_replace_variable_statement"
    match_grammar = Sequence(
        Ref.keyword("DECLARE"),
        Ref("OrReplaceGrammar", optional=True),
        Ref.keyword("VARIABLE", optional=True),
        Ref("SingleIdentifierGrammar"),  # Variable name
        Ref("DatatypeSegment", optional=True),  # Variable type
        Sequence(
            OneOf("DEFAULT", Ref("EqualsSegment")),
            Ref("ExpressionSegment"),
            optional=True,
        ),
    )


class CommentOnStatementSegment(BaseSegment):
    """`COMMENT ON` statement.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-comment.html
    """

    type = "comment_clause"

    match_grammar = Sequence(
        "COMMENT",
        "ON",
        OneOf(
            Sequence(
                "CATALOG",
                Ref("CatalogReferenceSegment"),
            ),
            Sequence(
                OneOf("DATABASE", "SCHEMA"),
                Ref("DatabaseReferenceSegment"),
            ),
            Sequence(
                "TABLE",
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "VOLUME",
                Ref("VolumeReferenceSegment"),
            ),
            Sequence(
                "COLUMN",
                Ref("ColumnReferenceSegment"),
            ),
            # TODO: Split out individual items if they have references
            Sequence(
                OneOf(
                    "CONNECTION",
                    "PROVIDER",
                    "RECIPIENT",
                    "SHARE",
                ),
                Ref("ObjectReferenceSegment"),
            ),
        ),
        "IS",
        OneOf(Ref("QuotedLiteralSegment"), "NULL"),
    )


class FunctionNameSegment(BaseSegment):
    """Function name, including any prefix bits, e.g. project or schema."""

    type = "function_name"
    match_grammar: Matchable = Sequence(
        # Project name, schema identifier, etc.
        AnyNumberOf(
            Sequence(
                Ref("SingleIdentifierGrammar"),
                Ref("DotSegment"),
            ),
            terminators=[Ref("BracketedSegment")],
        ),
        # Base function name
        Ref("FunctionNameIdentifierSegment", terminators=[Ref("BracketedSegment")]),
        allow_gaps=False,
    )


class MagicCellStatementSegment(BaseSegment):
    """Treat -- MAGIC %md/py/sh/... Cells as their own segments.

    N.B. This is a workaround, to make databricks notebooks
    with leading parsable by sqlfluff.

    https://learn.microsoft.com/en-us/azure/databricks/notebooks/notebooks-code#language-magic
    """

    type = "magic_cell_segment"
    match_grammar = Sequence(
        Ref("NotebookStart", optional=True),
        OneOf(
            Sequence(
                Ref("MagicStartGrammar", optional=True),
                AnyNumberOf(Ref("MagicLineGrammar"), optional=True),
            ),
            Ref("MagicSingleLineGrammar", optional=True),
            # One `bare_magic_cell` token per line (see the lexer subdivider).
            AnyNumberOf(
                OneOf(
                    Ref("BareMagicCellGrammar"),
                    Ref("NotebookStartBareMagicCellGrammar"),
                )
            ),
        ),
        terminators=[Ref("CommandCellSegment", optional=True)],
        reset_terminators=True,
    )


class ParameterizedSegment(BaseSegment):
    """Databricks named parameters to prevent SQL Injection.

    Supports both colon-based (:param) and pipeline (${param}) parameters:
    - Colon syntax:
      https://docs.databricks.com/aws/en/jobs/parameter-use
    - Pipeline syntax:
      https://docs.databricks.com/en/delta-live-tables/parameters.html
    """

    type = "parameterized_expression"
    match_grammar = OneOf(
        # Colon-based parameters: :param_name
        Sequence(
            Ref("ColonSegment"),
            Ref("NakedIdentifierSegment"),
            allow_gaps=False,
        ),
        # Pipeline parameters: ${param_name}
        Ref("PipelineParameterSegment"),
    )


class SetVariableStatementSegment(BaseSegment):
    """A `SET VARIABLE` statement used to set session variables.

    https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-aux-set-variable.html
    """

    type = "set_variable_statement"

    # set var v1=val, v2=val2;
    set_kv_pair = Sequence(
        Delimited(
            Ref("VariableNameIdentifierSegment"),
            Ref("EqualsSegment"),
            OneOf("DEFAULT", OptionallyBracketed(Ref("ExpressionSegment"))),
        )
    )
    # set var (v1,v2) = (values(100,200))
    set_bracketed = Sequence(
        Bracketed(
            Ref("VariableNameIdentifierSegment"),
        ),
        Ref("EqualsSegment"),
        Bracketed(
            OneOf(
                Ref("SelectStatementSegment"),
                Ref("ValuesClauseSegment"),
            )
        ),
    )

    match_grammar = Sequence(
        "SET",
        OneOf(
            "VAR",
            "VARIABLE",
        ),
        OneOf(
            set_kv_pair,
            set_bracketed,
        ),
        allow_gaps=True,
    )


class CDCSpecificationSegment(BaseSegment):
    """The segment shared by APPLY CHANGES INTO and CREATE FLOW...AUTO CDC INTO.

    Used for specifying the data location and rules for ingesting a CDC data source
    """

    type = "cdc_specification_segment"

    match_grammar = Sequence(
        Ref("FromClauseSegment"),
        Sequence(
            "KEYS",
            Indent,
            Ref("BracketedColumnReferenceListGrammar"),
            Dedent,
        ),
        Sequence("IGNORE", "NULL", "UPDATES", optional=True),
        Ref("WhereClauseSegment", optional=True),
        AnyNumberOf(
            Sequence(
                "APPLY",
                "AS",
                OneOf("DELETE", "TRUNCATE"),
                "WHEN",
                Ref("ColumnReferenceSegment"),
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            # NB: Setting max_times to allow for one instance
            #     of DELETE and TRUNCATE at most
            max_times=2,
        ),
        Sequence(
            "SEQUENCE",
            "BY",
            Ref("ColumnReferenceSegment"),
        ),
        Sequence(
            "COLUMNS",
            OneOf(
                Delimited(
                    Ref("ColumnReferenceSegment"),
                ),
                Sequence(
                    Ref("StarSegment"),
                    "EXCEPT",
                    Ref("BracketedColumnReferenceListGrammar"),
                ),
            ),
            optional=True,
        ),
        Sequence(
            "STORED",
            "AS",
            "SCD",
            "TYPE",
            Ref("NumericLiteralSegment"),
            optional=True,
        ),
        Sequence(
            "TRACK",
            "HISTORY",
            "ON",
            OneOf(
                Delimited(
                    Ref("ColumnReferenceSegment"),
                ),
                Sequence(
                    Ref("StarSegment"),
                    "EXCEPT",
                    Ref("BracketedColumnReferenceListGrammar"),
                ),
            ),
            optional=True,
        ),
    )


class ApplyChangesIntoStatementSegment(BaseSegment):
    """A statement to ingest CDC data into a target table.

    https://docs.databricks.com/workflows/delta-live-tables/delta-live-tables-cdc.html#sql
    """

    type = "apply_changes_into_statement"

    match_grammar = Sequence(
        Sequence(
            "APPLY",
            "CHANGES",
            "INTO",
        ),
        Indent,
        Ref("TableExpressionSegment"),
        Dedent,
        Ref("CDCSpecificationSegment"),
    )


class FlowReferenceSegment(ObjectReferenceSegment):
    """A reference to a flow."""

    type = "flow_reference"


class CreateFlowStatementSegment(BaseSegment):
    """A statement for creating a flow to ingest CDC data into a target table.

    https://docs.databricks.com/aws/en/ldp/flows
    https://docs.databricks.com/aws/en/ldp/developer/ldp-sql-ref-apply-changes-into
    """

    type = "create_flow_statement"

    match_grammar = Sequence(
        Sequence(
            "CREATE",
            "FLOW",
        ),
        Ref("FlowReferenceSegment"),
        Sequence(
            "AS",
            "AUTO",
            "CDC",
            "INTO",
        ),
        Indent,
        Ref("TableReferenceSegment"),
        Dedent,
        Ref("CDCSpecificationSegment"),
    )
