from __future__ import annotations

from sqlglot import exp
from sqlglot.dialects.dialect import (
    approx_count_distinct_sql,
    build_timestamp_trunc,
    rename_func,
    time_format,
    unit_to_str,
)
from sqlglot.dialects.mysql import MySQL


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 Doris(MySQL):
    DATE_FORMAT = "'yyyy-MM-dd'"
    DATEINT_FORMAT = "'yyyyMMdd'"
    TIME_FORMAT = "'yyyy-MM-dd HH:mm:ss'"

    class Parser(MySQL.Parser):
        FUNCTIONS = {
            **MySQL.Parser.FUNCTIONS,
            "COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
            "DATE_TRUNC": build_timestamp_trunc,
            "MONTHS_ADD": exp.AddMonths.from_arg_list,
            "REGEXP": exp.RegexpLike.from_arg_list,
            "TO_DATE": exp.TsOrDsToDate.from_arg_list,
        }

        FUNCTION_PARSERS = MySQL.Parser.FUNCTION_PARSERS.copy()
        FUNCTION_PARSERS.pop("GROUP_CONCAT")

    class Generator(MySQL.Generator):
        LAST_DAY_SUPPORTS_DATE_PART = False
        VARCHAR_REQUIRES_SIZE = False

        TYPE_MAPPING = {
            **MySQL.Generator.TYPE_MAPPING,
            exp.DataType.Type.TEXT: "STRING",
            exp.DataType.Type.TIMESTAMP: "DATETIME",
            exp.DataType.Type.TIMESTAMPTZ: "DATETIME",
        }

        CAST_MAPPING = {}
        TIMESTAMP_FUNC_TYPES = set()

        TRANSFORMS = {
            **MySQL.Generator.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.CurrentTimestamp: lambda self, _: self.func("NOW"),
            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, unit_to_str(e)),
            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.RegexpLike: rename_func("REGEXP"),
            exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
            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, unit_to_str(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",
        }
