from __future__ import annotations

import logging
from typing import Dict, Union

from typing_extensions import Literal

from graph.ingestion.dataiku.data_types import GraphDataikuDataType, NativeDataikuDataType

logger = logging.getLogger(__name__)

# All types defined in Kuzu: https://kuzudb.github.io/docs/cypher/data-types/
BooleanAliases = Union[
    Literal["BOOLEAN"],
    Literal[
        "BOOL"
    ],  # Not explicitly documented but appears in kuzu schema for boolean columns. Values are "true"/"false".
]
IntervalAliases = Union[
    Literal["INTERVAL"],
    Literal["DURATION"],
]
KuzuDataType = Union[
    Literal["INT8"],
    Literal["INT16"],
    Literal["INT32"],
    Literal["INT64"],
    Literal["INT128"],
    Literal["UINT8"],
    Literal["UINT16"],
    Literal["UINT32"],
    Literal["UINT64"],
    Literal["FLOAT"],
    Literal["DOUBLE"],
    Literal["DECIMAL"],
    BooleanAliases,
    Literal["UUID"],
    Literal["STRING"],
    Literal["NULL"],
    Literal["DATE"],
    Literal["TIMESTAMP"],
    IntervalAliases,
    Literal["STRUCT"],
    Literal["MAP"],
    Literal["UNION"],
    Literal["BLOB"],
    Literal["SERIAL"],
    Literal["NODE"],
    Literal["REL"],
    Literal["RECURSIVE_REL"],
    # Literal["LIST"], # This type does not exist as such, it is typed like this: INT64[]/NODE[] for an array of int64/nodes. Specific dev effort required.
    # Literal["ARRAY"], # This type does not exist as such, it is typed like this: INT64[3]/NODE[3] for an array of 3 int64/nodes. Specific dev effort required.
]

# ALL dataiku types mapped to kuzu types.
TYPE_MAP_DATAIKU_TO_KUZU: Dict[NativeDataikuDataType, KuzuDataType] = {
    # Native Dataiku types
    "string": "STRING",
    "boolean": "BOOLEAN",
    "date": "TIMESTAMP",  # https://kuzudb.github.io/docs/cypher/data-types/#timestamp todo: there might be issues with formatting as kuzu expects ISO-8601 format (YYYY-MM-DD hh:mm:ss[.zzzzzz][+-TT[:tt]])
    "datetimenotz": "TIMESTAMP",  # https://kuzudb.github.io/docs/cypher/data-types/#timestamp kuzu defaults to UTC timezone
    "dateonly": "DATE",  # https://kuzudb.github.io/docs/cypher/data-types/#date ISO-8601 format (YYYY-MM-DD)
    "tinyint": "INT8",
    "smallint": "INT16",
    "int": "INT32",
    "bigint": "INT64",
    "float": "FLOAT",
    "double": "DOUBLE",
    # The following are badly supported Dataiku types.
    # Some parts of the kuzu documentation are unavailable (LIST/ARRAY for example).
    "object": "STRING",  # STRING to simplify, JSON with installed extension in kuzu might work, requires dev.
    "array": "STRING",  # STRING to simplify, JSON with installed extension in kuzu might work, requires dev.
    "map": "STRING",  # Ingestion structure is unclear, probably necessitates custom developments.
    # The following types are not supported.
    # "geometry"
    # "geopoint"
}

# ALL kuzu types mapped to dataiku types.
_TYPE_MAP_KUZU_TO_DATAIKU: Dict[KuzuDataType, GraphDataikuDataType] = {
    "INT8": "tinyint",
    "INT16": "smallint",
    "INT32": "int",
    "INT64": "bigint",
    "INT128": "hugeint",  # Does not exist in Dataiku but we still surface this type. Consumer will have to decide how to handle it (e.g. map to string and parse as needed, or map to bigint and risk overflow).
    "UINT8": "smallint",
    "UINT16": "int",
    "UINT32": "bigint",
    "UINT64": "hugeint",  # Does not exist in Dataiku but we still surface this type. Consumer will have to decide how to handle it (e.g. map to string and parse as needed, or map to bigint and risk overflow).
    "FLOAT": "float",
    "DOUBLE": "double",
    "DECIMAL": "float",  # No decimal type in dataiku, mapping to float. Might cause precision issues for very large/small numbers or numbers with many decimals.
    "BOOL": "boolean",  # Not documented but appears in kuzu schema for boolean columns. Values are "true"/"false".
    "BOOLEAN": "boolean",
    "UUID": "string",  # no uuid type in dataiku, mapping to string
    "STRING": "string",
    "NULL": "string",  # no null type in dataiku, mapping to string
    "DATE": "dateonly",
    "TIMESTAMP": "date",
    "INTERVAL": "string",  # no interval type in dataiku, mapping to string.
    "DURATION": "string",  # no interval type in dataiku, mapping to string.
    "STRUCT": "object",  # Could happen if user returns using STRUCT_PACK() function.
    "MAP": "map",  # Could happen if user returns using map() function.
    "UNION": "object",  # I dont see how it could happen as we dont map to this type during ingestion. Edge case is if a user creates its own kuzu db -> unlikely
    "BLOB": "object",  # I dont see how it could happen as we dont map to this type during ingestion. Edge case is if a user creates its own kuzu db -> unlikely
    "SERIAL": "bigint",
    "NODE": "node",
    "REL": "edge",
    "RECURSIVE_REL": "path",  # Requires specific devlopments. For ex: MATCH p = (a:User)-[:Follows]->(b:User)WHERE a.name = 'Adam' AND b.name = 'Karissa' RETURN p
}


def map_kuzu_to_dataiku_type(kuzu_type: str) -> GraphDataikuDataType:
    normalized_type = kuzu_type.strip().upper()

    if normalized_type in _TYPE_MAP_KUZU_TO_DATAIKU:
        return _TYPE_MAP_KUZU_TO_DATAIKU[normalized_type]  # type: ignore[index]  # Ignore required for mypy + py39.

    if normalized_type.startswith("DECIMAL("):
        return _TYPE_MAP_KUZU_TO_DATAIKU["DECIMAL"]

    if normalized_type.startswith(("STRUCT(")):
        return _TYPE_MAP_KUZU_TO_DATAIKU["STRUCT"]

    if normalized_type.startswith("MAP("):
        return _TYPE_MAP_KUZU_TO_DATAIKU["MAP"]

    if normalized_type.startswith("UNION("):
        return _TYPE_MAP_KUZU_TO_DATAIKU["UNION"]

    logger.warning(f"Badly supported Kuzu '{kuzu_type}' type. Mapping to object.")

    return "object"
