"""CTF dataset downloader via WRDS PostgreSQL.

Downloads the three CTF Parquet files from Wharton Research Data Services (WRDS)
and saves them locally for use with ``run_local()`` and ``pipeline()``.

WRDS connection details::

    host : wrds-pgdata.wharton.upenn.edu
    port : 9737
    database : wrds
    schema : ctff
    SSL : required

The three tables downloaded:

    ctff_chars       — firm characteristics at each eom (~900+ columns)
    ctff_features    — feature name list (single column, small)
    ctff_daily_ret   — daily excess returns (id, date, ret_exc)

Credentials are loaded from the project ``.env`` file (or environment variables).
Set ``WRDS_USERNAME`` in ``.env`` — the ``wrds`` package handles password
authentication interactively on first login and caches it in ``~/.pgpass``.

Usage::

    from eapctf.ctf.data import download_ctf_data

    # wrds_username read from .env / WRDS_USERNAME env var automatically
    paths = download_ctf_data(output_dir="data/ctf/")
    print(paths["ctff_chars"])   # Path to saved Parquet file
"""

from __future__ import annotations

import os
from pathlib import Path

import pandas as pd


def _load_env() -> None:
    """Load .env from the project root (silently skip if not found or dotenv missing)."""
    try:
        from dotenv import load_dotenv
    except ImportError:
        return
    # Walk up from this file to find .env at the project root
    here = Path(__file__).resolve()
    for parent in here.parents:
        env_file = parent / ".env"
        if env_file.exists():
            load_dotenv(env_file, override=False)
            return


# Load .env once at import time
_load_env()


# ---------------------------------------------------------------------------
# WRDS connection constants
# ---------------------------------------------------------------------------

_WRDS_HOST = "wrds-pgdata.wharton.upenn.edu"
_WRDS_PORT = 9737
_WRDS_DB = "wrds"
_CTF_SCHEMA = "contrib_global_factor"

_TABLE_NAMES = ("ctff_chars", "ctff_features", "ctff_daily_ret")

# Chunk size for large tables (rows per fetch).  ctff_chars has ~900 columns
# so 50k rows × 900 cols already occupies ~200 MB in memory.
_DEFAULT_CHUNK_ROWS = 50_000


# ---------------------------------------------------------------------------
# Connection helpers
# ---------------------------------------------------------------------------


def _open_wrds(wrds_username: str) -> object:
    """Return an open ``wrds.Connection`` (preferred) or psycopg2 connection.

    The ``wrds`` package (v3.x) uses SQLAlchemy internally and exposes
    ``.engine`` for use with ``pd.read_sql()``.  When the package is not
    installed, falls back to a raw psycopg2 connection.

    Returns the connection object directly — callers must call ``.close()``
    when finished.
    """
    try:
        import wrds

        return wrds.Connection(wrds_username=wrds_username)
    except ImportError:
        pass

    # Fallback: direct psycopg2
    try:
        import psycopg2
    except ImportError:
        raise ImportError(
            "Either the 'wrds' package or 'psycopg2' is required.\n"
            "  pip install wrds          (recommended)\n"
            "  pip install psycopg2-binary"
        )

    import getpass

    password = getpass.getpass(f"WRDS password for '{wrds_username}': ")
    return psycopg2.connect(
        host=_WRDS_HOST,
        port=_WRDS_PORT,
        dbname=_WRDS_DB,
        user=wrds_username,
        password=password,
        sslmode="require",
        connect_timeout=30,
    )


# ---------------------------------------------------------------------------
# Table fetch helpers
# ---------------------------------------------------------------------------


def _fetch_table_chunked(
    db: object,
    table: str,
    schema: str,
    chunk_rows: int,
) -> pd.DataFrame:
    """Fetch an entire table in chunks; return a single concatenated DataFrame.

    Supports two backends:
    * ``wrds.Connection`` (v3.x, SQLAlchemy) — uses ``pd.read_sql()`` with
      the connection's ``.engine`` attribute and ``chunksize``.
    * Raw psycopg2 connection — uses a named server-side cursor.

    Parameters
    ----------
    db
        Open connection returned by ``_open_wrds()``.
    table : str
        Table name (without schema prefix).
    schema : str
        PostgreSQL schema name.
    chunk_rows : int
        Rows per fetch iteration.

    Returns
    -------
    pd.DataFrame
    """
    qualified = f'"{schema}"."{table}"'
    sql = f"SELECT * FROM {qualified}"
    print(f"  Fetching {qualified} ...", flush=True)

    chunks: list[pd.DataFrame] = []
    fetched = 0

    # wrds.Connection exposes .engine (SQLAlchemy Engine)
    if hasattr(db, "engine"):
        for chunk in pd.read_sql(sql, db.engine, chunksize=chunk_rows):
            chunks.append(chunk)
            fetched += len(chunk)
            print(f"\r  Fetching {qualified} ... {fetched:,} rows", end="", flush=True)
    else:
        # psycopg2 fallback: named server-side cursor
        import psycopg2.extras

        cursor_name = f"_eap_ctf_{table}"
        with db.cursor(cursor_name, cursor_factory=psycopg2.extras.DictCursor) as cur:  # type: ignore[attr-defined]
            cur.execute(sql)
            col_names: list[str] | None = None
            while True:
                rows = cur.fetchmany(chunk_rows)
                if not rows:
                    break
                if col_names is None:
                    col_names = [desc[0] for desc in cur.description]
                chunks.append(pd.DataFrame(list(rows), columns=col_names))
                fetched += len(rows)
                print(f"\r  Fetching {qualified} ... {fetched:,} rows", end="", flush=True)

    print(f"\r  Fetched  {qualified} — {fetched:,} rows             ")

    if not chunks:
        return pd.DataFrame()

    return pd.concat(chunks, ignore_index=True)


