"""Snowflake CLI connection helpers for data modeling."""
from __future__ import annotations

import os
from pathlib import Path
from typing import Dict, Optional


def _load_toml(path: Path) -> Dict:
    try:
        import tomllib  # type: ignore
        return tomllib.loads(path.read_text(encoding="utf-8"))
    except Exception:
        try:
            import tomli  # type: ignore
            return tomli.loads(path.read_text(encoding="utf-8"))
        except Exception:
            return {}


def get_snowflake_connection(connection_name: Optional[str] = None) -> Dict[str, str]:
    """Load a Snowflake CLI connection from ~/.snowflake/config.toml."""
    config_path = Path(os.environ.get("USERPROFILE", str(Path.home()))) / ".snowflake" / "config.toml"
    if not config_path.exists():
        return {}
    config = _load_toml(config_path)
    connections = config.get("connections", {}) if isinstance(config, dict) else {}

    name = connection_name or config.get("default_connection_name") or "default"
    conn = connections.get(name, {})
    if not isinstance(conn, dict):
        return {}
    return {k: str(v) for k, v in conn.items()}


def ensure_snowflake_env_from_cli(
    connection_name: Optional[str] = None,
    allow_override: bool = False,
) -> Dict[str, str]:
    """Populate SNOWFLAKE_* env vars from Snowflake CLI connection."""
    conn = get_snowflake_connection(connection_name)
    if not conn:
        return {}

    mapping = {
        "SNOWFLAKE_ACCOUNT": conn.get("account"),
        "SNOWFLAKE_USER": conn.get("user"),
        "SNOWFLAKE_PASSWORD": conn.get("password"),
        "SNOWFLAKE_ROLE": conn.get("role"),
        "SNOWFLAKE_WAREHOUSE": conn.get("warehouse"),
        "SNOWFLAKE_DATABASE": conn.get("database"),
        "SNOWFLAKE_SCHEMA": conn.get("schema"),
        "SNOWFLAKE_AUTHENTICATOR": conn.get("authenticator"),
    }
    for key, value in mapping.items():
        if value is None:
            continue
        if allow_override or not os.getenv(key):
            os.environ[key] = value
    return {k: v for k, v in mapping.items() if v}
