from __future__ import annotations

# Map standard SQL/DDL types to our desired Dask/Pandas working schema.
# We prioritize PyArrow-backed types for performance and better null handling.
SQLALCHEMY_TO_DASK_DTYPE: dict[str, str] = {
    # Integers: Use Pandas extension types (e.g., Int64) for nullability
    "INTEGER": "Int64[pyarrow]",
    "SMALLINT": "Int32[pyarrow]",
    "BIGINT": "Int64[pyarrow]",
    # Floats: Use Pandas extension types (Float64) for nullability
    "FLOAT": "Float64[pyarrow]",
    "DOUBLE": "Float64[pyarrow]",
    # Exact numbers: String is the safest choice to avoid floating point error
    "NUMERIC": "string[pyarrow]",
    "DECIMAL": "string[pyarrow]",
    # Boolean: Use Pandas extension type for nullability
    "BOOLEAN": "boolean[pyarrow]",
    # Strings: Use PyArrow-backed string (efficient storage/zero-copy potential)
    "VARCHAR": "string[pyarrow]",
    "CHAR": "string[pyarrow]",
    "TEXT": "string[pyarrow]",
    "STRING": "string[pyarrow]",
    "UUID": "string[pyarrow]",
    # Dates/Times: Enforce Standard Pandas UTC Dtype for consistency
    "DATE": "datetime64[ns, UTC]",
    "DATETIME": "datetime64[ns, UTC]",
    "TIMESTAMP": "datetime64[ns, UTC]",
    "TIME": "string[pyarrow]",
}

DASK_TO_CLICKHOUSE_DTYPE: dict[str, str] = {
    # --- 1. PyArrow-Backed Integer Types (Nullable) ---
    # These cover the Pandas extension type (Capital 'I') and the raw PyArrow string alias (lowercase 'i').
    "Int64[pyarrow]": "Nullable(Int64)",
    "Int32[pyarrow]": "Nullable(Int32)",
    "int64[pyarrow]": "Nullable(Int64)",  # Alias found in logs
    "int32[pyarrow]": "Nullable(Int32)",  # Alias found in logs
    # --- 2. PyArrow-Backed Float Types (Nullable) ---
    # These cover the Pandas extension type (Capital 'F') and common PyArrow string aliases.
    "Float64[pyarrow]": "Nullable(Float64)",
    "Float32[pyarrow]": "Nullable(Float32)",
    "double[pyarrow]": "Nullable(Float64)",  # Alias found in logs
    "float64[pyarrow]": "Nullable(Float64)",
    # --- 3. PyArrow-Backed String and Boolean Types (Nullable) ---
    "boolean[pyarrow]": "Nullable(Bool)",
    "string[pyarrow]": "String",  # ClickHouse String is inherently nullable
    # --- 4. PyArrow Timestamp and Datetime Types ---
    # We map both Pandas canonical datetime types and the raw PyArrow timestamp aliases.
    "datetime64[ns, UTC]": "DateTime64(9, 'UTC')",
    "datetime64[ns]": "DateTime64(9)",
    "timestamp[ns][pyarrow]": "DateTime64(9, 'UTC')",  # Alias found in logs
    "timestamp[us][pyarrow]": "DateTime64(6, 'UTC')",  # Microsecond precision
    "timestamp[ms][pyarrow]": "DateTime64(3, 'UTC')",  # Millisecond precision
    # --- 5. Standard Pandas/NumPy Dtypes (Fallback for non-PyArrow) ---
    # These types are non-nullable in NumPy, so we map them to non-Nullable CH types.
    # Note: If these appear, it might indicate an issue earlier in the pipeline.
    "int64": "Int64",
    "int32": "Int32",
    "float64": "Float64",
    "float32": "Float32",
    # Generic types that must default to String
    "object": "String",
    "category": "String",  # Base type for LowCardinality wrapper in DDL logic
    "bool": "Bool",
}

# Map intended data type to the Dask/Pandas type needed for safe, fast *ingestion*
CSV_INGESTION_DTYPE: dict[str, str] = {
    # Integers: Must use Pandas Extension Type to handle NULLs,
    # preventing the column from becoming float64.
    "INT_WITH_NULLS": "Int64[pyarrow]",
    "SMALL_INT_WITH_NULLS": "Int32[pyarrow]",
    # Floats: Standard float is usually fine for ingestion
    "FLOAT": "float64[pyarrow]",
    # High-Precision Decimals/Strings: Use PyArrow-backed string
    "DECIMAL_AS_TEXT": "string[pyarrow]",
    "STRING": "string[pyarrow]",
    "TEXT": "string[pyarrow]",
    # Dates: Read as object/string first, clean later in the pipeline
    "DATE_OR_DATETIME": "object",
    "TIMESTAMP": "object",
    # Boolean
    "BOOLEAN_WITH_NULLS": "boolean[pyarrow]",
}

DASK_TO_SQLALCHEMY_DTYPE: dict[str, str] = {
    # Integers: Prefer BIGINT as it accommodates both INTEGER and BIGINT sizes (safer default)
    "Int64[pyarrow]": "BIGINT",
    "Int32[pyarrow]": "INTEGER",
    # Floats: Prefer DOUBLE as it offers the highest precision (safer default)
    "Float64[pyarrow]": "DOUBLE",
    # Strings: Prefer VARCHAR (or TEXT) as a general-purpose string type
    # If the precision was critical (NUMERIC/DECIMAL), you'd usually write to TEXT
    "string[pyarrow]": "TEXT",
    "category": "VARCHAR",  # If exporting a categorical type
    # Boolean
    "boolean[pyarrow]": "BOOLEAN",
    # Dates/Times:
    "datetime64[ns, UTC]": "TIMESTAMP",
    "datetime64[ns]": "TIMESTAMP",  # Handle naive timestamp fallback
}
