"""Sample loader for data modeling inputs."""
from __future__ import annotations

import csv
import logging
from pathlib import Path
from typing import Dict, List, Optional

logger = logging.getLogger(__name__)


def _read_csv(path: Path, limit: int) -> Dict[str, List[dict]]:
    with path.open(newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    return {"columns": list(rows[0].keys()) if rows else [], "rows": rows[:limit]}


def load_csv_samples(paths: List[str], limit: int = 200) -> Dict[str, Dict[str, List[dict]]]:
    samples: Dict[str, Dict[str, List[dict]]] = {}
    for p in paths:
        path = Path(p)
        if not path.exists():
            continue
        table_name = path.stem.upper()
        samples[table_name] = _read_csv(path, limit)
    return samples


def load_csv_folder(folder: str, limit: int = 200) -> Dict[str, Dict[str, List[dict]]]:
    base = Path(folder)
    if not base.exists():
        return {}
    paths = [str(p) for p in base.glob("*.csv")]
    return load_csv_samples(paths, limit=limit)


def load_sqlalchemy_samples(
    connection_string: str,
    database: Optional[str],
    schema: Optional[str],
    tables: List[str],
    limit: int = 200,
) -> Dict[str, Dict[str, List[dict]]]:
    try:
        from sqlalchemy import create_engine, text
    except Exception as exc:
        raise RuntimeError("sqlalchemy not installed") from exc

    engine = create_engine(connection_string)
    samples: Dict[str, Dict[str, List[dict]]] = {}

    for table in tables:
        fqn = ".".join([p for p in [database, schema, table] if p])
        sql = text(f"SELECT * FROM {fqn} LIMIT {limit}")
        with engine.connect() as conn:
            result = conn.execute(sql)
            rows = [dict(r._mapping) for r in result]
        samples[table.upper()] = {
            "columns": list(rows[0].keys()) if rows else [],
            "rows": rows,
        }
    return samples


def load_snowflake_samples(
    database: str,
    schema: str,
    tables: List[str],
    limit: int = 200,
    sf_connection=None,
) -> Dict[str, Dict[str, List[dict]]]:
    """Load sample data from Snowflake tables.

    Args:
        sf_connection: Optional pre-existing Snowflake connection to reuse.
                       When provided, the connection is NOT closed after use.
    """
    own_conn = sf_connection is None

    if own_conn:
        # Use connection pool to avoid repeated SSO prompts
        from .snowflake_pool import sf_pool
        import os
        from .snowflake_utils import ensure_snowflake_env_from_cli

        if not os.getenv("SNOWFLAKE_ACCOUNT"):
            ensure_snowflake_env_from_cli()
        if not os.getenv("SNOWFLAKE_ACCOUNT"):
            raise RuntimeError("SNOWFLAKE_ACCOUNT not set")
        conn = sf_pool.get_from_env(database=database, schema=schema)
    else:
        conn = sf_connection

    samples: Dict[str, Dict[str, List[dict]]] = {}
    cur = conn.cursor()
    total = len(tables)
    for i, table in enumerate(tables, 1):
        logger.info("Sampling table %d/%d: %s", i, total, table)
        cur.execute(f"SELECT * FROM {database}.{schema}.{table} LIMIT {limit}")
        cols = [c[0] for c in cur.description] if cur.description else []
        rows = [dict(zip(cols, r)) for r in cur.fetchall()]
        samples[table.upper()] = {"columns": cols, "rows": rows}
        logger.info("  %s: %d rows, %d columns", table, len(rows), len(cols))
    cur.close()
    # Don't close the connection — pool manages lifecycle
    return samples
