"""The Snowflake dialect.

Inherits from ANSI.

Based on https://docs.snowflake.com/en/sql-reference-commands.html
"""

from sqlfluff.core.dialects import load_raw_dialect
from sqlfluff.core.parser import (
    AnyNumberOf,
    AnySetOf,
    Anything,
    BaseSegment,
    Bracketed,
    CodeSegment,
    CommentSegment,
    Dedent,
    Delimited,
    Indent,
    NamedParser,
    OneOf,
    OptionallyBracketed,
    Ref,
    RegexLexer,
    RegexParser,
    SegmentGenerator,
    Sequence,
    StartsWith,
    StringLexer,
    StringParser,
    SymbolSegment,
)
from sqlfluff.dialects.dialect_snowflake_keywords import (
    snowflake_reserved_keywords,
    snowflake_unreserved_keywords,
)

ansi_dialect = load_raw_dialect("ansi")
snowflake_dialect = ansi_dialect.copy_as("snowflake")

snowflake_dialect.patch_lexer_matchers(
    [
        # In snowflake, a double single quote resolves as a single quote in the string.
        # https://docs.snowflake.com/en/sql-reference/data-types-text.html#single-quoted-string-constants
        RegexLexer("single_quote", r"'([^'\\]|\\.|'')*'", CodeSegment),
        RegexLexer(
            "inline_comment",
            r"(--|#|//)[^\n]*",
            CommentSegment,
            segment_kwargs={"trim_start": ("--", "#", "//")},
        ),
    ]
)

snowflake_dialect.insert_lexer_matchers(
    [
        # Keyword assigner needed for keyword functions.
        StringLexer("parameter_assigner", "=>", CodeSegment),
        StringLexer("function_assigner", "->", CodeSegment),
        RegexLexer("stage_path", r"(?:@[^\s;)]+|'@[^']+')", CodeSegment),
        # Column selector
        # https://docs.snowflake.com/en/sql-reference/sql/select.html#parameters
        RegexLexer("column_selector", r"\$[0-9]+", CodeSegment),
        RegexLexer(
            "dollar_quote",
            r"\$\$.*\$\$",
            CodeSegment,
        ),
        RegexLexer(
            "dollar_literal",
            r"[$][a-zA-Z0-9_.]*",
            CodeSegment,
        ),
        RegexLexer("inline_dollar_sign", r"[a-zA-Z_][a-zA-Z0-9_$]*", CodeSegment),
        StringLexer("question_mark", "?", CodeSegment),
        StringLexer("exclude_bracket_open", "{-", CodeSegment),
        StringLexer("exclude_bracket_close", "-}", CodeSegment),
    ],
    before="like_operator",
)

snowflake_dialect.sets("bracket_pairs").add(
    ("exclude", "StartExcludeBracketSegment", "EndExcludeBracketSegment", True)
)

snowflake_dialect.add(
    # In snowflake, these are case sensitive even though they're not quoted
    # so they need a different `name` and `type` so they're not picked up
    # by other rules.
    ParameterAssignerSegment=StringParser(
        "=>", SymbolSegment, name="parameter_assigner", type="parameter_assigner"
    ),
    FunctionAssignerSegment=StringParser(
        "->", SymbolSegment, name="function_assigner", type="function_assigner"
    ),
    NakedSemiStructuredElementSegment=RegexParser(
        r"[A-Z0-9_]*",
        CodeSegment,
        name="naked_semi_structured_element",
        type="semi_structured_element",
    ),
    QuotedSemiStructuredElementSegment=NamedParser(
        "double_quote",
        CodeSegment,
        name="quoted_semi_structured_element",
        type="semi_structured_element",
    ),
    ColumnIndexIdentifierSegment=RegexParser(
        r"\$[0-9]+",
        CodeSegment,
        name="column_index_identifier_segment",
        type="identifier",
    ),
    LocalVariableNameSegment=RegexParser(
        r"[a-zA-Z0-9_]*",
        CodeSegment,
        name="declared_variable",
        type="variable",
    ),
    ReferencedVariableNameSegment=RegexParser(
        r"\$[A-Z][A-Z0-9_]*",
        CodeSegment,
        name="referenced_variable",
        type="variable",
        trim_chars=("$"),
    ),
    # We use a RegexParser instead of keywords as some (those with dashes) require
    # quotes:
    WarehouseSize=RegexParser(
        r"'?XSMALL'?|'?SMALL'?|'?MEDIUM'?|'?LARGE'?|'?XLARGE'?|'?XXLARGE'?|'?X2LARGE'?|"
        r"'?XXXLARGE'?|'?X3LARGE'?|'?X4LARGE'?|'?X5LARGE|'?X6LARGE'?|"
        r"'X-SMALL'|'X-LARGE'|'2X-LARGE'|'3X-LARGE'|'4X-LARGE'|'5X-LARGE'|'6X-LARGE'",
        CodeSegment,
        name="warehouse_size",
        type="warehouse_size",
    ),
    ValidationModeOptionSegment=RegexParser(
        r"'?RETURN_(?:\d+_ROWS|ERRORS|ALL_ERRORS)'?",
        CodeSegment,
        name="validation_mode_option",
        type="validation_mode_option",
    ),
    CopyOptionOnErrorSegment=RegexParser(
        r"'?CONTINUE'?|'?SKIP_FILE(?:_[0-9]+%?)?'?|'?ABORT_STATEMENT'?",
        CodeSegment,
        name="literal",
        type="literal",
    ),
    DoubleQuotedUDFBody=NamedParser(
        "double_quote",
        CodeSegment,
        name="udf_body",
        type="udf_body",
        trim_chars=('"',),
    ),
    SingleQuotedUDFBody=NamedParser(
        "single_quote",
        CodeSegment,
        name="udf_body",
        type="udf_body",
        trim_chars=("'",),
    ),
    DollarQuotedUDFBody=NamedParser(
        "dollar_quote",
        CodeSegment,
        name="udf_body",
        type="udf_body",
        trim_chars=("$",),
    ),
    StagePath=RegexParser(
        r"(?:@[^\s;)]+|'@[^']+')",
        CodeSegment,
        name="stage_path",
        type="identifier",
    ),
    S3Path=RegexParser(
        # https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html
        r"'s3://[a-z0-9][a-z0-9\.-]{1,61}[a-z0-9](?:/.+)?'",
        CodeSegment,
        name="s3_path",
        type="bucket_path",
    ),
    GCSPath=RegexParser(
        # https://cloud.google.com/storage/docs/naming-buckets
        r"'gcs://[a-z0-9][\w\.-]{1,61}[a-z0-9](?:/.+)?'",
        CodeSegment,
        name="gcs_path",
        type="bucket_path",
    ),
    AzureBlobStoragePath=RegexParser(
        # https://docs.microsoft.com/en-us/azure/azure-resource-manager/management/resource-name-rules#microsoftstorage
        r"'azure://[a-z0-9][a-z0-9-]{1,61}[a-z0-9]\.blob\.core\.windows\.net/[a-z0-9]"
        r"[a-z0-9\.-]{1,61}[a-z0-9](?:/.+)?'",
        CodeSegment,
        name="azure_blob_storage_path",
        type="bucket_path",
    ),
    SnowflakeEncryptionOption=RegexParser(
        r"'SNOWFLAKE_FULL'|'SNOWFLAKE_SSE'",
        CodeSegment,
        name="snowflake_encryption_option",
        type="stage_encryption_option",
    ),
    S3EncryptionOption=RegexParser(
        r"'AWS_CSE'|'AWS_SSE_S3'|'AWS_SSE_KMS'",
        CodeSegment,
        name="s3_encryption_option",
        type="stage_encryption_option",
    ),
    GCSEncryptionOption=RegexParser(
        r"'GCS_SSE_KMS'",
        CodeSegment,
        name="gcs_encryption_option",
        type="stage_encryption_option",
    ),
    AzureBlobStorageEncryptionOption=RegexParser(
        r"'AZURE_CSE'",
        CodeSegment,
        name="azure_blob_storage_encryption_option",
        type="stage_encryption_option",
    ),
    FileType=RegexParser(
        r"'?CSV'?|'?JSON'?|'?AVRO'?|'?ORC'?|'?PARQUET'?|'?XML'?",
        CodeSegment,
        name="file_type",
        type="file_type",
    ),
    GroupByContentsGrammar=Delimited(
        OneOf(
            Ref("ColumnReferenceSegment"),
            # Can `GROUP BY 1`
            Ref("NumericLiteralSegment"),
            # Can `GROUP BY coalesce(col, 1)`
            Ref("ExpressionSegment"),
        ),
        terminator=OneOf(
            "ORDER", "LIMIT", "FETCH", "OFFSET", "HAVING", "QUALIFY", "WINDOW"
        ),
    ),
    LimitLiteralGrammar=OneOf(
        Ref("NumericLiteralSegment"),
        "NULL",
        # '' and $$$$ are allowed as alternatives to NULL.
        Ref("QuotedLiteralSegment"),
    ),
    StartExcludeBracketSegment=StringParser(
        "{-", SymbolSegment, name="start_exclude_bracket", type="start_exclude_bracket"
    ),
    EndExcludeBracketSegment=StringParser(
        "-}", SymbolSegment, name="end_exclude_bracket", type="end_exclude_bracket"
    ),
    QuestionMarkSegment=StringParser(
        "?", SymbolSegment, name="question_mark", type="question_mark"
    ),
    CaretSegment=StringParser("^", SymbolSegment, name="caret", type="caret"),
    DollarSegment=StringParser("$", SymbolSegment, name="dollar", type="dollar"),
    PatternQuantifierGrammar=Sequence(
        OneOf(
            Ref("PositiveSegment"),
            Ref("StarSegment"),
            Ref("QuestionMarkSegment"),
            Bracketed(
                OneOf(
                    Ref("NumericLiteralSegment"),
                    Sequence(
                        Ref("NumericLiteralSegment"),
                        Ref("CommaSegment"),
                    ),
                    Sequence(
                        Ref("CommaSegment"),
                        Ref("NumericLiteralSegment"),
                    ),
                    Sequence(
                        Ref("NumericLiteralSegment"),
                        Ref("CommaSegment"),
                        Ref("NumericLiteralSegment"),
                    ),
                ),
                bracket_type="curly",
                bracket_pairs_set="bracket_pairs",
            ),
        ),
        # To put a quantifier into “reluctant mode”.
        Ref("QuestionMarkSegment", optional=True),
        allow_gaps=False,
    ),
    PatternSymbolGrammar=Sequence(
        Ref("SingleIdentifierGrammar"),
        Ref("PatternQuantifierGrammar", optional=True),
        allow_gaps=False,
    ),
    PatternOperatorGrammar=OneOf(
        Ref("PatternSymbolGrammar"),
        Sequence(
            OneOf(
                Bracketed(
                    OneOf(
                        AnyNumberOf(
                            Ref("PatternOperatorGrammar"),
                        ),
                        Delimited(
                            Ref("PatternOperatorGrammar"),
                            delimiter=Ref("BitwiseOrSegment"),
                        ),
                    ),
                    bracket_type="exclude",
                    bracket_pairs_set="bracket_pairs",
                ),
                Bracketed(
                    OneOf(
                        AnyNumberOf(
                            Ref("PatternOperatorGrammar"),
                        ),
                        Delimited(
                            Ref("PatternOperatorGrammar"),
                            delimiter=Ref("BitwiseOrSegment"),
                        ),
                    ),
                ),
                Sequence(
                    "PERMUTE",
                    Bracketed(
                        Delimited(
                            Ref("PatternSymbolGrammar"),
                        ),
                    ),
                ),
            ),
            # Operators can also be followed by a quantifier.
            Ref("PatternQuantifierGrammar", optional=True),
            allow_gaps=False,
        ),
    ),
    PatternGrammar=Sequence(
        # https://docs.snowflake.com/en/sql-reference/constructs/match_recognize.html#pattern-specifying-the-pattern-to-match
        Ref("CaretSegment", optional=True),
        OneOf(
            AnyNumberOf(
                Ref("PatternOperatorGrammar"),
            ),
            Delimited(
                Ref("PatternOperatorGrammar"),
                delimiter=Ref("BitwiseOrSegment"),
            ),
        ),
        Ref("DollarSegment", optional=True),
    ),
)

