from __future__ import annotations
from typing import Optional
import polars as pl


class TypeMapper:
    """Maps Alteryx field types to Polars dtypes and DuckDB type strings."""

    _POLARS: dict[str, pl.PolarsDataType] = {
        "Bool": pl.Boolean,
        "Byte": pl.UInt8,
        "Int16": pl.Int16,
        "Int32": pl.Int32,
        "Int64": pl.Int64,
        "Float": pl.Float64,
        "Double": pl.Float64,
        "String": pl.String,
        "V_String": pl.String,
        "WString": pl.String,
        "V_WString": pl.String,
        "Date": pl.Date,
        "Time": pl.Time,
        "DateTime": pl.Datetime,
        "SpatialObj": pl.String,
        "Blob": pl.Binary,
    }

    _DUCKDB: dict[str, str] = {
        "Bool": "BOOLEAN",
        "Byte": "UTINYINT",
        "Int16": "SMALLINT",
        "Int32": "INTEGER",
        "Int64": "BIGINT",
        "Float": "DOUBLE",
        "Double": "DOUBLE",
        "String": "VARCHAR",
        "V_String": "VARCHAR",
        "WString": "VARCHAR",
        "V_WString": "VARCHAR",
        "Date": "DATE",
        "Time": "TIME",
        "DateTime": "TIMESTAMP",
        "SpatialObj": "VARCHAR",
        "Blob": "BLOB",
        "FixedDecimal": "DECIMAL",
    }

    def map(self, alteryx_type: str, size: Optional[str] = None) -> pl.PolarsDataType:
        if alteryx_type == "FixedDecimal":
            if size:
                parts = size.split(".")
                precision = int(parts[0])
                scale = int(parts[1]) if len(parts) > 1 else 0
                return pl.Decimal(precision=precision, scale=scale)
            return pl.Decimal(precision=19, scale=2)
        return self._POLARS.get(alteryx_type, pl.String)

    def map_duckdb(self, alteryx_type: str, size: Optional[str] = None) -> str:
        if alteryx_type == "FixedDecimal":
            if size:
                parts = size.split(".")
                precision = int(parts[0])
                scale = int(parts[1]) if len(parts) > 1 else 0
                return f"DECIMAL({precision},{scale})"
            return "DECIMAL(19,2)"
        return self._DUCKDB.get(alteryx_type, "VARCHAR")
