from __future__ import annotations

import typing as t

from sqlglot import exp
from sqlglot.dialects.dialect import (
    approx_count_distinct_sql,
    property_sql,
    rename_func,
    time_format,
    weekstart_unit_to_str,
)
from sqlglot.generators.mysql import MySQLGenerator


def _lag_lead_sql(self, expression: exp.Lag | exp.Lead) -> str:
    return self.func(
        "LAG" if isinstance(expression, exp.Lag) else "LEAD",
        expression.this,
        expression.args.get("offset") or exp.Literal.number(1),
        expression.args.get("default") or exp.null(),
    )


class DorisGenerator(MySQLGenerator):
    LAST_DAY_SUPPORTS_DATE_PART = False
    SUPPORTS_ALTER_COLUMN_NULLABILITY = False
    VARCHAR_REQUIRES_SIZE = False
    WITH_PROPERTIES_PREFIX = "PROPERTIES"
    RENAME_TABLE_WITH_DB = False
    UPDATE_STATEMENT_SUPPORTS_FROM = True

    TYPE_MAPPING = {
        **MySQLGenerator.TYPE_MAPPING,
        exp.DType.TEXT: "STRING",
        exp.DType.TIMESTAMP: "DATETIME",
        exp.DType.TIMESTAMPTZ: "DATETIME",
    }

    PROPERTIES_LOCATION = {
        **MySQLGenerator.PROPERTIES_LOCATION,
        exp.UniqueKeyProperty: exp.Properties.Location.POST_SCHEMA,
        exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
        exp.BuildProperty: exp.Properties.Location.POST_SCHEMA,
    }

    CAST_MAPPING: t.ClassVar[dict[exp.DType, str]] = {}
    TIMESTAMP_FUNC_TYPES: t.ClassVar[set[exp.DType]] = set()

    TRANSFORMS = {
        **MySQLGenerator.TRANSFORMS,
        exp.AddMonths: rename_func("MONTHS_ADD"),
        exp.ApproxDistinct: approx_count_distinct_sql,
        exp.ArgMax: rename_func("MAX_BY"),
        exp.ArgMin: rename_func("MIN_BY"),
        exp.ArrayAgg: rename_func("COLLECT_LIST"),
        exp.ArrayToString: rename_func("ARRAY_JOIN"),
        exp.ArrayUniqueAgg: rename_func("COLLECT_SET"),
        exp.CurrentDate: lambda self, _: self.func("CURRENT_DATE"),
        exp.CurrentTimestamp: lambda self, _: self.func("NOW"),
        exp.DateTrunc: lambda self, e: self.func(
            "DATE_TRUNC", e.this, weekstart_unit_to_str(self, e)
        ),
        exp.EuclideanDistance: rename_func("L2_DISTANCE"),
        exp.GroupConcat: lambda self, e: self.func(
            "GROUP_CONCAT", e.this, e.args.get("separator") or exp.Literal.string(",")
        ),
        exp.JSONExtractScalar: lambda self, e: self.func("JSON_EXTRACT", e.this, e.expression),
        exp.Lag: _lag_lead_sql,
        exp.Lead: _lag_lead_sql,
        exp.Map: rename_func("ARRAY_MAP"),
        exp.Property: property_sql,
        exp.RegexpLike: rename_func("REGEXP"),
        exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
        exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
        exp.Split: rename_func("SPLIT_BY_STRING"),
        exp.StringToArray: rename_func("SPLIT_BY_STRING"),
        exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
        exp.TimeStrToDate: rename_func("TO_DATE"),
        exp.TsOrDsAdd: lambda self, e: self.func("DATE_ADD", e.this, e.expression),
        exp.TsOrDsToDate: lambda self, e: self.func("TO_DATE", e.this),
        exp.TimeToUnix: rename_func("UNIX_TIMESTAMP"),
        exp.TimestampTrunc: lambda self, e: self.func(
            "DATE_TRUNC", e.this, weekstart_unit_to_str(self, e)
        ),
        exp.UnixToStr: lambda self, e: self.func(
            "FROM_UNIXTIME", e.this, time_format("doris")(self, e)
        ),
        exp.UnixToTime: rename_func("FROM_UNIXTIME"),
    }

    # https://github.com/apache/doris/blob/e4f41dbf1ec03f5937fdeba2ee1454a20254015b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4#L93
    RESERVED_KEYWORDS = {
        "account_lock",
        "account_unlock",
        "add",
        "adddate",
        "admin",
        "after",
        "agg_state",
        "aggregate",
        "alias",
        "all",
        "alter",
        "analyze",
        "analyzed",
        "and",
        "anti",
        "append",
        "array",
        "array_range",
        "as",
        "asc",
        "at",
        "authors",
        "auto",
        "auto_increment",
        "backend",
        "backends",
        "backup",
        "begin",
        "belong",
        "between",
        "bigint",
        "bin",
        "binary",
        "binlog",
        "bitand",
        "bitmap",
        "bitmap_union",
        "bitor",
        "bitxor",
        "blob",
        "boolean",
        "brief",
        "broker",
        "buckets",
        "build",
        "builtin",
        "bulk",
        "by",
        "cached",
        "call",
        "cancel",
        "case",
        "cast",
        "catalog",
        "catalogs",
        "chain",
        "char",
        "character",
        "charset",
        "check",
        "clean",
        "cluster",
        "clusters",
        "collate",
        "collation",
        "collect",
        "column",
        "columns",
        "comment",
        "commit",
        "committed",
        "compact",
        "complete",
        "config",
        "connection",
        "connection_id",
        "consistent",
        "constraint",
        "constraints",
        "convert",
        "copy",
        "count",
        "create",
        "creation",
        "cron",
        "cross",
        "cube",
        "current",
        "current_catalog",
        "current_date",
        "current_time",
        "current_timestamp",
        "current_user",
        "data",
        "database",
        "databases",
        "date",
        "date_add",
        "date_ceil",
        "date_diff",
        "date_floor",
        "date_sub",
        "dateadd",
        "datediff",
        "datetime",
        "datetimev2",
        "datev2",
        "datetimev1",
        "datev1",
        "day",
        "days_add",
        "days_sub",
        "decimal",
        "decimalv2",
        "decimalv3",
        "decommission",
        "default",
        "deferred",
        "delete",
        "demand",
        "desc",
        "describe",
        "diagnose",
        "disk",
        "distinct",
        "distinctpc",
        "distinctpcsa",
        "distributed",
        "distribution",
        "div",
        "do",
        "doris_internal_table_id",
        "double",
        "drop",
        "dropp",
        "dual",
        "duplicate",
        "dynamic",
        "else",
        "enable",
        "encryptkey",
        "encryptkeys",
        "end",
        "ends",
        "engine",
        "engines",
        "enter",
        "errors",
        "events",
        "every",
        "except",
        "exclude",
        "execute",
        "exists",
        "expired",
        "explain",
        "export",
        "extended",
        "external",
        "extract",
        "failed_login_attempts",
        "false",
        "fast",
        "feature",
        "fields",
        "file",
        "filter",
        "first",
        "float",
        "follower",
        "following",
        "for",
        "foreign",
        "force",
        "format",
        "free",
        "from",
        "frontend",
        "frontends",
        "full",
        "function",
        "functions",
        "generic",
        "global",
        "grant",
        "grants",
        "graph",
        "group",
        "grouping",
        "groups",
        "hash",
        "having",
        "hdfs",
        "help",
        "histogram",
        "hll",
        "hll_union",
        "hostname",
        "hour",
        "hub",
        "identified",
        "if",
        "ignore",
        "immediate",
        "in",
        "incremental",
        "index",
        "indexes",
        "infile",
        "inner",
        "insert",
        "install",
        "int",
        "integer",
        "intermediate",
        "intersect",
        "interval",
        "into",
        "inverted",
        "ipv4",
        "ipv6",
        "is",
        "is_not_null_pred",
        "is_null_pred",
        "isnull",
        "isolation",
        "job",
        "jobs",
        "join",
        "json",
        "jsonb",
        "key",
        "keys",
        "kill",
        "label",
        "largeint",
        "last",
        "lateral",
        "ldap",
        "ldap_admin_password",
        "left",
        "less",
        "level",
        "like",
        "limit",
        "lines",
        "link",
        "list",
        "load",
        "local",
        "localtime",
        "localtimestamp",
        "location",
        "lock",
        "logical",
        "low_priority",
        "manual",
        "map",
        "match",
        "match_all",
        "match_any",
        "match_phrase",
        "match_phrase_edge",
        "match_phrase_prefix",
        "match_regexp",
        "materialized",
        "max",
        "maxvalue",
        "memo",
        "merge",
        "migrate",
        "migrations",
        "min",
        "minus",
        "minute",
        "modify",
        "month",
        "mtmv",
        "name",
        "names",
        "natural",
        "negative",
        "never",
        "next",
        "ngram_bf",
        "no",
        "non_nullable",
        "not",
        "null",
        "nulls",
        "observer",
        "of",
        "offset",
        "on",
        "only",
        "open",
        "optimized",
        "or",
        "order",
        "outer",
        "outfile",
        "over",
        "overwrite",
        "parameter",
        "parsed",
        "partition",
        "partitions",
        "password",
        "password_expire",
        "password_history",
        "password_lock_time",
        "password_reuse",
        "path",
        "pause",
        "percent",
        "period",
        "permissive",
        "physical",
        "plan",
        "process",
        "plugin",
        "plugins",
        "policy",
        "preceding",
        "prepare",
        "primary",
        "proc",
        "procedure",
        "processlist",
        "profile",
        "properties",
        "property",
        "quantile_state",
        "quantile_union",
        "query",
        "quota",
        "random",
        "range",
        "read",
        "real",
        "rebalance",
        "recover",
        "recycle",
        "refresh",
        "references",
        "regexp",
        "release",
        "rename",
        "repair",
        "repeatable",
        "replace",
        "replace_if_not_null",
        "replica",
        "repositories",
        "repository",
        "resource",
        "resources",
        "restore",
        "restrictive",
        "resume",
        "returns",
        "revoke",
        "rewritten",
        "right",
        "rlike",
        "role",
        "roles",
        "rollback",
        "rollup",
        "routine",
        "row",
        "rows",
        "s3",
        "sample",
        "schedule",
        "scheduler",
        "schema",
        "schemas",
        "second",
        "select",
        "semi",
        "sequence",
        "serializable",
        "session",
        "set",
        "sets",
        "shape",
        "show",
        "signed",
        "skew",
        "smallint",
        "snapshot",
        "soname",
        "split",
        "sql_block_rule",
        "start",
        "starts",
        "stats",
        "status",
        "stop",
        "storage",
        "stream",
        "streaming",
        "string",
        "struct",
        "subdate",
        "sum",
        "superuser",
        "switch",
        "sync",
        "system",
        "table",
        "tables",
        "tablesample",
        "tablet",
        "tablets",
        "task",
        "tasks",
        "temporary",
        "terminated",
        "text",
        "than",
        "then",
        "time",
        "timestamp",
        "timestampadd",
        "timestampdiff",
        "tinyint",
        "to",
        "transaction",
        "trash",
        "tree",
        "triggers",
        "trim",
        "true",
        "truncate",
        "type",
        "type_cast",
        "types",
        "unbounded",
        "uncommitted",
        "uninstall",
        "union",
        "unique",
        "unlock",
        "unsigned",
        "update",
        "use",
        "user",
        "using",
        "value",
        "values",
        "varchar",
        "variables",
        "variant",
        "vault",
        "verbose",
        "version",
        "view",
        "warnings",
        "week",
        "when",
        "where",
        "whitelist",
        "with",
        "work",
        "workload",
        "write",
        "xor",
        "year",
    }

    def uniquekeyproperty_sql(
        self, expression: exp.UniqueKeyProperty, prefix: str = "UNIQUE KEY"
    ) -> str:
        create_stmt = expression.find_ancestor(exp.Create)
        if create_stmt and create_stmt.args["properties"].find(exp.MaterializedProperty):
            return super().uniquekeyproperty_sql(expression, prefix="KEY")

        return super().uniquekeyproperty_sql(expression)

    def partitionrange_sql(self, expression: exp.PartitionRange) -> str:
        name = self.sql(expression, "this")
        values = expression.expressions

        if len(values) != 1:
            # Multiple values: use VALUES [ ... )
            if values and isinstance(values[0], list):
                values_sql = ", ".join(
                    f"({', '.join(self.sql(v) for v in inner)})" for inner in values
                )
            else:
                values_sql = ", ".join(f"({self.sql(v)})" for v in values)

            return f"PARTITION {name} VALUES [{values_sql})"

        return f"PARTITION {name} VALUES LESS THAN ({self.sql(values[0])})"

    def partitionbyrangepropertydynamic_sql(
        self, expression: exp.PartitionByRangePropertyDynamic
    ) -> str:
        # Generates: FROM ("start") TO ("end") INTERVAL N UNIT
        start = self.sql(expression, "start")
        end = self.sql(expression, "end")
        every = expression.args.get("every")

        if every:
            number = self.sql(every, "this")
            interval = f"INTERVAL {number} {self.sql(every, 'unit')}"
        else:
            interval = ""

        return f"FROM ({start}) TO ({end}) {interval}"

    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
        this = expression.this
        if isinstance(this, exp.Schema):
            return f"PARTITION BY ({self.expressions(this, flat=True)})"
        return f"PARTITION BY ({self.sql(this)})"

    def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
        """Override table_sql to avoid AS keyword in UPDATE and DELETE statements."""
        ancestor = expression.find_ancestor(exp.Update, exp.Delete, exp.Select)
        if not isinstance(ancestor, exp.Select):
            sep = " "
        return super().table_sql(expression, sep=sep)