snowflake_dialect.replace(
    NakedIdentifierSegment=SegmentGenerator(
        # Generate the anti template from the set of reserved keywords
        lambda dialect: RegexParser(
            # See https://docs.snowflake.com/en/sql-reference/identifiers-syntax.html
            r"[a-zA-Z_][a-zA-Z0-9_$]*",
            CodeSegment,
            name="naked_identifier",
            type="identifier",
            anti_template=r"^(" + r"|".join(dialect.sets("reserved_keywords")) + r")$",
        )
    ),
    LiteralGrammar=ansi_dialect.get_grammar("LiteralGrammar").copy(
        insert=[
            Ref("ReferencedVariableNameSegment"),
        ]
    ),
    Accessor_Grammar=AnyNumberOf(
        Ref("ArrayAccessorSegment"),
        # Add in semi structured expressions
        Ref("SemiStructuredAccessorSegment"),
    ),
    PreTableFunctionKeywordsGrammar=OneOf(Ref("LateralKeywordSegment")),
    FunctionContentsExpressionGrammar=OneOf(
        Ref("DatetimeUnitSegment"),
        Ref("NamedParameterExpressionSegment"),
        Ref("ReferencedVariableNameSegment"),
        Sequence(
            Ref("ExpressionSegment"),
            Sequence(OneOf("IGNORE", "RESPECT"), "NULLS", optional=True),
        ),
    ),
    JoinLikeClauseGrammar=Sequence(
        AnySetOf(
            Ref("MatchRecognizeClauseSegment"),
            Ref("ChangesClauseSegment"),
            Ref("ConnectByClauseSegment"),
            Ref("FromAtExpressionSegment"),
            Ref("FromBeforeExpressionSegment"),
            Ref("FromPivotExpressionSegment"),
            Ref("FromUnpivotExpressionSegment"),
            Ref("SamplingExpressionSegment"),
            min_times=1,
        ),
        Ref("AliasExpressionSegment", optional=True),
    ),
    SingleIdentifierGrammar=OneOf(
        Ref("NakedIdentifierSegment"),
        Ref("QuotedIdentifierSegment"),
        Ref("ColumnIndexIdentifierSegment"),
        Ref("ReferencedVariableNameSegment"),
        Ref("StagePath"),
    ),
    PostFunctionGrammar=Sequence(
        Ref("WithinGroupClauseSegment", optional=True),
        Sequence(OneOf("IGNORE", "RESPECT"), "NULLS", optional=True),
        Ref("OverClauseSegment", optional=True),
    ),
    TemporaryGrammar=Sequence(
        OneOf("LOCAL", "GLOBAL", optional=True),
        OneOf("TEMP", "TEMPORARY", optional=True),
        Sequence("VOLATILE", optional=True),
        optional=True,
    ),
    TemporaryTransientGrammar=OneOf(Ref("TemporaryGrammar"), "TRANSIENT"),
    BaseExpressionElementGrammar=OneOf(
        # Allow use of CONNECT_BY_ROOT pseudo-columns.
        # https://docs.snowflake.com/en/sql-reference/constructs/connect-by.html#:~:text=Snowflake%20supports%20the%20CONNECT_BY_ROOT,the%20Examples%20section%20below.
        Sequence("CONNECT_BY_ROOT", Ref("ColumnReferenceSegment")),
        Ref("LiteralGrammar"),
        Ref("BareFunctionSegment"),
        Ref("IntervalExpressionSegment"),
        Ref("FunctionSegment"),
        Ref("ColumnReferenceSegment"),
        Ref("ExpressionSegment"),
    ),
    QuotedLiteralSegment=OneOf(
        # https://docs.snowflake.com/en/sql-reference/data-types-text.html#string-constants
        NamedParser(
            "single_quote",
            CodeSegment,
            name="quoted_literal",
            type="literal",
        ),
        NamedParser(
            "dollar_quote",
            CodeSegment,
            name="quoted_literal",
            type="literal",
        ),
    ),
    LikeGrammar=OneOf(
        # https://docs.snowflake.com/en/sql-reference/functions/like.html
        Sequence("LIKE", OneOf("ALL", "ANY", optional=True)),
        "RLIKE",
        Sequence("ILIKE", Ref.keyword("ANY", optional=True)),
        "REGEXP",
    ),
    SelectClauseElementTerminatorGrammar=OneOf(
        "FROM",
        "WHERE",
        Sequence("ORDER", "BY"),
        "LIMIT",
        "FETCH",
        "OFFSET",
        Ref("CommaSegment"),
        Ref("SetOperatorSegment"),
    ),
    FromClauseTerminatorGrammar=OneOf(
        "WHERE",
        "LIMIT",
        "FETCH",
        "OFFSET",
        Sequence("GROUP", "BY"),
        Sequence("ORDER", "BY"),
        "HAVING",
        "QUALIFY",
        "WINDOW",
        Ref("SetOperatorSegment"),
        Ref("WithNoSchemaBindingClauseSegment"),
        Ref("WithDataClauseSegment"),
    ),
    WhereClauseTerminatorGrammar=OneOf(
        "LIMIT",
        "FETCH",
        "OFFSET",
        Sequence("GROUP", "BY"),
        Sequence("ORDER", "BY"),
        "HAVING",
        "QUALIFY",
        "WINDOW",
        "OVERLAPS",
    ),
)

# Add all Snowflake keywords
snowflake_dialect.sets("unreserved_keywords").clear()
snowflake_dialect.sets("unreserved_keywords").update(
    [n.strip().upper() for n in snowflake_unreserved_keywords.split("\n")]
)

snowflake_dialect.sets("reserved_keywords").clear()
snowflake_dialect.sets("reserved_keywords").update(
    [n.strip().upper() for n in snowflake_reserved_keywords.split("\n")]
)

# Add datetime units and their aliases from
# https://docs.snowflake.com/en/sql-reference/functions-date-time.html#label-supported-date-time-parts
snowflake_dialect.sets("datetime_units").clear()
snowflake_dialect.sets("datetime_units").update(
    [
        "YEAR",
        "Y",
        "YY",
        "YYY",
        "YYYY",
        "YR",
        "YEARS",
        "YRS",
        "MONTH",
        "MM",
        "MON",
        "MONS",
        "MONTHS",
        "DAY",
        "D",
        "DD",
        "DAYS",
        "DAYOFMONTH",
        "DAYOFWEEK",
        "WEEKDAY",
        "DOW",
        "DW",
        "DAYOFWEEKISO",
        "WEEKDAY_ISO",
        "DOW_ISO",
        "DW_ISO",
        "DAYOFYEAR",
        "YEARDAY",
        "DOY",
        "DY",
        "WEEK",
        "W",
        "WK",
        "WEEKOFYEAR",
        "WOY",
        "WY",
        "WEEKISO",
        "WEEK_ISO",
        "WEEKOFYEARISO",
        "WEEKOFYEAR_ISO",
        "QUARTER",
        "Q",
        "QTR",
        "QTRS",
        "QUARTERS",
        "YEAROFWEEK",
        "YEAROFWEEKISO",
        "HOUR",
        "H",
        "HH",
        "HR",
        "HOURS",
        "HRS",
        "MINUTE",
        "M",
        "MI",
        "MIN",
        "MINUTES",
        "MINS",
        "SECOND",
        "S",
        "SEC",
        "SECONDS",
        "SECS",
        "MILLISECOND",
        "MS",
        "MSEC",
        "MILLISECONDS",
        "MICROSECOND",
        "US",
        "USEC",
        "MICROSECONDS",
        "NANOSECOND",
        "NS",
        "NSEC",
        "NANOSEC",
        "NSECOND",
        "NANOSECONDS",
        "NANOSECS",
        "NSECONDS",
        "EPOCH_SECOND",
        "EPOCH",
        "EPOCH_SECONDS",
        "EPOCH_MILLISECOND",
        "EPOCH_MILLISECONDS",
        "EPOCH_MICROSECOND",
        "EPOCH_MICROSECONDS",
        "EPOCH_NANOSECOND",
        "EPOCH_NANOSECONDS",
        "TIMEZONE_HOUR",
        "TZH",
        "TIMEZONE_MINUTE",
        "TZM",
    ]
)


