from collections import defaultdict
import json
import os
from pathlib import Path
import textwrap

from relationalai.config.constants import RAI_BUILD_DIR_ENV_VAR
from relationalai.util.error import source, warn, SourcePos
from relationalai.util.identifiers import normalize_qualified_name, quote_identifier_if_needed, split_qualified_name, unquote_identifier

#------------------------------------------------------
# General
#------------------------------------------------------

_BUILD_DIR = Path(os.environ.get(RAI_BUILD_DIR_ENV_VAR, "build"))
CACHE_PATH = _BUILD_DIR / "cache" / "schemas.json"
CACHE = {}

def _convert_types(result: dict[str, type | str]) -> dict[str, str]:
    """Convert Python types to string representations for the cache."""
    type_map = {str: "String", int: "Int64", float: "Float", bool: "Bool", bytes: "Binary"}
    return {
        col: type_map.get(t, "String") if isinstance(t, type) else t
        for col, t in result.items()
    }

def fetch(fqn: str, model=None) -> tuple[bool, dict[str, str]]:
    """
    Fetch column metadata for a fully-qualified table name.

    Args:
        fqn: Fully-qualified table name (database.schema.table)
        model: Optional Model instance. Uses the model's executor for schema
               fetching when in install_mode, otherwise falls back to Snowflake.
    """
    global CACHE
    fqn = normalize_qualified_name(fqn)
    parts = split_qualified_name(fqn)
    if len(parts) != 3:
        raise ValueError(f"Expected a fully-qualified table name with 3 parts, got: {fqn!r}")
    database, schema, table = parts
    table_name = unquote_identifier(table)

    # Load or initialize cache
    if not CACHE and CACHE_PATH.exists():
        with CACHE_PATH.open("r", encoding="utf-8") as f:
            str_cache: dict[str, dict[str, str]] = json.load(f)
            CACHE = {k: {ck: cv for ck, cv in v.items()} for k, v in str_cache.items()}

    # If already cached, return it
    if fqn in CACHE:
        return True, CACHE[fqn]

    # In install_mode, use the model's Executor (supports DuckDB and Snowflake)
    if model is not None:
        from ..semantics.backends.sql_executor import SQLExecutor
        executor = model._get_executor()
        if isinstance(executor, SQLExecutor):
            result = executor.dialect_executor.fetch_schema(database, schema, table_name)
            if result:
                CACHE[fqn] = _convert_types(result)
                _save_cache()
                return False, CACHE[fqn]

    # Fall back to Snowflake provider
    config = model.config if model else None
    result = fetch_snowflake(database, schema, [table_name], config=config)
    if table_name not in result:
        return False, {}

    CACHE[fqn] = result[table_name]
    _save_cache()

    return False, CACHE[fqn]


def _save_cache():
    """Save the cache to disk."""
    global CACHE
    try:
        CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
        with CACHE_PATH.open("w", encoding="utf-8") as f:
            str_cache = {k: {ck: cv for ck, cv in v.items()} for k, v in CACHE.items()}
            json.dump(str_cache, f, indent=2)
    except OSError as e:
        warn("Schema Cache",
             f"Could not persist schema cache to {CACHE_PATH}. Using in-memory cache",
             [source(SourcePos.new()), f"[red]{e}"]
        )

#------------------------------------------------------
# Snowflake
#------------------------------------------------------

SUPPORTED_SNOWFLAKE_TYPES = [
    'CHAR', 'STRING', 'VARCHAR', 'BINARY', 'NUMBER', 'FLOAT', 'REAL',
    'BOOLEAN', 'DATE', 'FIXED', 'TEXT', 'TIME', 'TIMESTAMP_LTZ',
    'TIMESTAMP_NTZ', 'TIMESTAMP_TZ'
]