def _infer_dtypes(df: pd.DataFrame, table: str) -> pd.DataFrame:
    """Apply known type casts to CTF tables to reduce memory use.

    For ctff_daily_ret and ctff_chars: cast known date columns to datetime,
    cast id columns to int32, cast float columns to float32 where safe.
    """
    df = df.copy()

    date_cols = {"date", "eom", "eom_ret"}
    id_cols = {"id"}
    bool_cols = {"ctff_test"}

    for col in df.columns:
        if col in date_cols:
            df[col] = pd.to_datetime(df[col], errors="coerce")
        elif col in id_cols:
            try:
                df[col] = df[col].astype("int32")
            except (ValueError, TypeError):
                pass
        elif col in bool_cols:
            try:
                df[col] = df[col].astype(bool)
            except (ValueError, TypeError):
                pass

    return df


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def download_ctf_data(
    wrds_username: str | None = None,
    output_dir: str | Path = "data/ctf",
    wrds_password: str | None = None,
    chunk_rows: int = _DEFAULT_CHUNK_ROWS,
    overwrite: bool = False,
    tables: tuple[str, ...] = _TABLE_NAMES,
) -> dict[str, Path]:
    """Download CTF datasets from WRDS and save as Parquet files.

    Parameters
    ----------
    wrds_username : str | None
        WRDS username.  If None, reads ``WRDS_USERNAME`` from the environment
        (loaded from ``.env`` automatically).
    output_dir : str | Path
        Directory to write Parquet files.  Created if it does not exist.
    wrds_password : str | None
        WRDS password.  Normally left as None — the ``wrds`` package prompts
        on first login and caches credentials in ``~/.pgpass``.
    chunk_rows : int
        Rows per server-side cursor fetch.  Reduce if running low on memory.
    overwrite : bool
        If False (default) and a file already exists, skip that table.
        Set True to re-download all tables.
    tables : tuple[str, ...]
        Table names to download.  Defaults to all three CTF tables.

    Returns
    -------
    dict[str, Path]
        Mapping from table name to the saved Parquet file path.

    Raises
    ------
    ValueError
        If WRDS username cannot be determined from argument or environment.
    ImportError
        If neither ``wrds`` nor ``psycopg2`` is installed.
    """
    # Resolve username: argument → env var → error
    if wrds_username is None:
        wrds_username = os.environ.get("WRDS_USERNAME", "").strip() or None
    if not wrds_username:
        raise ValueError(
            "WRDS username not set.  Either pass wrds_username= or set "
            "WRDS_USERNAME in your .env file."
        )

    output_dir = Path(output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    paths: dict[str, Path] = {}
    to_download: list[str] = []

    for table in tables:
        dest = output_dir / f"{table}.parquet"
        paths[table] = dest
        if dest.exists() and not overwrite:
            print(f"  Skipping {table} — already exists at {dest}")
        else:
            to_download.append(table)

    if not to_download:
        print("All CTF tables already downloaded.  Use overwrite=True to re-fetch.")
        return paths

    print(f"Connecting to WRDS as '{wrds_username}' ...")
    db = _open_wrds(wrds_username)

    try:
        for table in to_download:
            dest = output_dir / f"{table}.parquet"
            df = _fetch_table_chunked(db, table, _CTF_SCHEMA, chunk_rows)
            df = _infer_dtypes(df, table)
            df.to_parquet(dest, index=False, engine="pyarrow")
            size_mb = dest.stat().st_size / (1024 ** 2)
            print(f"  Saved  {dest}  ({size_mb:.1f} MB,  {len(df):,} rows)")
            paths[table] = dest
    finally:
        try:
            db.close()  # type: ignore[attr-defined]
        except Exception:
            pass

    print("Download complete.")
    return paths