@snowflake_dialect.segment()
class ConnectByClauseSegment(BaseSegment):
    """A `CONNECT BY` clause.

    https://docs.snowflake.com/en/sql-reference/constructs/connect-by.html
    """

    type = "connectby_clause"
    match_grammar = Sequence(
        "START",
        "WITH",
        Ref("ExpressionSegment"),
        "CONNECT",
        "BY",
        Delimited(
            Sequence(
                Ref.keyword("PRIOR", optional=True),
                Ref("ColumnReferenceSegment"),
                Ref("EqualsSegment"),
                Ref.keyword("PRIOR", optional=True),
                Ref("ColumnReferenceSegment"),
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class GroupByClauseSegment(BaseSegment):
    """A `GROUP BY` clause like in `SELECT`.

    Snowflake supports Cube, Rollup, and Grouping Sets

    https://docs.snowflake.com/en/sql-reference/constructs/group-by.html
    """

    type = "groupby_clause"
    match_grammar = StartsWith(
        Sequence("GROUP", "BY"),
        terminator=OneOf(
            "ORDER", "LIMIT", "FETCH", "OFFSET", "HAVING", "QUALIFY", "WINDOW"
        ),
        enforce_whitespace_preceding_terminator=True,
    )
    parse_grammar = Sequence(
        "GROUP",
        "BY",
        Indent,
        OneOf(
            Sequence(
                OneOf("CUBE", "ROLLUP", Sequence("GROUPING", "SETS")),
                Bracketed(
                    Ref("GroupByContentsGrammar"),
                ),
            ),
            Ref("GroupByContentsGrammar"),
        ),
        Dedent,
    )


@snowflake_dialect.segment(replace=True)
class ValuesClauseSegment(BaseSegment):
    """A `VALUES` clause like in `INSERT`."""

    type = "values_clause"
    match_grammar = Sequence(
        "VALUES",
        Delimited(
            Bracketed(
                Delimited(
                    # DEFAULT and NULL keywords used in
                    # INSERT INTO statement.
                    "DEFAULT",
                    "NULL",
                    Ref("ExpressionSegment"),
                    ephemeral_name="ValuesClauseElements",
                )
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class InsertStatementSegment(BaseSegment):
    """An `INSERT` statement.

    https://docs.snowflake.com/en/sql-reference/sql/insert.html
    https://docs.snowflake.com/en/sql-reference/sql/insert-multi-table.html
    """

    type = "insert_statement"
    match_grammar = Sequence(
        "INSERT",
        Ref.keyword("OVERWRITE", optional=True),
        OneOf(
            # Single table INSERT INTO.
            Sequence(
                "INTO",
                Ref("TableReferenceSegment"),
                Ref("BracketedColumnReferenceListGrammar", optional=True),
                Ref("SelectableGrammar"),
            ),
            # Unconditional multi-table INSERT INTO.
            Sequence(
                "ALL",
                AnyNumberOf(
                    Sequence(
                        "INTO",
                        Ref("TableReferenceSegment"),
                        Ref("BracketedColumnReferenceListGrammar", optional=True),
                        Ref("ValuesClauseSegment", optional=True),
                    ),
                    min_times=1,
                ),
                Ref("SelectStatementSegment"),
            ),
            # Conditional multi-table INSERT INTO.
            Sequence(
                OneOf(
                    "FIRST",
                    "ALL",
                ),
                AnyNumberOf(
                    Sequence(
                        "WHEN",
                        Ref("ExpressionSegment"),
                        "THEN",
                        AnyNumberOf(
                            Sequence(
                                "INTO",
                                Ref("TableReferenceSegment"),
                                Ref(
                                    "BracketedColumnReferenceListGrammar", optional=True
                                ),
                                Ref("ValuesClauseSegment", optional=True),
                            ),
                            min_times=1,
                        ),
                    ),
                    min_times=1,
                ),
                Sequence(
                    "ELSE",
                    "INTO",
                    Ref("TableReferenceSegment"),
                    Ref("BracketedColumnReferenceListGrammar", optional=True),
                    Ref("ValuesClauseSegment", optional=True),
                    optional=True,
                ),
                Ref("SelectStatementSegment"),
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class FunctionDefinitionGrammar(BaseSegment):
    """This is the body of a `CREATE FUNCTION AS` statement."""

    type = "function_definition"
    match_grammar = Sequence(
        "AS",
        Ref("QuotedLiteralSegment"),
        Sequence(
            "LANGUAGE",
            # Not really a parameter, but best fit for now.
            Ref("ParameterNameSegment"),
            optional=True,
        ),
    )


@snowflake_dialect.segment(replace=True)
class StatementSegment(ansi_dialect.get_segment("StatementSegment")):  # type: ignore
    """A generic segment, to any of its child subsegments."""

    parse_grammar = ansi_dialect.get_segment("StatementSegment").parse_grammar.copy(
        insert=[
            Ref("CreateStatementSegment"),
            Ref("CreateTaskSegment"),
            Ref("CreateCloneStatementSegment"),
            Ref("CreateProcedureStatementSegment"),
            Ref("ShowStatementSegment"),
            Ref("AlterUserSegment"),
            Ref("AlterSessionStatementSegment"),
            Ref("AlterTaskStatementSegment"),
            Ref("SetAssignmentStatementSegment"),
            Ref("CallStoredProcedureSegment"),
            Ref("MergeStatementSegment"),
            Ref("CopyIntoTableStatementSegment"),
            Ref("AlterWarehouseStatementSegment"),
            Ref("CreateExternalTableSegment"),
            Ref("CreateSchemaStatementSegment"),
            Ref("AlterSchemaStatementSegment"),
            Ref("CreateFunctionStatementSegment"),
            Ref("AlterFunctionStatementSegment"),
            Ref("CreateStageSegment"),
            Ref("AlterStageSegment"),
            Ref("UnsetStatementSegment"),
            Ref("UndropStatementSegment"),
            Ref("CommentStatementSegment"),
            Ref("CallStatementSegment"),
            Ref("AlterViewStatementSegment"),
        ],
        remove=[
            Ref("CreateTypeStatementSegment"),
            Ref("CreateExtensionStatementSegment"),
            Ref("CreateIndexStatementSegment"),
            Ref("DropIndexStatementSegment"),
        ],
    )


@snowflake_dialect.segment()
class SetAssignmentStatementSegment(BaseSegment):
    """A `SET` statement.

    https://docs.snowflake.com/en/sql-reference/sql/set.html
    """

    type = "set_statement"

    match_grammar = OneOf(
        Sequence(
            "SET",
            Ref("LocalVariableNameSegment"),
            Ref("EqualsSegment"),
            Ref("ExpressionSegment"),
        ),
        Sequence(
            "SET",
            Bracketed(
                Delimited(
                    Ref("LocalVariableNameSegment"), delimiter=Ref("CommaSegment")
                )
            ),
            Ref("EqualsSegment"),
            Bracketed(
                Delimited(
                    Ref("ExpressionSegment"),
                    delimiter=Ref("CommaSegment"),
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class CallStoredProcedureSegment(BaseSegment):
    """This is a CALL statement used to execute a stored procedure.

    https://docs.snowflake.com/en/sql-reference/sql/call.html
    """

    type = "call_segment"

    match_grammar = Sequence(
        "CALL",
        Ref("FunctionSegment"),
    )


@snowflake_dialect.segment()
class WithinGroupClauseSegment(BaseSegment):
    """An WITHIN GROUP clause for window functions.

    https://docs.snowflake.com/en/sql-reference/functions/listagg.html.
    https://docs.snowflake.com/en/sql-reference/functions/array_agg.html.
    """

    type = "withingroup_clause"
    match_grammar = Sequence(
        "WITHIN",
        "GROUP",
        Bracketed(Anything(optional=True)),
    )

    parse_grammar = Sequence(
        "WITHIN",
        "GROUP",
        Bracketed(Ref("OrderByClauseSegment", optional=True)),
    )


@snowflake_dialect.segment()
class MatchRecognizeClauseSegment(BaseSegment):
    """A `MATCH_RECOGNIZE` clause.

    https://docs.snowflake.com/en/sql-reference/constructs/match_recognize.html
    """

    type = "match_recognize_clause"
    match_grammar = Sequence(
        "MATCH_RECOGNIZE",
        Bracketed(
            Ref("PartitionClauseSegment", optional=True),
            Ref("OrderByClauseSegment", optional=True),
            Sequence(
                "MEASURES",
                Delimited(
                    Sequence(
                        # The edges of the window frame can be specified
                        # by using either RUNNING or FINAL semantics.
                        # https://docs.snowflake.com/en/sql-reference/constructs/match_recognize.html#expressions-in-define-and-measures-clauses
                        OneOf(
                            "FINAL",
                            "RUNNING",
                            optional=True,
                        ),
                        Ref("ExpressionSegment"),
                        Ref("AliasExpressionSegment"),
                    ),
                ),
                optional=True,
            ),
            OneOf(
                Sequence(
                    "ONE",
                    "ROW",
                    "PER",
                    "MATCH",
                ),
                Sequence(
                    "ALL",
                    "ROWS",
                    "PER",
                    "MATCH",
                    OneOf(
                        Sequence(
                            "SHOW",
                            "EMPTY",
                            "MATCHES",
                        ),
                        Sequence(
                            "OMIT",
                            "EMPTY",
                            "MATCHES",
                        ),
                        Sequence(
                            "WITH",
                            "UNMATCHED",
                            "ROWS",
                        ),
                        optional=True,
                    ),
                ),
                optional=True,
            ),
            Sequence(
                "AFTER",
                "MATCH",
                "SKIP",
                OneOf(
                    Sequence(
                        "PAST",
                        "LAST",
                        "ROW",
                    ),
                    Sequence(
                        "TO",
                        "NEXT",
                        "ROW",
                    ),
                    Sequence(
                        "TO",
                        OneOf("FIRST", "LAST", optional=True),
                        Ref("SingleIdentifierGrammar"),
                    ),
                ),
                optional=True,
            ),
            "PATTERN",
            Bracketed(
                Ref("PatternGrammar"),
            ),
            "DEFINE",
            Delimited(
                Sequence(
                    Ref("SingleIdentifierGrammar"),
                    "AS",
                    Ref("ExpressionSegment"),
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class ChangesClauseSegment(BaseSegment):
    """A `CHANGES` clause.

    https://docs.snowflake.com/en/sql-reference/constructs/changes.html
    """

    type = "changes_clause"
    match_grammar = Sequence(
        "CHANGES",
        Bracketed(
            "INFORMATION",
            Ref("ParameterAssignerSegment"),
            OneOf("DEFAULT", "APPEND_ONLY"),
        ),
        OneOf(
            Sequence(
                "AT",
                Bracketed(
                    OneOf("TIMESTAMP", "OFFSET", "STATEMENT"),
                    Ref("ParameterAssignerSegment"),
                    Ref("ExpressionSegment"),
                ),
            ),
            Sequence(
                "BEFORE",
                Bracketed(
                    "STATEMENT",
                    Ref("ParameterAssignerSegment"),
                    Ref("ExpressionSegment"),
                ),
            ),
        ),
        Sequence(
            "END",
            Bracketed(
                OneOf("TIMESTAMP", "OFFSET", "STATEMENT"),
                Ref("ParameterAssignerSegment"),
                Ref("ExpressionSegment"),
            ),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class FromAtExpressionSegment(BaseSegment):
    """An AT expression."""

    type = "from_at_expression"
    match_grammar = Sequence("AT", Bracketed(Anything()))

    parse_grammar = Sequence(
        "AT",
        Bracketed(
            OneOf("TIMESTAMP", "OFFSET", "STATEMENT"),
            Ref("ParameterAssignerSegment"),
            Ref("ExpressionSegment"),
        ),
    )


@snowflake_dialect.segment()
class FromBeforeExpressionSegment(BaseSegment):
    """A BEFORE expression."""

    type = "from_before_expression"
    match_grammar = Sequence("BEFORE", Bracketed(Anything()))

    parse_grammar = Sequence(
        "BEFORE",
        Bracketed(
            OneOf("TIMESTAMP", "OFFSET", "STATEMENT"),
            Ref("ParameterAssignerSegment"),
            Ref("ExpressionSegment"),
        ),
    )


@snowflake_dialect.segment()
class FromPivotExpressionSegment(BaseSegment):
    """A PIVOT expression."""

    type = "from_pivot_expression"
    match_grammar = Sequence(
        "PIVOT",
        Bracketed(
            Ref("FunctionSegment"),
            "FOR",
            Ref("SingleIdentifierGrammar"),
            "IN",
            Bracketed(Delimited(Ref("LiteralGrammar"), delimiter=Ref("CommaSegment"))),
        ),
    )


@snowflake_dialect.segment()
class FromUnpivotExpressionSegment(BaseSegment):
    """An UNPIVOT expression."""

    type = "from_unpivot_expression"
    match_grammar = Sequence(
        "UNPIVOT",
        Bracketed(
            Ref("SingleIdentifierGrammar"),
            "FOR",
            Ref("SingleIdentifierGrammar"),
            "IN",
            Bracketed(
                Delimited(Ref("SingleIdentifierGrammar"), delimiter=Ref("CommaSegment"))
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class SamplingExpressionSegment(BaseSegment):
    """A sampling expression."""

    type = "sample_expression"
    match_grammar = Sequence(
        OneOf("SAMPLE", "TABLESAMPLE"),
        OneOf("BERNOULLI", "ROW", "SYSTEM", "BLOCK", optional=True),
        Bracketed(Ref("NumericLiteralSegment"), Ref.keyword("ROWS", optional=True)),
        Sequence(
            OneOf("REPEATABLE", "SEED"),
            Bracketed(Ref("NumericLiteralSegment")),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class NamedParameterExpressionSegment(BaseSegment):
    """A keyword expression.

    e.g. 'input => custom_fields'

    """

    type = "snowflake_keyword_expression"
    match_grammar = Sequence(
        Ref("ParameterNameSegment"),
        Ref("ParameterAssignerSegment"),
        OneOf(
            Ref("LiteralGrammar"),
            Ref("ColumnReferenceSegment"),
            Ref("ExpressionSegment"),
        ),
    )


@snowflake_dialect.segment()
class SemiStructuredAccessorSegment(BaseSegment):
    """A semi-structured data accessor segment.

    https://docs.snowflake.com/en/user-guide/semistructured-considerations.html
    """

    type = "snowflake_semi_structured_expression"
    match_grammar = Sequence(
        OneOf(
            # If a field is already a VARIANT, this could
            # be initiated by a colon or a dot. This is particularly
            # useful when a field is an ARRAY of objects.
            Ref("DotSegment"),
            Ref("ColonSegment"),
        ),
        OneOf(
            Ref("NakedSemiStructuredElementSegment"),
            Ref("QuotedSemiStructuredElementSegment"),
        ),
        Ref("ArrayAccessorSegment", optional=True),
        AnyNumberOf(
            Sequence(
                OneOf(
                    # Can be delimited by dots or colons
                    Ref("DotSegment"),
                    Ref("ColonSegment"),
                ),
                OneOf(
                    Ref("NakedSemiStructuredElementSegment"),
                    Ref("QuotedSemiStructuredElementSegment"),
                ),
                allow_gaps=True,
            ),
            Ref("ArrayAccessorSegment", optional=True),
            allow_gaps=True,
        ),
        allow_gaps=True,
    )


@snowflake_dialect.segment()
class QualifyClauseSegment(BaseSegment):
    """A `QUALIFY` clause like in `SELECT`.

    https://docs.snowflake.com/en/sql-reference/constructs/qualify.html
    """

    type = "qualify_clause"
    match_grammar = StartsWith(
        "QUALIFY",
        terminator=OneOf(
            Sequence("ORDER", "BY"),
            "LIMIT",
            "FETCH",
            "OFFSET",
        ),
    )
    parse_grammar = Sequence(
        "QUALIFY",
        Indent,
        OneOf(
            Bracketed(
                Ref("ExpressionSegment"),
            ),
            Ref("ExpressionSegment"),
        ),
        Dedent,
    )


@snowflake_dialect.segment(replace=True)
class SelectStatementSegment(
    ansi_dialect.get_segment("SelectStatementSegment")  # type: ignore
):
    """A snowflake `SELECT` statement including optional Qualify.

    https://docs.snowflake.com/en/sql-reference/constructs/qualify.html
    """

    type = "select_statement"
    match_grammar = StartsWith(
        # NB: In bigquery, the select clause may include an EXCEPT, which
        # will also match the set operator, but by starting with the whole
        # select clause rather than just the SELECT keyword, we normally
        # mitigate that here. But this isn't BigQuery! So we can be more
        # efficient and just just the keyword.
        "SELECT",
        terminator=Ref("SetOperatorSegment"),
    )

    parse_grammar = ansi_dialect.get_segment(
        "SelectStatementSegment"
    ).parse_grammar.copy(
        insert=[Ref("QualifyClauseSegment", optional=True)],
        before=Ref("OrderByClauseSegment", optional=True),
    )


@snowflake_dialect.segment(replace=True)
class SelectClauseModifierSegment(BaseSegment):
    """Things that come after SELECT but before the columns, specifically for Snowflake.

    https://docs.snowflake.com/en/sql-reference/constructs.html
    """

    type = "select_clause_modifier"
    match_grammar = Sequence(
        OneOf("DISTINCT", "ALL", optional=True),
        # TOP N is unique to Snowflake, and we can optionally add DISTINCT/ALL in front
        # of it.
        Sequence("TOP", Ref("NumericLiteralSegment"), optional=True),
    )


@snowflake_dialect.segment(replace=True)
class AlterTableStatementSegment(BaseSegment):
    """An `ALTER TABLE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-table.html
    If possible, please keep the order below the same as Snowflake's doc:
    """

    type = "alter_table_statement"

    match_grammar = Sequence(
        "ALTER",
        "TABLE",
        Ref("TableReferenceSegment"),
        OneOf(
            # Rename
            Sequence(
                "RENAME",
                "TO",
                Ref("TableReferenceSegment"),
            ),
            # Swap With
            Sequence(
                "SWAP",
                "WITH",
                Ref("TableReferenceSegment"),
            ),
            # searchOptimizationAction
            # N.B. Since SEARCH and OPTIMIZATION are unreserved keywords
            # we move this above AlterTableTableColumnActionSegment
            # in order to avoid matching these as columns.
            Sequence(
                OneOf(
                    "ADD",
                    "DROP",
                ),
                "SEARCH",
                "OPTIMIZATION",
            ),
            Ref("AlterTableClusteringActionSegment"),
            Ref("AlterTableTableColumnActionSegment"),
            # @TODO: constraintAction
            # @TODO: extTableColumnAction
            # SET Table options
            # @TODO: Restrict the list of parameters supported per Snowflake doc.
            Sequence(
                Ref.keyword("SET"),
                OneOf(
                    Ref("ParameterNameSegment"),
                    Ref.keyword("COMMENT"),
                ),
                Ref("EqualsSegment", optional=True),
                OneOf(
                    Ref("LiteralGrammar"),
                    Ref("NakedIdentifierSegment"),
                    Ref("QuotedLiteralSegment"),
                ),
            ),
            # @TODO: Set/unset TAG
            # @TODO: Unset table options
            # @TODO: Add/drop row access policies
        ),
    )


@snowflake_dialect.segment()
class AlterTableTableColumnActionSegment(BaseSegment):
    """ALTER TABLE `tableColumnAction` per defined in Snowflake's grammar.

    https://docs.snowflake.com/en/sql-reference/sql/alter-table.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-table-column.html

    If possible, please match the order of this sequence with what's defined in
    Snowflake's tableColumnAction grammar.
    """

    type = "alter_table_table_column_action"

    match_grammar = OneOf(
        # Add Column
        Sequence(
            "ADD",
            "COLUMN",
            # Handle Multiple Columns
            Delimited(
                Sequence(
                    Ref("ColumnReferenceSegment"),
                    Ref("DatatypeSegment"),
                    OneOf(
                        # Default
                        Sequence(
                            "DEFAULT",
                            Ref("ExpressionSegment"),
                        ),
                        # Auto-increment/identity column
                        Sequence(
                            OneOf(
                                "AUTOINCREMENT",
                                "IDENTITY",
                            ),
                            OneOf(
                                # ( <start_num>, <step_num> )
                                Bracketed(
                                    Ref("NumericLiteralSegment"),
                                    Ref("CommaSegment"),
                                    Ref("NumericLiteralSegment"),
                                ),
                                # START <num> INCREMENT <num>
                                Sequence(
                                    "START",
                                    Ref("NumericLiteralSegment"),
                                    "INCREMENT",
                                    Ref("NumericLiteralSegment"),
                                ),
                                optional=True,
                            ),
                        ),
                        optional=True,
                    ),
                    # @TODO: Add support for `inlineConstraint`
                    Sequence(
                        Ref.keyword("WITH", optional=True),
                        "MASKING",
                        "POLICY",
                        Ref("FunctionNameSegment"),
                        Sequence(
                            "USING",
                            Bracketed(
                                Delimited(
                                    OneOf(
                                        Ref("ColumnReferenceSegment"),
                                        Ref("ExpressionSegment"),
                                    )
                                ),
                            ),
                            optional=True,
                        ),
                        optional=True,
                    ),
                    Ref("CommentClauseSegment", optional=True),
                ),
            ),
        ),
        # Rename column
        Sequence(
            "RENAME",
            "COLUMN",
            Ref("ColumnReferenceSegment"),
            "TO",
            Ref("ColumnReferenceSegment"),
        ),
        # Alter/Modify column(s)
        Sequence(
            OneOf("ALTER", "MODIFY"),
            OptionallyBracketed(
                Delimited(
                    OneOf(
                        # Add things
                        Sequence(
                            Ref.keyword("COLUMN", optional=True),
                            Ref("ColumnReferenceSegment"),
                            OneOf(
                                Sequence("DROP", "DEFAULT"),
                                Sequence(
                                    "SET",
                                    "DEFAULT",
                                    Ref("NakedIdentifierSegment"),
                                    Ref("DotSegment"),
                                    "NEXTVAL",
                                ),
                                Sequence(
                                    OneOf("SET", "DROP", optional=True),
                                    "NOT",
                                    "NULL",
                                ),
                                Sequence(
                                    Sequence(
                                        Sequence("SET", "DATA", optional=True),
                                        "TYPE",
                                        optional=True,
                                    ),
                                    Ref("DatatypeSegment"),
                                ),
                                Ref("CommentClauseSegment"),
                            ),
                        ),
                        Sequence(
                            "COLUMN",
                            Ref("ColumnReferenceSegment"),
                            "SET",
                            "MASKING",
                            "POLICY",
                            Ref("FunctionNameSegment"),
                            Sequence(
                                "USING",
                                Bracketed(
                                    Delimited(
                                        OneOf(
                                            Ref("ColumnReferenceSegment"),
                                            Ref("ExpressionSegment"),
                                        )
                                    ),
                                ),
                                optional=True,
                            ),
                        ),
                        Sequence(
                            "COLUMN",
                            Ref("ColumnReferenceSegment"),
                            "UNSET",
                            "MASKING",
                            "POLICY",
                        ),
                        # @TODO: Set/Unset TAG support
                    ),
                ),
            ),
        ),
        # Drop column
        Sequence(
            "DROP",
            Ref.keyword("COLUMN", optional=True),
            Delimited(Ref("ColumnReferenceSegment")),
        ),
        # @TODO: Drop columns
        # vvvvv COPIED FROM ANSI vvvvv
        # @TODO: Removed these once `tableColumnAction` is properly supported.
        Sequence(
            OneOf("ADD", "MODIFY"),
            Ref.keyword("COLUMN", optional=True),
            Ref("ColumnDefinitionSegment"),
            OneOf(
                Sequence(OneOf("FIRST", "AFTER"), Ref("ColumnReferenceSegment")),
                # Bracketed Version of the same
                Ref("BracketedColumnReferenceListGrammar"),
                optional=True,
            ),
        ),
    )


@snowflake_dialect.segment()
class AlterTableClusteringActionSegment(BaseSegment):
    """ALTER TABLE `clusteringAction` per defined in Snowflake's grammar.

    https://docs.snowflake.com/en/sql-reference/sql/alter-table.html#clustering-actions-clusteringaction
    """

    type = "alter_table_clustering_action"

    match_grammar = OneOf(
        Sequence(
            "CLUSTER",
            "BY",
            Bracketed(
                Delimited(Ref("ExpressionSegment")),
            ),
        ),
        # N.B. RECLUSTER is deprecated:
        # https://docs.snowflake.com/en/user-guide/tables-clustering-manual.html
        Sequence(
            "RECLUSTER",
            Sequence(
                "MAX_SIZE",
                Ref("EqualsSegment"),
                Ref("NumericLiteralSegment"),
                optional=True,
            ),
            Ref("WhereClauseSegment", optional=True),
        ),
        Sequence(
            OneOf(
                "SUSPEND",
                "RESUME",
            ),
            "RECLUSTER",
        ),
        Sequence(
            "DROP",
            "CLUSTERING",
            "KEY",
        ),
    )


@snowflake_dialect.segment()
class AlterWarehouseStatementSegment(BaseSegment):
    """An `ALTER WAREHOUSE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-warehouse.html

    """

    type = "alter_warehouse_statement"
    match_grammar = Sequence(
        "ALTER",
        "WAREHOUSE",
        Sequence("IF", "EXISTS", optional=True),
        OneOf(
            Sequence(
                Ref("NakedIdentifierSegment", optional=True),
                OneOf(
                    "SUSPEND",
                    Sequence(
                        "RESUME",
                        Sequence("IF", "SUSPENDED", optional=True),
                    ),
                ),
            ),
            Sequence(
                Ref("NakedIdentifierSegment", optional=True),
                Sequence(
                    "ABORT",
                    "ALL",
                    "QUERIES",
                ),
            ),
            Sequence(
                Ref("NakedIdentifierSegment"),
                "RENAME",
                "TO",
                Ref("NakedIdentifierSegment"),
            ),
            Sequence(
                Ref("NakedIdentifierSegment"),
                "SET",
                OneOf(
                    AnyNumberOf(
                        Ref("WarehouseObjectPropertiesSegment"),
                        Ref("CommentEqualsClauseSegment"),
                        Ref("WarehouseObjectParamsSegment"),
                    ),
                    Ref("TagEqualsSegment"),
                ),
            ),
            Sequence(
                Ref("NakedIdentifierSegment"),
                "UNSET",
                OneOf(
                    Delimited(Ref("NakedIdentifierSegment")),
                    Sequence("TAG", Delimited(Ref("NakedIdentifierSegment"))),
                ),
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class CommentClauseSegment(BaseSegment):
    """A comment clause.

    e.g. COMMENT 'column description'
    """

    type = "comment_clause"
    match_grammar = Sequence("COMMENT", Ref("QuotedLiteralSegment"))


@snowflake_dialect.segment()
class CommentEqualsClauseSegment(BaseSegment):
    """A comment clause.

    e.g. COMMENT = 'view/table description'
    """

    type = "comment_equals_clause"
    match_grammar = Sequence(
        "COMMENT", Ref("EqualsSegment"), Ref("QuotedLiteralSegment")
    )


@snowflake_dialect.segment()
class TagBracketedEqualsSegment(BaseSegment):
    """A tag clause.

    e.g. TAG (tag1 = 'value1', tag2 = 'value2')
    """

    type = "tag_bracketed_equals"
    match_grammar = Sequence(
        Sequence("WITH", optional=True),
        "TAG",
        Bracketed(
            Delimited(
                Sequence(
                    Ref("NakedIdentifierSegment"),
                    Ref("EqualsSegment"),
                    Ref("QuotedLiteralSegment"),
                )
            ),
        ),
    )


@snowflake_dialect.segment()
class TagEqualsSegment(BaseSegment):
    """A tag clause.

    e.g. TAG tag1 = 'value1', tag2 = 'value2'
    """

    type = "tag_equals"
    match_grammar = Sequence(
        "TAG",
        Delimited(
            Sequence(
                Ref("NakedIdentifierSegment"),
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            )
        ),
    )


@snowflake_dialect.segment(replace=True)
class UnorderedSelectStatementSegment(
    ansi_dialect.get_segment("SelectStatementSegment")  # type: ignore
):
    """A snowflake unordered `SELECT` statement including optional Qualify.

    https://docs.snowflake.com/en/sql-reference/constructs/qualify.html
    """

    type = "select_statement"
    match_grammar = ansi_dialect.get_segment(
        "UnorderedSelectStatementSegment"
    ).match_grammar.copy()

    parse_grammar = ansi_dialect.get_segment(
        "UnorderedSelectStatementSegment"
    ).parse_grammar.copy(
        insert=[Ref("QualifyClauseSegment", optional=True)],
        before=Ref("OverlapsClauseSegment", optional=True),
    )


@snowflake_dialect.segment()
class CreateCloneStatementSegment(BaseSegment):
    """A snowflake `CREATE ... CLONE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-clone.html
    """

    type = "create_clone_statement"
    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        OneOf(
            "DATABASE",
            "SCHEMA",
            "TABLE",
            "SEQUENCE",
            Sequence("FILE", "FORMAT"),
            "STAGE",
            "STREAM",
            "TASK",
        ),
        Sequence("IF", "NOT", "EXISTS", optional=True),
        Ref("SingleIdentifierGrammar"),
        "CLONE",
        Ref("SingleIdentifierGrammar"),
        OneOf(
            Ref("FromAtExpressionSegment"),
            Ref("FromBeforeExpressionSegment"),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class CreateProcedureStatementSegment(BaseSegment):
    """A snowflake `CREATE ... PROCEDURE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-procedure.html
    """

    type = "create_procedure_statement"
    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        "PROCEDURE",
        Ref("FunctionNameSegment"),
        Ref("FunctionParameterListGrammar"),
        "RETURNS",
        Ref("DatatypeSegment"),
        Sequence("NOT", "NULL", optional=True),
        "LANGUAGE",
        "JAVASCRIPT",
        OneOf(
            Sequence("CALLED", "ON", "NULL", "INPUT"),
            Sequence("RETURNS", "NULL", "ON", "NULL", "INPUT"),
            "STRICT",
            optional=True,
        ),
        OneOf("VOLATILE", "IMMUTABLE", optional=True),
        Ref("CommentEqualsClauseSegment", optional=True),
        Sequence("EXECUTE", "AS", OneOf("CALLER", "OWNER"), optional=True),
        "AS",
        OneOf(
            Ref("DoubleQuotedUDFBody"),
            Ref("SingleQuotedUDFBody"),
            Ref("DollarQuotedUDFBody"),
        ),
    )


@snowflake_dialect.segment(replace=True)
class CreateFunctionStatementSegment(BaseSegment):
    """A snowflake `CREATE ... FUNCTION` statement for SQL and JavaScript functions.

    https://docs.snowflake.com/en/sql-reference/sql/create-function.html
    """

    type = "create_function_statement"
    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        Sequence("SECURE", optional=True),
        "FUNCTION",
        Ref("FunctionNameSegment"),
        Ref("FunctionParameterListGrammar"),
        "RETURNS",
        OneOf(
            Ref("DatatypeSegment"),
            Sequence("TABLE", Bracketed(Delimited(Ref("ColumnDefinitionSegment")))),
        ),
        Sequence("NOT", "NULL", optional=True),
        OneOf("VOLATILE", "IMMUTABLE", optional=True),
        Sequence("LANGUAGE", "JAVASCRIPT", optional=True),
        OneOf(
            Sequence("CALLED", "ON", "NULL", "INPUT"),
            Sequence("RETURNS", "NULL", "ON", "NULL", "INPUT"),
            "STRICT",
            optional=True,
        ),
        OneOf("VOLATILE", "IMMUTABLE", optional=True),
        Ref("CommentEqualsClauseSegment", optional=True),
        "AS",
        OneOf(
            Ref("DoubleQuotedUDFBody"),
            Ref("SingleQuotedUDFBody"),
            Ref("DollarQuotedUDFBody"),
        ),
    )


@snowflake_dialect.segment()
class AlterFunctionStatementSegment(BaseSegment):
    """A snowflake `ALTER ... FUNCTION` statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-function.html
    """

    type = "alter_function_statement"
    match_grammar = Sequence(
        "ALTER",
        "FUNCTION",
        Sequence("IF", "EXISTS", optional=True),
        Ref("FunctionNameSegment"),
        Ref("FunctionParameterListGrammar"),
        OneOf(
            Sequence("RENAME", "TO", Ref("FunctionNameSegment")),
            Sequence("SET", OneOf("SECURE", Ref("CommentEqualsClauseSegment"))),
            Sequence("UNSET", OneOf("SECURE", "COMMENT")),
        ),
    )


@snowflake_dialect.segment()
class WarehouseObjectPropertiesSegment(BaseSegment):
    """A snowflake Warehouse Object Properties segment.

    https://docs.snowflake.com/en/sql-reference/sql/create-warehouse.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-warehouse.html

    Note: comments are handled seperately so not incorrectly marked as
    warehouse object.
    """

    type = "warehouse_object_properties"

    match_grammar = AnySetOf(
        Sequence(
            "WAREHOUSE_SIZE",
            Ref("EqualsSegment"),
            Ref("WarehouseSize"),
        ),
        Sequence(
            "WAIT_FOR_COMPLETION",
            Ref("EqualsSegment"),
            Ref("BooleanLiteralGrammar"),
        ),
        Sequence(
            "MAX_CLUSTER_COUNT",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "MIN_CLUSTER_COUNT",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "SCALING_POLICY",
            Ref("EqualsSegment"),
            OneOf(
                "STANDARD",
                "ECONOMY",
            ),
        ),
        Sequence(
            "AUTO_SUSPEND",
            Ref("EqualsSegment"),
            OneOf(
                Ref("NumericLiteralSegment"),
                "NULL",
            ),
        ),
        Sequence(
            "AUTO_RESUME",
            Ref("EqualsSegment"),
            Ref("BooleanLiteralGrammar"),
        ),
        Sequence(
            "INITIALLY_SUSPENDED",
            Ref("EqualsSegment"),
            Ref("BooleanLiteralGrammar"),
        ),
        Sequence(
            "RESOURCE_MONITOR",
            Ref("EqualsSegment"),
            Ref("NakedIdentifierSegment"),
        ),
    )


@snowflake_dialect.segment()
class WarehouseObjectParamsSegment(BaseSegment):
    """A snowflake Warehouse Object Param segment.

    https://docs.snowflake.com/en/sql-reference/sql/create-warehouse.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-warehouse.html
    """

    type = "warehouse_object_properties"

    match_grammar = AnySetOf(
        Sequence(
            "MAX_CONCURRENCY_LEVEL",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "STATEMENT_QUEUED_TIMEOUT_IN_SECONDS",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "STATEMENT_TIMEOUT_IN_SECONDS",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
    )


@snowflake_dialect.segment()
class ConstraintPropertiesSegment(BaseSegment):
    """CONSTRAINT clause for CREATE TABLE or ALTER TABLE command.

    https://docs.snowflake.com/en/sql-reference/constraints-properties.html
    """

    type = "constraint_properties_segment"
    match_grammar = Sequence(
        Sequence("CONSTRAINT", Ref("QuotedLiteralSegment"), optional=True),
        OneOf(
            Sequence("UNIQUE", Bracketed(Ref("ColumnReferenceSegment"), optional=True)),
            Sequence(
                Ref("PrimaryKeyGrammar"),
                Bracketed(Ref("ColumnReferenceSegment"), optional=True),
            ),
            Sequence(
                Sequence(
                    Ref("ForeignKeyGrammar"),
                    Bracketed(Ref("ColumnReferenceSegment"), optional=True),
                    optional=True,
                ),
                "REFERENCES",
                Ref("TableReferenceSegment"),
                Bracketed(Ref("ColumnReferenceSegment")),
            ),
        ),
        AnySetOf(
            OneOf(Sequence("NOT", optional=True), "ENFORCED"),
            OneOf(Sequence("NOT", optional=True), "DEFERRABLE"),
            OneOf("INITIALLY", OneOf("DEFERRED", "IMMEDIATE")),
        ),
    )


@snowflake_dialect.segment(replace=True)
class ColumnConstraintSegment(BaseSegment):
    """A column option; each CREATE TABLE column can have 0 or more.

    https://docs.snowflake.com/en/sql-reference/sql/create-table.html
    """

    type = "column_constraint_segment"
    match_grammar = AnySetOf(
        Sequence("COLLATE", Ref("QuotedLiteralSegment")),
        Sequence(
            "DEFAULT",
            OneOf(
                Ref("QuotedLiteralSegment"),
                # https://docs.snowflake.com/en/sql-reference/functions/current_timestamp.html
                Sequence(
                    "CURRENT_TIMESTAMP",
                    Bracketed(
                        Ref("NumericLiteralSegment", optional=True), optional=True
                    ),
                ),
                # https://docs.snowflake.com/en/sql-reference/functions/sysdate.html
                Sequence("SYSDATE", Bracketed()),
            ),
        ),
        Sequence(
            OneOf("AUTOINCREMENT", "IDENTITY"),
            OneOf(
                Bracketed(Delimited(Ref("NumericLiteralSegment"))),
                Sequence(
                    "START",
                    Ref("NumericLiteralSegment"),
                    "INCREMENT",
                    Ref("NumericLiteralSegment"),
                ),
                optional=True,
            ),
        ),
        Sequence(Ref.keyword("NOT", optional=True), "NULL"),  # NOT NULL or NULL
        Sequence(
            Sequence("WITH", optional=True),
            "MASKING",
            "POLICY",
            Ref("FunctionNameSegment"),
            Sequence(
                "USING",
                Bracketed(
                    Delimited(
                        OneOf(
                            Ref("ColumnReferenceSegment"),
                            Ref("ExpressionSegment"),
                        )
                    ),
                ),
                optional=True,
            ),
        ),
        Ref("TagBracketedEqualsSegment", optional=True),
        Ref("ConstraintPropertiesSegment"),
        Sequence("DEFAULT", Ref("QuotedLiteralSegment")),
        Sequence("CHECK", Bracketed(Ref("ExpressionSegment"))),
        Sequence(  # DEFAULT <value>
            "DEFAULT",
            OneOf(
                Ref("LiteralGrammar"),
                Ref("FunctionSegment"),
                # ?? Ref('IntervalExpressionSegment')
            ),
        ),
        Sequence(  # REFERENCES reftable [ ( refcolumn) ]
            "REFERENCES",
            Ref("ColumnReferenceSegment"),
            # Foreign columns making up FOREIGN KEY constraint
            Ref("BracketedColumnReferenceListGrammar", optional=True),
        ),
    )


@snowflake_dialect.segment()
class CopyOptionsSegment(BaseSegment):
    """A Snowflake CopyOptions statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-table.html
    """

    type = "copy_options"
    match_grammar = AnySetOf(
        Sequence("ON_ERROR", Ref("EqualsSegment"), Ref("CopyOptionOnErrorSegment")),
        Sequence("SIZE_LIMIT", Ref("EqualsSegment"), Ref("LiteralNumericSegment")),
        Sequence("PURGE", Ref("EqualsSegment"), Ref("BooleanLiteralGrammar")),
        Sequence(
            "RETURN_FAILED_ONLY", Ref("EqualsSegment"), Ref("BooleanLiteralGrammar")
        ),
        Sequence(
            "MATCH_BY_COLUMN_NAME",
            Ref("EqualsSegment"),
            OneOf("CASE_SENSITIVE", "CASE_INSENSITIVE", "NONE"),
        ),
        Sequence("ENFORCE_LENGTH", Ref("EqualsSegment"), Ref("BooleanLiteralGrammar")),
        Sequence("TRUNCATECOLUMNS", Ref("EqualsSegment"), Ref("BooleanLiteralGrammar")),
        Sequence("FORCE", Ref("EqualsSegment"), Ref("BooleanLiteralGrammar")),
    )


@snowflake_dialect.segment(replace=True)
class CreateSchemaStatementSegment(BaseSegment):
    """A `CREATE SCHEMA` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-schema.html
    """

    type = "create_schema_statement"
    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        Ref("TemporaryTransientGrammar", optional=True),
        "SCHEMA",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("SchemaReferenceSegment"),
        Sequence("WITH", "MANAGED", "ACCESS", optional=True),
        Ref("SchemaObjectParamsSegment", optional=True),
        Ref("TagBracketedEqualsSegment", optional=True),
    )


@snowflake_dialect.segment()
class AlterSchemaStatementSegment(BaseSegment):
    """An `ALTER SCHEMA` statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-schema.html

    """

    type = "alter_schema_statement"
    match_grammar = Sequence(
        "ALTER",
        "SCHEMA",
        Sequence("IF", "EXISTS", optional=True),
        Ref("SchemaReferenceSegment"),
        OneOf(
            Sequence(
                "RENAME",
                "TO",
                Ref("SchemaReferenceSegment"),
            ),
            Sequence(
                "SWAP",
                "WITH",
                Ref("SchemaReferenceSegment"),
            ),
            Sequence(
                "SET",
                OneOf(Ref("SchemaObjectParamsSegment"), Ref("TagEqualsSegment")),
            ),
            Sequence(
                "UNSET",
                OneOf(
                    Delimited(
                        "DATA_RETENTION_TIME_IN_DAYS",
                        "MAX_DATA_EXTENSION_TIME_IN_DAYS",
                        "DEFAULT_DDL_COLLATION",
                        "COMMENT",
                    ),
                    Sequence("TAG", Delimited(Ref("NakedIdentifierSegment"))),
                ),
            ),
            Sequence(OneOf("ENABLE", "DISABLE"), Sequence("MANAGED", "ACCESS")),
        ),
    )


@snowflake_dialect.segment()
class SchemaObjectParamsSegment(BaseSegment):
    """A Snowflake Schema Object Param segment.

    https://docs.snowflake.com/en/sql-reference/sql/create-schema.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-schema.html
    """

    type = "schema_object_properties"

    match_grammar = AnySetOf(
        Sequence(
            "DATA_RETENTION_TIME_IN_DAYS",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "MAX_DATA_EXTENSION_TIME_IN_DAYS",
            Ref("EqualsSegment"),
            Ref("NumericLiteralSegment"),
        ),
        Sequence(
            "DEFAULT_DDL_COLLATION",
            Ref("EqualsSegment"),
            Ref("QuotedLiteralSegment"),
        ),
        Ref("CommentEqualsClauseSegment"),
    )


@snowflake_dialect.segment(replace=True)
class CreateTableStatementSegment(BaseSegment):
    """A `CREATE TABLE` statement.

    A lot more options than ANSI
    https://docs.snowflake.com/en/sql-reference/sql/create-table.html
    """

    type = "create_table_statement"
    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        Ref("TemporaryTransientGrammar", optional=True),
        "TABLE",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("TableReferenceSegment"),
        # Columns and comment syntax:
        Sequence(
            Bracketed(
                Delimited(
                    Sequence(
                        OneOf(
                            Ref("TableConstraintSegment"),
                            Ref("ColumnDefinitionSegment"),
                            Ref("SingleIdentifierGrammar"),
                        ),
                        Ref("CommentClauseSegment", optional=True),
                    ),
                ),
            ),
            optional=True,
        ),
        Sequence(
            "CLUSTER",
            "BY",
            Bracketed(Delimited(Ref("ExpressionSegment"))),
            optional=True,
        ),
        Sequence(
            "STAGE_FILE_FORMAT",
            Ref("EqualsSegment"),
            Ref("FileFormatSegment"),
            optional=True,
        ),
        Sequence(
            "STAGE_COPY_OPTIONS",
            Ref("EqualsSegment"),
            Bracketed(Ref("CopyOptionsSegment")),
            optional=True,
        ),
        Sequence(
            "DATA_RETENTION_TIME_IN_DAYS",
            Ref("EqualsSegment"),
            Ref("LiteralNumericSegment"),
            optional=True,
        ),
        Sequence(
            "MAX_DATA_EXTENSION_TIME_IN_DAYS",
            Ref("EqualsSegment"),
            Ref("LiteralNumericSegment"),
            optional=True,
        ),
        Sequence(
            "CHANGE_TRACKING",
            Ref("EqualsSegment"),
            Ref("BooleanLiteralGrammar"),
            optional=True,
        ),
        Sequence(
            "DEFAULT_DDL_COLLATION",
            Ref("EqualsSegment"),
            Ref("QuotedLiteralGrammar"),
            optional=True,
        ),
        Sequence(
            "COPY",
            "GRANTS",
            optional=True,
        ),
        Sequence(
            Sequence("WITH", optional=True),
            "ROW",
            "ACCESS",
            "POLICY",
            Ref("NakedIdentifierSegment"),
            "ON",
            Bracketed(Delimited(Ref("ColumnReferenceSegment"))),
            optional=True,
        ),
        Ref("TagBracketedEqualsSegment", optional=True),
        Ref("CommentEqualsClauseSegment", optional=True),
        OneOf(
            # Create AS syntax:
            Sequence(
                "AS",
                OptionallyBracketed(Ref("SelectableGrammar")),
            ),
            # Create like syntax
            Sequence("LIKE", Ref("TableReferenceSegment")),
            # Create clone syntax
            Sequence("ClONE", Ref("TableReferenceSegment")),
            Sequence("USING", "TEMPLATE", Ref("SelectableGrammar")),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class CreateTaskSegment(BaseSegment):
    """A snowflake `CREATE TASK` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-task.html
    """

    type = "create_task_statement"

    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        "TASK",
        Sequence("IF", "NOT", "EXISTS", optional=True),
        Ref("ObjectReferenceSegment"),
        Indent,
        AnyNumberOf(
            OneOf(
                Sequence(
                    "WAREHOUSE",
                    Ref("EqualsSegment"),
                    Ref("ObjectReferenceSegment"),
                ),
                Sequence(
                    "USER_TASK_MANAGED_INITIAL_WAREHOUSE_SIZE",
                    Ref("EqualsSegment"),
                    Ref("WarehouseSize"),
                ),
            ),
            Sequence(
                "SCHEDULE",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            Sequence(
                "ALLOW_OVERLAPPING_EXECUTION",
                Ref("EqualsSegment"),
                Ref("BooleanLiteralGrammar"),
            ),
            Sequence(
                "USER_TASK_TIMEOUT_MS",
                Ref("EqualsSegment"),
                Ref("NumericLiteralSegment"),
            ),
            Delimited(
                Sequence(
                    Ref("ParameterNameSegment"),
                    Ref("EqualsSegment"),
                    OneOf(
                        Ref("BooleanLiteralGrammar"),
                        Ref("QuotedLiteralSegment"),
                        Ref("NumericLiteralSegment"),
                    ),
                ),
            ),
            Sequence(
                "COPY",
                "GRANTS",
            ),
            Ref("CommentEqualsClauseSegment"),
        ),
        Sequence(
            "AFTER",
            Ref("ObjectReferenceSegment"),
            optional=True,
        ),
        Dedent,
        Sequence(
            "WHEN",
            Indent,
            Ref("ExpressionSegment"),
            Dedent,
            optional=True,
        ),
        Sequence(
            Ref.keyword("AS"),
            Indent,
            Ref("StatementSegment"),
            Dedent,
        ),
    )


@snowflake_dialect.segment()
class CreateStatementSegment(BaseSegment):
    """A snowflake `CREATE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create.html
    """

    type = "create_statement"

    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        OneOf(
            Sequence("NETWORK", "POLICY"),
            Sequence("RESOURCE", "MONITOR"),
            "SHARE",
            "ROLE",
            "USER",
            "WAREHOUSE",
            Sequence("NOTIFICATION", "INTEGRATION"),
            Sequence("SECURITY", "INTEGRATION"),
            Sequence("STORAGE", "INTEGRATION"),
            Sequence("MATERIALIZED", "VIEW"),
            Sequence("MASKING", "POLICY"),
            "PIPE",
            Sequence("EXTERNAL", "FUNCTION"),
            # Objects that also support clone
            "DATABASE",
            "SEQUENCE",
            Sequence("FILE", "FORMAT"),
            "STREAM",
        ),
        Ref("IfNotExistsGrammar", optional=True),
        Ref("ObjectReferenceSegment"),
        # Next set are Pipe statements
        # https://docs.snowflake.com/en/sql-reference/sql/create-pipe.html
        Sequence(
            Sequence(
                "AUTO_INGEST",
                Ref("EqualsSegment"),
                Ref("BooleanLiteralGrammar"),
                optional=True,
            ),
            Sequence(
                "AWS_SNS_TOPIC",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
                optional=True,
            ),
            Sequence(
                "INTEGRATION",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
                optional=True,
            ),
            optional=True,
        ),
        # Next are WAREHOUSE options
        # https://docs.snowflake.com/en/sql-reference/sql/create-warehouse.html
        Sequence(
            Sequence("WITH", optional=True),
            AnyNumberOf(
                Ref("WarehouseObjectPropertiesSegment"),
                Ref("WarehouseObjectParamsSegment"),
            ),
            Ref("TagBracketedEqualsSegment", optional=True),
            optional=True,
        ),
        Ref("CommentEqualsClauseSegment", optional=True),
        Ref.keyword("AS", optional=True),
        OneOf(
            Ref("SelectStatementSegment"),
            Sequence(
                Bracketed(Ref("FunctionContentsGrammar"), optional=True),
                "RETURNS",
                Ref("DatatypeSegment"),
                Ref("FunctionAssignerSegment"),
                Ref("ExpressionSegment"),
                Sequence(
                    "COMMENT",
                    Ref("EqualsSegment"),
                    Ref("QuotedLiteralSegment"),
                    optional=True,
                ),
                optional=True,
            ),
            Ref("CopyIntoTableStatementSegment"),
            optional=True,
        ),
    )


@snowflake_dialect.segment(replace=True)
class CreateViewStatementSegment(BaseSegment):
    """A `CREATE VIEW` statement, specifically for Snowflake's dialect.

    https://docs.snowflake.com/en/sql-reference/sql/create-view.html
    """

    type = "create_view_statement"

    match_grammar = Sequence(
        "CREATE",
        Ref("OrReplaceGrammar", optional=True),
        AnySetOf(
            "SECURE",
            "RECURSIVE",
        ),
        "VIEW",
        Ref("IfNotExistsGrammar", optional=True),
        Ref("TableReferenceSegment"),
        AnySetOf(
            Bracketed(
                Delimited(
                    Sequence(
                        Ref("ColumnReferenceSegment"),
                        Ref("CommentClauseSegment", optional=True),
                    ),
                ),
            ),
            Sequence(
                Ref.keyword("WITH", optional=True),
                "ROW",
                "ACCESS",
                "POLICY",
                Ref("NakedIdentifierSegment"),
                "ON",
                Bracketed(
                    Delimited(Ref("ColumnReferenceSegment")),
                ),
            ),
            Ref("TagBracketedEqualsSegment"),
            Sequence("COPY", "GRANTS"),
            Ref("CommentEqualsClauseSegment"),
            # @TODO: Support column-level masking policy & tagging.
        ),
        "AS",
        OptionallyBracketed(Ref("SelectableGrammar")),
    )


@snowflake_dialect.segment()
class AlterViewStatementSegment(BaseSegment):
    """An `ALTER VIEW` statement, specifically for Snowflake's dialect.

    https://docs.snowflake.com/en/sql-reference/sql/alter-view.html
    """

    type = "alter_view_statement"

    match_grammar = Sequence(
        "ALTER",
        "VIEW",
        Ref("IfExistsGrammar", optional=True),
        Ref("TableReferenceSegment"),
        OneOf(
            Sequence(
                "RENAME",
                "TO",
                Ref("TableReferenceSegment"),
            ),
            Sequence(
                "COMMENT",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            Sequence(
                "UNSET",
                "COMMENT",
            ),
            Sequence(
                OneOf("SET", "UNSET"),
                "SECURE",
            ),
            Sequence("SET", Ref("TagEqualsSegment")),
            Sequence("UNSET", "TAG", Delimited(Ref("NakedIdentifierSegment"))),
            Delimited(
                Sequence(
                    "ADD",
                    "ROW",
                    "ACCESS",
                    "POLICY",
                    Ref("FunctionNameSegment"),
                    "ON",
                    Bracketed(Delimited(Ref("ColumnReferenceSegment"))),
                ),
                Sequence(
                    "DROP",
                    "ROW",
                    "ACCESS",
                    "POLICY",
                    Ref("FunctionNameSegment"),
                ),
            ),
            Sequence(
                OneOf("ALTER", "MODIFY"),
                OneOf(
                    Delimited(
                        Sequence(
                            Ref.keyword("COLUMN", optional=True),
                            Ref("ColumnReferenceSegment"),
                            OneOf(
                                Sequence(
                                    "SET",
                                    "MASKING",
                                    "POLICY",
                                    Ref("FunctionNameSegment"),
                                    Sequence(
                                        "USING",
                                        Bracketed(
                                            Delimited(Ref("ColumnReferenceSegment"))
                                        ),
                                        optional=True,
                                    ),
                                ),
                                Sequence("UNSET", "MASKING", "POLICY"),
                                Sequence("SET", Ref("TagEqualsSegment")),
                            ),
                        ),
                        Sequence(
                            "COLUMN",
                            Ref("ColumnReferenceSegment"),
                            "UNSET",
                            "TAG",
                            Delimited(Ref("NakedIdentifierSegment")),
                        ),
                    ),
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class FileFormatSegment(BaseSegment):
    """A Snowflake FILE_FORMAT Segment.

    https://docs.snowflake.com/en/sql-reference/sql/create-table.html
    https://docs.snowflake.com/en/sql-reference/sql/create-external-table.html
    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    """

    type = "file_format_segment"

    match_grammar = OneOf(
        OneOf(Ref("NakedIdentifierSegment"), Ref("QuotedLiteralSegment")),
        Bracketed(
            OneOf(
                Sequence(
                    "FORMAT_NAME",
                    Ref("EqualsSegment"),
                    OneOf(Ref("NakedIdentifierSegment"), Ref("QuotedLiteralSegment")),
                ),
                Sequence(
                    "TYPE",
                    Ref("EqualsSegment"),
                    Ref("FileType"),
                    # formatTypeOptions - To Do to make this more specific
                    Ref("FormatTypeOptionsSegment", optional=True),
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class FormatTypeOptionsSegment(BaseSegment):
    """A snowflake `formatTypeOptions` Segment.

    https://docs.snowflake.com/en/sql-reference/sql/create-table.html
    https://docs.snowflake.com/en/sql-reference/sql/create-external-table.html
    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    """

    type = "format_type_options_segment"

    match_grammar = AnyNumberOf(
        # formatTypeOptions - To Do to make this more specific
        Ref("NakedIdentifierSegment"),
        Ref("EqualsSegment"),
        OneOf(
            Ref("NakedIdentifierSegment"),
            Ref("QuotedLiteralSegment"),
            Ref("NumericLiteralSegment"),
            Bracketed(
                Delimited(
                    Ref("QuotedLiteralSegment"),
                )
            ),
        ),
    )


@snowflake_dialect.segment()
class CreateExternalTableSegment(BaseSegment):
    """A snowflake `CREATE EXTERNAL TABLE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-external-table.html
    """

    type = "create_external_table_statement"

    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        "EXTERNAL",
        "TABLE",
        Sequence("IF", "NOT", "EXISTS", optional=True),
        Ref("TableReferenceSegment"),
        # Columns:
        Bracketed(
            Delimited(
                Sequence(
                    Ref("SingleIdentifierGrammar"),
                    Ref("DatatypeSegment"),
                    "AS",
                    OptionallyBracketed(
                        Sequence(
                            Ref("ExpressionSegment"),
                            Ref("TableConstraintSegment", optional=True),
                        )
                    ),
                )
            ),
            optional=True,
        ),
        # The use of AnySetOf is not strictly correct here, because LOCATION and
        # FILE_FORMAT are required parameters. They can however be in arbitrary order
        # with the other parameters.
        AnySetOf(
            Sequence("INTEGRATION", Ref("EqualsSegment"), Ref("QuotedLiteralSegment")),
            Sequence(
                "PARTITION",
                "BY",
                Bracketed(Delimited(Ref("SingleIdentifierGrammar"))),
            ),
            Sequence(
                Sequence("WITH", optional=True),
                "LOCATION",
                Ref("EqualsSegment"),
                Ref("StagePath"),
            ),
            Sequence(
                "REFRESH_ON_CREATE",
                Ref("EqualsSegment"),
                Ref("BooleanLiteralGrammar"),
            ),
            Sequence(
                "AUTO_REFRESH",
                Ref("EqualsSegment"),
                Ref("BooleanLiteralGrammar"),
            ),
            Sequence(
                "PATTERN",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            Sequence(
                "FILE_FORMAT",
                Ref("EqualsSegment"),
                Ref("FileFormatSegment"),
            ),
            Sequence(
                "AWS_SNS_TOPIC",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            Sequence(
                "COPY",
                "GRANTS",
            ),
            Sequence(
                Sequence("WITH", optional=True),
                "ROW",
                "ACCESS",
                "POLICY",
                Ref("NakedIdentifierSegment"),
            ),
            Ref("TagBracketedEqualsSegment"),
            Ref("CommentEqualsClauseSegment"),
        ),
    )


@snowflake_dialect.segment(replace=True)
class TableExpressionSegment(BaseSegment):
    """The main table expression e.g. within a FROM clause."""

    type = "table_expression"
    match_grammar = OneOf(
        Ref("BareFunctionSegment"),
        Ref("FunctionSegment"),
        Ref("TableReferenceSegment"),
        # Nested Selects
        Bracketed(Ref("SelectableGrammar")),
        # Values clause?
        Ref("StagePath"),
    )


@snowflake_dialect.segment()
class CopyIntoTableStatementSegment(BaseSegment):
    """A Snowflake `COPY INTO <table>` statement.

    # https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html
    """

    type = "copy_into_statement"

    match_grammar = Sequence(
        "COPY",
        "INTO",
        Ref("TableReferenceSegment"),
        Bracketed(Delimited(Ref("ColumnReferenceSegment")), optional=True),
        Sequence(
            "FROM",
            OneOf(
                Ref("StorageLocation"),
                Bracketed(Ref("SelectStatementSegment")),
            ),
            optional=True,
        ),
        AnySetOf(
            Sequence(
                "FILES",
                Ref("EqualsSegment"),
                Bracketed(
                    Delimited(
                        Ref("QuotedLiteralSegment"),
                    ),
                ),
            ),
            Sequence(
                "PATTERN",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
            ),
            Sequence(
                "FILE_FORMAT",
                Ref("EqualsSegment"),
                Ref("FileFormatSegment"),
            ),
            Ref("CopyOptionsSegment"),
        ),
        Sequence(
            "VALIDATION_MODE",
            Ref("EqualsSegment"),
            Ref("ValidationModeOptionSegment"),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class StorageLocation(BaseSegment):
    """A Snowflake storage location.

    https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#syntax
    """

    type = "storage_location"

    match_grammar = OneOf(
        Ref("StagePath"), Ref("S3Path"), Ref("GCSPath"), Ref("AzureBlobStoragePath")
    )


@snowflake_dialect.segment()
class InternalStageParameters(BaseSegment):
    """Parameters for an internal stage in Snowflake.

    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-stage.html
    """

    name = "internal_stage_parameters"
    type = "stage_parameters"

    match_grammar = Sequence(
        Sequence(
            "ENCRYPTION",
            Ref("EqualsSegment"),
            Bracketed(
                "TYPE",
                Ref("EqualsSegment"),
                Ref("SnowflakeEncryptionOption"),
            ),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class S3ExternalStageParameters(BaseSegment):
    """Parameters for an S3 external stage in Snowflake.

    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-stage.html
    """

    name = "s3_external_stage_parameters"
    type = "stage_parameters"

    match_grammar = Sequence(
        OneOf(
            Sequence(
                "STORAGE_INTEGRATION",
                Ref("EqualsSegment"),
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "CREDENTIALS",
                Ref("EqualsSegment"),
                Bracketed(
                    OneOf(
                        Sequence(
                            "AWS_KEY_ID",
                            Ref("EqualsSegment"),
                            Ref("QuotedLiteralSegment"),
                            "AWS_SECRET_KEY",
                            Ref("EqualsSegment"),
                            Ref("QuotedLiteralSegment"),
                            Sequence(
                                "AWS_TOKEN",
                                Ref("EqualsSegment"),
                                Ref("QuotedLiteralSegment"),
                                optional=True,
                            ),
                        ),
                        Sequence(
                            "AWS_ROLE",
                            Ref("EqualsSegment"),
                            Ref("QuotedLiteralSegment"),
                        ),
                    )
                ),
            ),
            optional=True,
        ),
        Sequence(
            "ENCRYPTION",
            Ref("EqualsSegment"),
            Bracketed(
                OneOf(
                    Sequence(
                        Sequence(
                            "TYPE",
                            Ref("EqualsSegment"),
                            Ref("S3EncryptionOption"),
                            optional=True,
                        ),
                        "MASTER_KEY",
                        Ref("EqualsSegment"),
                        Ref("QuotedLiteralSegment"),
                    ),
                    Sequence("TYPE", Ref("EqualsSegment"), Ref("S3EncryptionOption")),
                    Sequence(
                        "TYPE",
                        Ref("EqualsSegment"),
                        Ref("S3EncryptionOption"),
                        Sequence(
                            "KMS_KEY_ID",
                            Ref("EqualsSegment"),
                            Ref("QuotedLiteralSegment"),
                            optional=True,
                        ),
                    ),
                    Sequence("TYPE", Ref("EqualsSegment"), "NONE"),
                )
            ),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class GCSExternalStageParameters(BaseSegment):
    """Parameters for a GCS external stage in Snowflake.

    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-stage.html
    """

    name = "gcs_external_stage_parameters"
    type = "stage_parameters"

    match_grammar = Sequence(
        Sequence(
            "STORAGE_INTEGRATION",
            Ref("EqualsSegment"),
            Ref("ObjectReferenceSegment"),
            optional=True,
        ),
        Sequence(
            "ENCRYPTION",
            Ref("EqualsSegment"),
            Bracketed(
                Sequence(
                    "TYPE",
                    Ref("EqualsSegment"),
                    OneOf(
                        Sequence(
                            Ref("GCSEncryptionOption"),
                            Sequence(
                                "KMS_KEY_ID",
                                Ref("EqualsSegment"),
                                Ref("QuotedLiteralSegment"),
                                optional=True,
                            ),
                        ),
                        "NONE",
                    ),
                )
            ),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class AzureBlobStorageExternalStageParameters(BaseSegment):
    """Parameters for an Azure Blob Storage external stage in Snowflake.

    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    https://docs.snowflake.com/en/sql-reference/sql/alter-stage.html
    """

    name = "azure_blob_storage_external_stage_parameters"
    type = "stage_parameters"

    match_grammar = Sequence(
        OneOf(
            Sequence(
                "STORAGE_INTEGRATION",
                Ref("EqualsSegment"),
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "CREDENTIALS",
                Ref("EqualsSegment"),
                Bracketed(
                    Sequence("AZURE_SAS_TOKEN"),
                    Ref("EqualsSegment"),
                    Ref("QuotedLiteralSegment"),
                ),
            ),
            optional=True,
        ),
        Sequence(
            "ENCRYPTION",
            Ref("EqualsSegment"),
            Bracketed(
                Sequence(
                    "TYPE",
                    Ref("EqualsSegment"),
                    OneOf(
                        Sequence(
                            Ref("AzureBlobStorageEncryptionOption"),
                            Sequence(
                                "MASTER_KEY",
                                Ref("EqualsSegment"),
                                Ref("QuotedLiteralSegment"),
                                optional=True,
                            ),
                        ),
                        "NONE",
                    ),
                )
            ),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class CreateStageSegment(BaseSegment):
    """A Snowflake CREATE STAGE statement.

    https://docs.snowflake.com/en/sql-reference/sql/create-stage.html
    """

    type = "create_stage_statement"

    match_grammar = Sequence(
        "CREATE",
        Sequence("OR", "REPLACE", optional=True),
        Ref.keyword("TEMPORARY", optional=True),
        "STAGE",
        Sequence("IF", "NOT", "EXISTS", optional=True),
        Ref("ObjectReferenceSegment"),
        Indent,
        OneOf(
            # Internal stages
            Sequence(
                Ref("InternalStageParameters", optional=True),
                Sequence(
                    "DIRECTORY",
                    Ref("EqualsSegment"),
                    Bracketed(
                        Sequence(
                            "ENABLE",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                        )
                    ),
                    optional=True,
                ),
            ),
            # External S3 stage
            Sequence(
                "URL",
                Ref("EqualsSegment"),
                Ref("S3Path"),
                Ref("S3ExternalStageParameters", optional=True),
                Sequence(
                    "DIRECTORY",
                    Ref("EqualsSegment"),
                    Bracketed(
                        Sequence(
                            "ENABLE",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                        ),
                        Sequence(
                            "AUTO_REFRESH",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                            optional=True,
                        ),
                    ),
                    optional=True,
                ),
            ),
            # External GCS stage
            Sequence(
                "URL",
                Ref("EqualsSegment"),
                Ref("GCSPath"),
                Ref("GCSExternalStageParameters", optional=True),
                Sequence(
                    "DIRECTORY",
                    Ref("EqualsSegment"),
                    Bracketed(
                        Sequence(
                            "ENABLE",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                        ),
                        Sequence(
                            "AUTO_REFRESH",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                            optional=True,
                        ),
                        Sequence(
                            "NOTIFICATION_INTEGRATION",
                            Ref("EqualsSegment"),
                            OneOf(
                                Ref("NakedIdentifierSegment"),
                                Ref("QuotedLiteralSegment"),
                            ),
                            optional=True,
                        ),
                    ),
                    optional=True,
                ),
            ),
            # External Azure Blob Storage stage
            Sequence(
                "URL",
                Ref("EqualsSegment"),
                Ref("AzureBlobStoragePath"),
                Ref("AzureBlobStorageExternalStageParameters", optional=True),
                Sequence(
                    "DIRECTORY",
                    Ref("EqualsSegment"),
                    Bracketed(
                        Sequence(
                            "ENABLE",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                        ),
                        Sequence(
                            "AUTO_REFRESH",
                            Ref("EqualsSegment"),
                            Ref("BooleanLiteralGrammar"),
                            optional=True,
                        ),
                        Sequence(
                            "NOTIFICATION_INTEGRATION",
                            Ref("EqualsSegment"),
                            OneOf(
                                Ref("NakedIdentifierSegment"),
                                Ref("QuotedLiteralSegment"),
                            ),
                            optional=True,
                        ),
                    ),
                    optional=True,
                ),
            ),
            optional=True,
        ),
        Sequence(
            "FILE_FORMAT", Ref("EqualsSegment"), Ref("FileFormatSegment"), optional=True
        ),
        Sequence(
            "COPY_OPTIONS",
            Ref("EqualsSegment"),
            Bracketed(Ref("CopyOptionsSegment")),
            optional=True,
        ),
        Ref("TagBracketedEqualsSegment", optional=True),
        Ref("CommentEqualsClauseSegment", optional=True),
        Dedent,
    )


@snowflake_dialect.segment()
class AlterStageSegment(BaseSegment):
    """A Snowflake ALTER STAGE statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-stage.html
    """

    type = "alter_stage_statement"

    match_grammar = Sequence(
        "ALTER",
        "STAGE",
        Sequence("IF", "EXISTS", optional=True),
        Ref("ObjectReferenceSegment"),
        OneOf(
            Sequence("RENAME", "TO", Ref("ObjectReferenceSegment")),
            Sequence(
                "SET",
                Indent,
                OneOf(
                    Sequence(
                        OneOf(
                            Ref("InternalStageParameters"),
                            Sequence(
                                Sequence(
                                    "URL",
                                    Ref("EqualsSegment"),
                                    Ref("S3Path"),
                                    optional=True,
                                ),
                                Ref(
                                    "S3ExternalStageParameters",
                                    optional=True,
                                ),
                            ),
                            Sequence(
                                Sequence(
                                    "URL",
                                    Ref("EqualsSegment"),
                                    Ref("GCSPath"),
                                    optional=True,
                                ),
                                Ref(
                                    "GCSExternalStageParameters",
                                    optional=True,
                                ),
                            ),
                            Sequence(
                                Sequence(
                                    "URL",
                                    Ref("EqualsSegment"),
                                    Ref("AzureBlobStoragePath"),
                                    optional=True,
                                ),
                                Ref(
                                    "AzureBlobStorageExternalStageParameters",
                                    optional=True,
                                ),
                            ),
                            optional=True,
                        ),
                        Sequence(
                            "FILE_FORMAT",
                            Ref("EqualsSegment"),
                            Ref("FileFormatSegment"),
                            optional=True,
                        ),
                        Sequence(
                            "COPY_OPTIONS",
                            Ref("EqualsSegment"),
                            Bracketed(Ref("CopyOptionsSegment")),
                            optional=True,
                        ),
                        Ref("CommentEqualsClauseSegment", optional=True),
                    ),
                    Ref("TagEqualsSegment"),
                ),
                Dedent,
            ),
            Sequence(
                "REFRESH",
                Sequence(
                    "SUBPATH",
                    Ref("EqualsSegment"),
                    Ref("QuotedLiteralSegment"),
                    optional=True,
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class ShowStatementSegment(BaseSegment):
    """A snowflake `SHOW` statement.

    https://docs.snowflake.com/en/sql-reference/sql/show.html
    """

    _object_types_plural = OneOf(
        "PARAMETERS",
        Sequence("GLOBAL", "ACCOUNTS"),
        "REGIONS",
        Sequence("REPLICATION", "ACCOUNTS"),
        Sequence("REPLICATION", "DATABASES"),
        "PARAMETERS",
        "VARIABLES",
        "TRANSACTIONS",
        "LOCKS",
        "PARAMETERS",
        "FUNCTIONS",
        Sequence("NETWORK", "POLICIES"),
        "SHARES",
        "ROLES",
        "GRANTS",
        "USERS",
        "WAREHOUSES",
        "DATABASES",
        Sequence(
            OneOf("API", "NOTIFICATION", "SECURITY", "STORAGE", optional=True),
            "INTEGRATIONS",
        ),
        "SCHEMAS",
        "OBJECTS",
        "TABLES",
        Sequence("EXTERNAL", "TABLES"),
        "VIEWS",
        Sequence("MATERIALIZED", "VIEWS"),
        Sequence("MASKING", "POLICIES"),
        "COLUMNS",
        Sequence("FILE", "FORMATS"),
        "SEQUENCES",
        "STAGES",
        "PIPES",
        "STREAMS",
        "TASKS",
        Sequence("USER", "FUNCTIONS"),
        Sequence("EXTERNAL", "FUNCTIONS"),
        "PROCEDURES",
        Sequence("FUTURE", "GRANTS"),
    )

    _object_scope_types = OneOf(
        "ACCOUNT",
        "SESSION",
        Sequence(
            OneOf(
                "DATABASE",
                "SCHEMA",
                "SHARE",
                "ROLE",
                "TABLE",
                "TASK",
                "USER",
                "WAREHOUSE",
                "VIEW",
            ),
            Ref("ObjectReferenceSegment", optional=True),
        ),
    )

    type = "show_statement"

    match_grammar = Sequence(
        "SHOW",
        OneOf("TERSE", optional=True),
        _object_types_plural,
        OneOf("HISTORY", optional=True),
        Sequence("LIKE", Ref("QuotedLiteralSegment"), optional=True),
        Sequence(
            OneOf("ON", "TO", "OF", "IN"),
            OneOf(
                Sequence(_object_scope_types),
                Ref("ObjectReferenceSegment"),
            ),
            optional=True,
        ),
        Sequence("STARTS", "WITH", Ref("QuotedLiteralSegment"), optional=True),
        Sequence("WITH", "PRIMARY", Ref("ObjectReferenceSegment"), optional=True),
        Sequence(
            Ref("LimitClauseSegment"),
            Sequence("FROM", Ref("QuotedLiteralSegment"), optional=True),
            optional=True,
        ),
    )


@snowflake_dialect.segment()
class AlterUserSegment(BaseSegment):
    """`ALTER USER` statement.

    https://docs.snowflake.com/en/sql-reference/sql/alter-user.html

    All user parameters can be found here
    https://docs.snowflake.com/en/sql-reference/parameters.html
    """

    type = "alter_user"

    match_grammar = StartsWith(
        Sequence("ALTER", "USER"),
    )
    parse_grammar = Sequence(
        "ALTER",
        "USER",
        Sequence("IF", "EXISTS", optional=True),
        Ref("ObjectReferenceSegment"),
        OneOf(
            Sequence("RENAME", "TO", Ref("ObjectReferenceSegment")),
            Sequence("RESET", "PASSWORD"),
            Sequence("ABORT", "ALL", "QUERIES"),
            Sequence(
                "ADD",
                "DELEGATED",
                "AUTHORIZATION",
                "OF",
                "ROLE",
                Ref("ObjectReferenceSegment"),
                "TO",
                "SECURITY",
                "INTEGRATION",
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "REMOVE",
                "DELEGATED",
                OneOf(
                    Sequence(
                        "AUTHORIZATION", "OF", "ROLE", Ref("ObjectReferenceSegment")
                    ),
                    "AUTHORIZATIONS",
                ),
                "FROM",
                "SECURITY",
                "INTEGRATION",
                Ref("ObjectReferenceSegment"),
            ),
            # Snowflake supports the SET command with space delimitted parameters, but
            # it also supports using commas which is better supported by `Delimited`, so
            # we will just use that.
            Sequence(
                "SET",
                Delimited(
                    Sequence(
                        Ref("ParameterNameSegment"),
                        Ref("EqualsSegment"),
                        OneOf(Ref("LiteralGrammar"), Ref("ObjectReferenceSegment")),
                    ),
                ),
            ),
            Sequence("UNSET", Delimited(Ref("ParameterNameSegment"))),
        ),
    )


@snowflake_dialect.segment(replace=True)
class CreateRoleStatementSegment(BaseSegment):
    """A `CREATE ROLE` statement.

    Redefined because it's much simpler than postgres.
    https://docs.snowflake.com/en/sql-reference/sql/create-role.html
    """

    type = "create_role_statement"
    match_grammar = Sequence(
        "CREATE",
        Sequence(
            "OR",
            "REPLACE",
            optional=True,
        ),
        "ROLE",
        Sequence(
            "IF",
            "NOT",
            "EXISTS",
            optional=True,
        ),
        Ref("ObjectReferenceSegment"),
        Sequence(
            "COMMENT",
            Ref("EqualsSegment"),
            Ref("QuotedLiteralSegment"),
            optional=True,
        ),
    )


@snowflake_dialect.segment(replace=True)
class ExplainStatementSegment(
    ansi_dialect.get_segment("ExplainStatementSegment")  # type: ignore
):
    """An `Explain` statement.

    EXPLAIN [ USING { TABULAR | JSON | TEXT } ] <statement>

    https://docs.snowflake.com/en/sql-reference/sql/explain.html
    """

    match_grammar = Sequence(
        "EXPLAIN",
        Sequence(
            "USING",
            OneOf("TABULAR", "JSON", "TEXT"),
            optional=True,
        ),
        ansi_dialect.get_segment("ExplainStatementSegment").explainable_stmt,
    )


@snowflake_dialect.segment()
class AlterSessionStatementSegment(BaseSegment):
    """Snowflake's ALTER SESSION statement.

    ```
    ALTER SESSION SET <param_name> = <param_value>;
    ALTER SESSION UNSET <param_name>, [ , <param_name> , ... ];
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-session.html
    """

    type = "alter_session_statement"

    match_grammar = Sequence(
        "ALTER",
        "SESSION",
        OneOf(
            Ref("AlterSessionSetClauseSegment"),
            Ref("AlterSessionUnsetClauseSegment"),
        ),
    )


@snowflake_dialect.segment()
class AlterSessionSetClauseSegment(BaseSegment):
    """Snowflake's ALTER SESSION SET clause.

    ```
    [ALTER SESSION] SET <param_name> = <param_value>;
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-session.html
    """

    type = "alter_session_set_statement"

    match_grammar = Sequence(
        "SET",
        Ref("ParameterNameSegment"),
        Ref("EqualsSegment"),
        OneOf(
            Ref("BooleanLiteralGrammar"),
            Ref("QuotedLiteralSegment"),
            Ref("NumericLiteralSegment"),
        ),
    )


@snowflake_dialect.segment()
class AlterSessionUnsetClauseSegment(BaseSegment):
    """Snowflake's ALTER SESSION UNSET clause.

    ```
    [ALTER SESSION] UNSET <param_name>, [ , <param_name> , ... ];
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-session.html
    """

    type = "alter_session_unset_clause"

    match_grammar = Sequence(
        "UNSET",
        Delimited(Ref("ParameterNameSegment"), delimiter=Ref("CommaSegment")),
    )


@snowflake_dialect.segment()
class AlterTaskStatementSegment(BaseSegment):
    """Snowflake's ALTER TASK statement.

    ```
    ALTER TASK [IF EXISTS] <name> RESUME;
    ALTER TASK [IF EXISTS] <name> SUSPEND;
    ALTER TASK [IF EXISTS] <name> REMOVE AFTER <value>;
    ALTER TASK [IF EXISTS] <name> ADD AFTER <value>;
    ALTER TASK [IF EXISTS] <name> SET
        [WAREHOUSE = <value>]
        [SCHEDULE = <value>]
        [ALLOW_OVERLAPPING_EXECUTION = TRUE|FALSE];
    ALTER TASK [IF EXISTS] <name> SET
        <param_name> = <param_value> [ , <param_name> = <param_value> , ...];
    ALTER TASK [IF EXISTS] <name> UNSET <param_name> [ , <param_name> , ... ];
    ALTER TASK [IF EXISTS] <name> MODIFY AS <sql>;
    ALTER TASK [IF EXISTS] <name> MODIFY WHEN <boolean>;
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-task.html
    """

    type = "alter_task_statement"

    match_grammar = Sequence(
        "ALTER",
        "TASK",
        Sequence("IF", "EXISTS", optional=True),
        Ref("ObjectReferenceSegment"),
        OneOf(
            "RESUME",
            "SUSPEND",
            Sequence("REMOVE", "AFTER", Ref("ObjectReferenceSegment")),
            Sequence("ADD", "AFTER", Ref("ObjectReferenceSegment")),
            Ref("AlterTaskSpecialSetClauseSegment"),
            Ref("AlterTaskSetClauseSegment"),
            Ref("AlterTaskUnsetClauseSegment"),
            Sequence(
                "MODIFY",
                "AS",
                ansi_dialect.get_segment("ExplainStatementSegment").explainable_stmt,
            ),
            Sequence("MODIFY", "WHEN", Ref("BooleanLiteralGrammar")),
        ),
    )


@snowflake_dialect.segment()
class AlterTaskSpecialSetClauseSegment(BaseSegment):
    """Snowflake's ALTER TASK special SET clause.

    ```
    [ALTER TASK <name>] SET
        [WAREHOUSE = <value>]
        [SCHEDULE = <value>]
        [ALLOW_OVERLAPPING_EXECUTION = TRUE|FALSE];
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-task.html
    """

    type = "alter_task_special_set_clause"

    match_grammar = Sequence(
        "SET",
        AnySetOf(
            Sequence(
                "WAREHOUSE",
                Ref("EqualsSegment"),
                Ref("ObjectReferenceSegment"),
                optional=True,
            ),
            Sequence(
                "SCHEDULE",
                Ref("EqualsSegment"),
                Ref("QuotedLiteralSegment"),
                optional=True,
            ),
            Sequence(
                "ALLOW_OVERLAPPING_EXECUTION",
                Ref("EqualsSegment"),
                Ref("BooleanLiteralGrammar"),
                optional=True,
            ),
            min_times=1,
        ),
    )


@snowflake_dialect.segment()
class AlterTaskSetClauseSegment(BaseSegment):
    """Snowflake's ALTER TASK SET clause.

    ```
    [ALTER TASK <name>] SET
        <param_name> = <param_value> [ , <param_name> = <param_value> , ...];
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-task.html
    """

    type = "alter_task_set_clause"

    match_grammar = Sequence(
        "SET",
        Delimited(
            Sequence(
                Ref("ParameterNameSegment"),
                Ref("EqualsSegment"),
                OneOf(
                    Ref("BooleanLiteralGrammar"),
                    Ref("QuotedLiteralSegment"),
                    Ref("NumericLiteralSegment"),
                ),
            ),
            delimiter=Ref("CommaSegment"),
        ),
    )


@snowflake_dialect.segment()
class AlterTaskUnsetClauseSegment(BaseSegment):
    """Snowflake's ALTER TASK UNSET clause.

    ```
    [ALTER TASK <name>] UNSET <param_name> [ , <param_name> , ... ];
    ```

    https://docs.snowflake.com/en/sql-reference/sql/alter-task.html
    """

    type = "alter_task_unset_clause"

    match_grammar = Sequence(
        "UNSET",
        Delimited(Ref("ParameterNameSegment"), delimiter=Ref("CommaSegment")),
    )


############################
# MERGE
############################
@snowflake_dialect.segment(replace=True)
class MergeUpdateClauseSegment(BaseSegment):
    """`UPDATE` clause within the `MERGE` statement."""

    type = "merge_update_clause"
    match_grammar = Sequence(
        "UPDATE",
        Ref("SetClauseListSegment"),
        Ref("WhereClauseSegment", optional=True),
    )


@snowflake_dialect.segment(replace=True)
class MergeDeleteClauseSegment(BaseSegment):
    """`DELETE` clause within the `MERGE` statement."""

    type = "merge_delete_clause"
    match_grammar = Sequence(
        "DELETE",
        Ref("WhereClauseSegment", optional=True),
    )


@snowflake_dialect.segment(replace=True)
class MergeInsertClauseSegment(BaseSegment):
    """`INSERT` clause within the `MERGE` statement."""

    type = "merge_insert_clause"
    match_grammar = Sequence(
        "INSERT",
        Ref("BracketedColumnReferenceListGrammar", optional=True),
        "VALUES",
        Bracketed(
            Delimited(
                OneOf(
                    "DEFAULT",
                    Ref("ExpressionSegment"),
                ),
            )
        ),
        Ref("WhereClauseSegment", optional=True),
    )


@snowflake_dialect.segment(replace=True)
class DeleteStatementSegment(BaseSegment):
    """A `DELETE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/delete.html
    """

    type = "delete_statement"
    match_grammar = Sequence(
        "DELETE",
        "FROM",
        Ref("TableReferenceSegment"),
        Ref("AliasExpressionSegment", optional=True),
        Sequence(
            "USING",
            Indent,
            Delimited(
                Sequence(
                    Ref("TableExpressionSegment"),
                    Ref("AliasExpressionSegment", optional=True),
                ),
            ),
            Dedent,
            optional=True,
        ),
        Ref("WhereClauseSegment", optional=True),
    )


@snowflake_dialect.segment(replace=True)
class DescribeStatementSegment(BaseSegment):
    """`DESCRIBE` statement grammar.

    https://docs.snowflake.com/en/sql-reference/sql/desc.html
    """

    type = "describe_statement"
    match_grammar = Sequence(
        OneOf("DESCRIBE", "DESC"),
        OneOf(
            # https://docs.snowflake.com/en/sql-reference/sql/desc-result.html
            Sequence(
                "RESULT",
                OneOf(
                    Ref("QuotedLiteralSegment"),
                    Sequence("LAST_QUERY_ID", Bracketed()),
                ),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-network-policy.html
            Sequence(
                "NETWORK",
                "POLICY",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-share.html
            Sequence(
                "SHARE",
                Ref("ObjectReferenceSegment"),
                Sequence(
                    Ref("DotSegment"),
                    Ref("ObjectReferenceSegment"),
                    optional=True,
                ),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-user.html
            Sequence(
                "USER",
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "WAREHOUSE",
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "DATABASE",
                Ref("DatabaseReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-integration.html
            Sequence(
                OneOf(
                    "API",
                    "NOTIFICATION",
                    "SECURITY",
                    "STORAGE",
                ),
                "INTEGRATION",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-session-policy.html
            Sequence(
                "SESSION",
                "POLICY",
                Ref("ObjectReferenceSegment"),
            ),
            Sequence(
                "SCHEMA",
                Ref("SchemaReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-table.html
            Sequence(
                "TABLE",
                Ref("TableReferenceSegment"),
                Sequence(
                    "TYPE",
                    Ref("EqualsSegment"),
                    OneOf("COLUMNS", "STAGE"),
                    optional=True,
                ),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-external-table.html
            Sequence(
                "EXTERNAL",
                "TABLE",
                Ref("TableReferenceSegment"),
                Sequence(
                    "TYPE",
                    Ref("EqualsSegment"),
                    OneOf("COLUMNS", "STAGE"),
                    optional=True,
                ),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-view.html
            Sequence(
                "VIEW",
                Ref("TableReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-materialized-view.html
            Sequence(
                "MATERIALIZED",
                "VIEW",
                Ref("TableReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-sequence.html
            Sequence(
                "SEQUENCE",
                Ref("SequenceReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-masking-policy.html
            Sequence(
                "MASKING",
                "POLICY",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-row-access-policy.html
            Sequence(
                "ROW",
                "ACCESS",
                "POLICY",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-file-format.html
            Sequence(
                "FILE",
                "FORMAT",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-stage.html
            Sequence(
                "STAGE",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-pipe.html
            Sequence(
                "PIPE",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-stream.html
            Sequence(
                "STREAM",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-task.html
            Sequence(
                "TASK",
                Ref("ObjectReferenceSegment"),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-function.html
            Sequence(
                "FUNCTION",
                Ref("FunctionNameSegment"),
                Bracketed(
                    Delimited(
                        Ref("DatatypeSegment"),
                        optional=True,
                    ),
                ),
            ),
            # https://docs.snowflake.com/en/sql-reference/sql/desc-procedure.html
            Sequence(
                "PROCEDURE",
                Ref("FunctionNameSegment"),
                Bracketed(
                    Delimited(
                        Ref("DatatypeSegment"),
                        optional=True,
                    ),
                ),
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class TransactionStatementSegment(BaseSegment):
    """`BEGIN`, `START TRANSACTION`, `COMMIT`, AND `ROLLBACK` statement grammar.

    Overwrites ANSI to match correct Snowflake grammar.

    https://docs.snowflake.com/en/sql-reference/sql/begin.html
    https://docs.snowflake.com/en/sql-reference/sql/commit.html
    https://docs.snowflake.com/en/sql-reference/sql/rollback.html
    """

    type = "transaction_statement"
    match_grammar = OneOf(
        Sequence(
            "BEGIN",
            OneOf("WORK", "TRANSACTION", optional=True),
            Sequence("NAME", Ref("ObjectReferenceSegment"), optional=True),
        ),
        Sequence(
            "START",
            "TRANSACTION",
            Sequence("NAME", Ref("ObjectReferenceSegment"), optional=True),
        ),
        "COMMIT",
        "ROLLBACK",
    )


@snowflake_dialect.segment(replace=True)
class TruncateStatementSegment(BaseSegment):
    """`TRUNCATE TABLE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/truncate-table.html
    """

    type = "truncate_table"
    match_grammar = Sequence(
        "TRUNCATE",
        Ref.keyword("TABLE", optional=True),
        Sequence("IF", "EXISTS", optional=True),
        Ref("TableReferenceSegment"),
    )


@snowflake_dialect.segment()
class UnsetStatementSegment(BaseSegment):
    """An `UNSET` statement.

    https://docs.snowflake.com/en/sql-reference/sql/unset.html
    """

    type = "unset_statement"

    match_grammar = Sequence(
        "UNSET",
        OneOf(
            Ref("LocalVariableNameSegment"),
            Bracketed(
                Delimited(
                    Ref("LocalVariableNameSegment"),
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class UndropStatementSegment(BaseSegment):
    """`UNDROP` statement.

    DATABASE: https://docs.snowflake.com/en/sql-reference/sql/undrop-database.html
    SCHEMA: https://docs.snowflake.com/en/sql-reference/sql/undrop-schema.html
    TABLE: https://docs.snowflake.com/en/sql-reference/sql/undrop-table.html
    """

    type = "undrop_statement"
    match_grammar = Sequence(
        "UNDROP",
        OneOf(
            Sequence(
                "DATABASE",
                Ref("DatabaseReferenceSegment"),
            ),
            Sequence(
                "SCHEMA",
                Ref("SchemaReferenceSegment"),
            ),
            Sequence(
                "TABLE",
                Ref("TableReferenceSegment"),
            ),
        ),
    )


@snowflake_dialect.segment()
class CommentStatementSegment(BaseSegment):
    """`COMMENT` statement grammar.

    https://docs.snowflake.com/en/sql-reference/sql/comment.html

    N.B. this applies to all objects, so there may be some I've missed
    here so add any others to the OneOf grammar below.
    """

    type = "comment_statement"
    match_grammar = Sequence(
        "COMMENT",
        Sequence(
            "IF",
            "EXISTS",
            optional=True,
        ),
        "ON",
        OneOf(
            "COLUMN",
            "TABLE",
            "VIEW",
            "SCHEMA",
            "DATABASE",
            "WAREHOUSE",
            "USER",
            "STAGE",
            "FUNCTION",
            "PROCEDURE",
            "SEQUENCE",
            "SHARE",
            "PIPE",
            "STREAM",
            "TASK",
            Sequence(
                "NETWORK",
                "POLICY",
            ),
            Sequence(
                OneOf(
                    "API",
                    "NOTIFICATION",
                    "SECURITY",
                    "STORAGE",
                ),
                "INTEGRATION",
            ),
            Sequence(
                "SESSION",
                "POLICY",
            ),
            Sequence(
                "EXTERNAL",
                "TABLE",
            ),
            Sequence(
                "MATERIALIZED",
                "VIEW",
            ),
            Sequence(
                "MASKING",
                "POLICY",
            ),
            Sequence(
                "ROW",
                "ACCESS",
                "POLICY",
            ),
            Sequence(
                "FILE",
                "FORMAT",
            ),
        ),
        Ref("ObjectReferenceSegment"),
        "IS",
        Ref("QuotedLiteralSegment"),
    )


@snowflake_dialect.segment(replace=True)
class UseStatementSegment(BaseSegment):
    """A `USE` statement.

    https://docs.snowflake.com/en/sql-reference/sql/use.html
    """

    type = "use_statement"
    match_grammar = Sequence(
        "USE",
        OneOf(
            Sequence("ROLE", Ref("ObjectReferenceSegment")),
            Sequence("WAREHOUSE", Ref("ObjectReferenceSegment")),
            Sequence(
                Ref.keyword("DATABASE", optional=True),
                Ref("DatabaseReferenceSegment"),
            ),
            Sequence(
                Ref.keyword("SCHEMA", optional=True),
                Ref("SchemaReferenceSegment"),
            ),
            Sequence(
                "SECONDARY",
                "ROLES",
                OneOf(
                    "ALL",
                    "NONE",
                ),
            ),
        ),
    )


@snowflake_dialect.segment()
class CallStatementSegment(BaseSegment):
    """`CALL` statement.

    https://docs.snowflake.com/en/sql-reference/sql/call.html
    """

    type = "call_statement"
    match_grammar = Sequence(
        "CALL",
        Sequence(
            Ref("FunctionNameSegment"),
            Bracketed(
                Ref(
                    "FunctionContentsGrammar",
                    # The brackets might be empty for some functions...
                    optional=True,
                    ephemeral_name="FunctionContentsGrammar",
                ),
            ),
        ),
    )


@snowflake_dialect.segment(replace=True)
class LimitClauseSegment(BaseSegment):
    """A `LIMIT` clause.

    https://docs.snowflake.com/en/sql-reference/constructs/limit.html
    """

    type = "limit_clause"
    match_grammar = OneOf(
        Sequence(
            "LIMIT",
            Indent,
            Ref("LimitLiteralGrammar"),
            Dedent,
            Sequence(
                "OFFSET",
                Indent,
                Ref("LimitLiteralGrammar"),
                Dedent,
                optional=True,
            ),
        ),
        Sequence(
            Sequence(
                "OFFSET",
                Indent,
                Ref("LimitLiteralGrammar"),
                OneOf(
                    "ROW",
                    "ROWS",
                    optional=True,
                ),
                Dedent,
                optional=True,
            ),
            "FETCH",
            Indent,
            OneOf(
                "FIRST",
                "NEXT",
                optional=True,
            ),
            Ref("LimitLiteralGrammar"),
            OneOf(
                "ROW",
                "ROWS",
                optional=True,
            ),
            Ref.keyword("ONLY", optional=True),
            Dedent,
        ),
    )


@snowflake_dialect.segment(replace=True)
class SelectClauseSegment(BaseSegment):
    """A group of elements in a select target statement."""

    type = "select_clause"
    match_grammar = ansi_dialect.get_segment("SelectClauseSegment").match_grammar.copy()
    match_grammar.terminator = match_grammar.terminator.copy(
        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET")],
    )
    parse_grammar = ansi_dialect.get_segment("SelectClauseSegment").parse_grammar.copy()


@snowflake_dialect.segment(replace=True)
class OrderByClauseSegment(BaseSegment):
    """An `ORDER BY` clause.

    https://docs.snowflake.com/en/sql-reference/constructs/order-by.html
    """

    type = "orderby_clause"
    match_grammar = ansi_dialect.get_segment(
        "OrderByClauseSegment"
    ).match_grammar.copy()
    match_grammar.terminator = match_grammar.terminator.copy(
        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET"), Ref.keyword("MEASURES")],
    )
    parse_grammar = Sequence(
        "ORDER",
        "BY",
        Indent,
        Delimited(
            Sequence(
                OneOf(
                    Ref("ColumnReferenceSegment"),
                    # Can `ORDER BY 1`
                    Ref("NumericLiteralSegment"),
                    # Can order by an expression
                    Ref("ExpressionSegment"),
                ),
                OneOf("ASC", "DESC", optional=True),
                Sequence("NULLS", OneOf("FIRST", "LAST"), optional=True),
            ),
            terminator=OneOf("LIMIT", "FETCH", "OFFSET", Ref("FrameClauseUnitGrammar")),
        ),
        Dedent,
    )


@snowflake_dialect.segment(replace=True)
class HavingClauseSegment(BaseSegment):
    """A `HAVING` clause."""

    type = "having_clause"
    match_grammar = ansi_dialect.get_segment("HavingClauseSegment").match_grammar.copy()
    match_grammar.terminator = match_grammar.terminator.copy(
        insert=[Ref.keyword("FETCH"), Ref.keyword("OFFSET")],
    )
    parse_grammar = ansi_dialect.get_segment("HavingClauseSegment").parse_grammar.copy()