SFTypes = {
    "TEXT": "String",
    "FIXED": "Decimal",
    "DATE": "Date",
    "TIME": "DateTime",
    "TIMESTAMP": "DateTime",
    "TIMESTAMP_LTZ": "DateTime",
    "TIMESTAMP_TZ": "DateTime",
    "TIMESTAMP_NTZ": "DateTime",
    "FLOAT": "Float",
    "REAL": "Float",
    "BOOLEAN": "Bool",
}

def quoted(ident: str) -> str:
    return quote_identifier_if_needed(ident)

def get_provider(config=None):
    from ..shims.executor import get_provider as _get_provider
    return _get_provider(config)

def fetch_snowflake(database:str, schema:str, table_names:list[str], config=None) -> dict[str, dict[str, str]]:
    name_lookup = {name.upper(): name for name in table_names}
    name_lookup.update({name: name for name in table_names})

    tables = ", ".join(f"'{name.upper()}', '{name}'" for name in table_names)

    query = textwrap.dedent(f"""
        begin
            SHOW COLUMNS IN SCHEMA {quoted(database)}.{quoted(schema)};
            let r resultset := (
                select "table_name", "column_name", "data_type"
                from table(result_scan(-1)) as t
                where "table_name" in ({tables})
            );
            return table(r);
        end;
    """)

    # with debugging.span("fetch_schema", sql=query):
    # allow exception to propagate, as this will be caught and handled by the caller
    columns = get_provider(config).sql(query)
    assert isinstance(columns, list)

    unsupported_columns: dict[str, dict[str, str]] = defaultdict(dict)

    schemas: dict[str, dict[str, str]] = defaultdict(dict)
    for row in columns:
        table_name, column_name, data_type = row
        table_name = name_lookup.get(table_name, table_name)
        assert table_name is not None

        sf_type_info = json.loads(data_type)
        typ = sf_type_info.get("type")

        if typ not in SUPPORTED_SNOWFLAKE_TYPES:
            unsupported_columns[table_name][column_name] = typ
            continue

        if typ == "FIXED":
            concept_name = sf_numeric_to_type_str(column_name, sf_type_info)
        else:
            concept_name = SFTypes[typ]
        schemas[table_name][column_name] = concept_name

    for table_name, cols in unsupported_columns.items():
        col_str = ", ".join(f"{col}({typ})" for col, typ in cols.items())
        warn("Unsupported Column",
            f"The following columns in '{database}.{schema}.{table_name}' have unsupported types: {col_str}",
            [
                "These columns will not be accessible in your model.",
                "For the list of supported column types see: https://docs.relational.ai/api/cli/imports/stream/#supported-column-types"
            ])

    return schemas


def digits_to_bits(precision)-> int:
    """
    Transform from a number of base 10 digits to the number of bits necessary to represent
    that. If the precision is larger than 38, return None as that is not supported.

    For example, a number with 38 digits requires 128 bits.
    """
    if precision <= 2:
        return 8
    elif precision <= 4:
        return 16
    elif precision <= 9:
        return 32
    elif precision <= 18:
        return 64
    elif precision <= 38:
        return 128
    raise ValueError(f"Invalid numeric precision '{precision}'")

def sf_numeric_to_type_str(column_name: str, sf_type_info: dict) -> str:
    """
    Computes the appropriate type to use for this column. This code reflects exactly
    the logic currently used by RAI's CDC implementation to ensure we map to the exact
    same number types.
    """
    if "scale" not in sf_type_info or "precision" not in sf_type_info:
        raise ValueError(
            f"Invalid definition for column '{column_name}': "
            "'scale' or 'precision' missing"
        )

    precision = sf_type_info["precision"]
    scale = sf_type_info["scale"]

    if scale > precision or scale < 0 or scale > 37:
        raise ValueError(
            f"Invalid numeric scale '{scale}' for column '{column_name}'"
        )

    if scale == 0:
        # Integers (load_csv only supports these two (and not 8/16/32 bit ints)
        bits = digits_to_bits(precision)
        # return Int128 if bits == 128 else Int64
        return "Decimal(38,0)" if bits == 128 else "Decimal(18,0)"

    return f"Decimal({precision},{scale})"
