"""
data_loader.py -- Data ingestion and validation for hedonic pricing analysis.

Provides a uniform interface for loading the cleaned sample CSV, full
Inside Airbnb extracts (listings, calendar, reviews, neighbourhoods),
and any pre-processed interim artefacts.  Every loader gracefully
handles missing files and returns ``None`` when data is unavailable.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd

from src.config import (
    CLEANED_SAMPLE_CSV,
    LISTINGS_CSV_GZ,
    CALENDAR_CSV_GZ,
    REVIEWS_CSV_GZ,
    NEIGHBOURHOODS_CSV,
    NEIGHBOURHOODS_GEOJSON,
    DATA_INTERIM,
    DATA_PROCESSED,
    QUEBEC_LAT_RANGE,
    QUEBEC_LON_RANGE,
    LOG_LEVEL,
    LOG_FORMAT,
    LOG_DATE_FORMAT,
)

# ---------------------------------------------------------------------------
# Logger
# ---------------------------------------------------------------------------

logging.basicConfig(format=LOG_FORMAT, datefmt=LOG_DATE_FORMAT, level=LOG_LEVEL)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Schema expectations (column name -> expected dtype family)
# ---------------------------------------------------------------------------

_EXPECTED_COLUMNS_SAMPLE: Dict[str, str] = {
    "name":              "object",
    "city":              "object",
    "property_type":     "object",
    "price_numeric":     "float",
    "price_category":    "object",
    "rating":            "float",
    "rating_category":   "object",
    "bedrooms":          "int",
    "bathrooms":         "float",
    "guests_count":      "int",
    "amenities_count":   "int",
    "num_images":        "int",
    "num_reviews":       "float",
    "quality_score":     "float",
    "is_superhost":      "bool",
    "is_guest_favorite": "bool",
    "pets_allowed":      "bool",
    "lat":               "float",
    "long":              "float",
}

# Columns that must be present for the pipeline to proceed
_REQUIRED_COLUMNS: List[str] = [
    "price_numeric", "bedrooms", "bathrooms", "guests_count",
    "lat", "long", "property_type",
]

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _file_available(path: Path) -> bool:
    """Return True if *path* exists and is non-empty."""
    return path.is_file() and path.stat().st_size > 0


def _infer_separator(path: Path) -> str:
    """Sniff the first line to decide between comma and tab."""
    with open(path, "r", encoding="utf-8", errors="replace") as fh:
        first_line = fh.readline()
    if "\t" in first_line:
        return "\t"
    return ","


def _safe_read_csv(
    path: Path,
    *,
    compression: Optional[str] = None,
    nrows: Optional[int] = None,
    usecols: Optional[List[str]] = None,
    dtype: Optional[Dict[str, Any]] = None,
    low_memory: bool = False,
) -> Optional[pd.DataFrame]:
    """Read a CSV/TSV with graceful error handling.

    Returns ``None`` when the file does not exist or cannot be parsed.
    """
    if not _file_available(path):
        logger.info("File not found or empty: %s", path)
        return None

    sep = "," if compression else _infer_separator(path)

    try:
        df = pd.read_csv(
            path,
            sep=sep,
            compression=compression or "infer",
            nrows=nrows,
            usecols=usecols,
            dtype=dtype,
            low_memory=low_memory,
            encoding="utf-8",
            on_bad_lines="warn",
        )
        logger.info(
            "Loaded %s rows x %s cols from %s",
            f"{len(df):,}",
            len(df.columns),
            path.name,
        )
        return df
    except Exception:
        logger.exception("Failed to read %s", path)
        return None


# ---------------------------------------------------------------------------
# Public loaders
# ---------------------------------------------------------------------------


def load_listings(
    *,
    path: Optional[Union[str, Path]] = None,
    nrows: Optional[int] = None,
    usecols: Optional[List[str]] = None,
) -> pd.DataFrame:
    """Load the primary listings dataset.

    Resolution order:
    1. Explicit *path* argument.
    2. Full Inside Airbnb ``listings.csv.gz`` in ``data/raw/``.
    3. ``cleaned_dataset_sample.csv`` at project root.

    Raises
    ------
    FileNotFoundError
        If no listing data can be found at any location.
    """
    if path is not None:
        p = Path(path)
        df = _safe_read_csv(p, nrows=nrows, usecols=usecols)
        if df is not None:
            return df
        raise FileNotFoundError(f"Specified listing file not found: {p}")

    # Try full extract first
    if _file_available(LISTINGS_CSV_GZ):
        df = _safe_read_csv(
            LISTINGS_CSV_GZ,
            compression="gzip",
            nrows=nrows,
            usecols=usecols,
            low_memory=True,
        )
        if df is not None:
            return df

    # Fall back to cleaned sample
    if _file_available(CLEANED_SAMPLE_CSV):
        df = _safe_read_csv(CLEANED_SAMPLE_CSV, nrows=nrows, usecols=usecols)
        if df is not None:
            return df

    raise FileNotFoundError(
        "No listing data found.  Place listings.csv.gz in data/raw/ "
        "or ensure cleaned_dataset_sample.csv exists at the project root."
    )


def load_calendar(
    *,
    path: Optional[Union[str, Path]] = None,
    nrows: Optional[int] = None,
) -> Optional[pd.DataFrame]:
    """Load Inside Airbnb calendar data (availability & price over time).

    Returns ``None`` if the file is not available.
    """
    p = Path(path) if path else CALENDAR_CSV_GZ
    compression = "gzip" if p.suffix == ".gz" else None
    df = _safe_read_csv(p, compression=compression, nrows=nrows, low_memory=True)
    if df is not None and "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"], errors="coerce")
    return df


def load_reviews(
    *,
    path: Optional[Union[str, Path]] = None,
    nrows: Optional[int] = None,
) -> Optional[pd.DataFrame]:
    """Load Inside Airbnb reviews data.

    Returns ``None`` if the file is not available.
    """
    p = Path(path) if path else REVIEWS_CSV_GZ
    compression = "gzip" if p.suffix == ".gz" else None
    df = _safe_read_csv(p, compression=compression, nrows=nrows, low_memory=True)
    if df is not None and "date" in df.columns:
        df["date"] = pd.to_datetime(df["date"], errors="coerce")
    return df


def load_neighbourhoods(
    *,
    path: Optional[Union[str, Path]] = None,
    prefer_geojson: bool = True,
) -> Optional[pd.DataFrame]:
    """Load neighbourhood boundaries or names.

    If *prefer_geojson* is True (default), attempts GeoJSON first, then
    CSV.  Returns a plain ``DataFrame`` (not GeoDataFrame) unless
    geopandas is installed, in which case it returns a ``GeoDataFrame``.

    Returns ``None`` if no file is available.
    """
    csv_path = Path(path) if path else NEIGHBOURHOODS_CSV
    geojson_path = NEIGHBOURHOODS_GEOJSON

    if prefer_geojson and _file_available(geojson_path):
        try:
            import geopandas as gpd  # type: ignore[import-untyped]

            gdf = gpd.read_file(geojson_path)
            logger.info(
                "Loaded %s neighbourhoods from GeoJSON", f"{len(gdf):,}"
            )
            return gdf
        except ImportError:
            logger.warning(
                "geopandas not installed; falling back to CSV neighbourhoods"
            )
        except Exception:
            logger.exception("Failed to read GeoJSON: %s", geojson_path)

    if _file_available(csv_path):
        return _safe_read_csv(csv_path)

    logger.info("No neighbourhood data available")
    return None


def load_interim(name: str) -> Optional[pd.DataFrame]:
    """Load a parquet or CSV from ``data/interim/`` by stem name."""
    parquet = DATA_INTERIM / f"{name}.parquet"
    csv = DATA_INTERIM / f"{name}.csv"

    if _file_available(parquet):
        try:
            df = pd.read_parquet(parquet)
            logger.info("Loaded interim parquet: %s (%s rows)", name, f"{len(df):,}")
            return df
        except Exception:
            logger.exception("Failed to read interim parquet: %s", parquet)

    if _file_available(csv):
        return _safe_read_csv(csv)

    logger.info("Interim artefact '%s' not found", name)
    return None


def load_processed(name: str) -> Optional[pd.DataFrame]:
    """Load a parquet or CSV from ``data/processed/`` by stem name."""
    parquet = DATA_PROCESSED / f"{name}.parquet"
    csv = DATA_PROCESSED / f"{name}.csv"

    if _file_available(parquet):
        try:
            df = pd.read_parquet(parquet)
            logger.info("Loaded processed parquet: %s (%s rows)", name, f"{len(df):,}")
            return df
        except Exception:
            logger.exception("Failed to read processed parquet: %s", parquet)

    if _file_available(csv):
        return _safe_read_csv(csv)

    logger.info("Processed artefact '%s' not found", name)
    return None


def load_all_data(
    *,
    nrows: Optional[int] = None,
) -> Dict[str, Optional[pd.DataFrame]]:
    """Master loader -- returns a dict of every available dataset.

    Keys
    ----
    ``"listings"``
        Always present (raises if not).
    ``"calendar"``, ``"reviews"``, ``"neighbourhoods"``
        ``None`` when the corresponding raw file is missing.

    Parameters
    ----------
    nrows : int, optional
        Limit rows for quick exploratory loads.
    """
    data: Dict[str, Optional[pd.DataFrame]] = {}

    data["listings"] = load_listings(nrows=nrows)
    data["calendar"] = load_calendar(nrows=nrows)
    data["reviews"] = load_reviews(nrows=nrows)
    data["neighbourhoods"] = load_neighbourhoods()

    available = [k for k, v in data.items() if v is not None]
    logger.info("Loaded datasets: %s", ", ".join(available))
    return data


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


def validate_data(
    df: pd.DataFrame,
    *,
    strict: bool = False,
) -> Dict[str, Any]:
    """Validate a listings DataFrame against expected schema.

    Parameters
    ----------
    df : pd.DataFrame
        The listings data to validate.
    strict : bool
        If True, raise ``ValueError`` on critical failures instead of
        just logging warnings.

    Returns
    -------
    dict
        ``{"valid": bool, "warnings": [...], "errors": [...]}``.
    """
    warnings_list: List[str] = []
    errors_list: List[str] = []

    # -- Required columns --------------------------------------------------
    missing_required = [c for c in _REQUIRED_COLUMNS if c not in df.columns]
    if missing_required:
        msg = f"Missing required columns: {missing_required}"
        errors_list.append(msg)
        logger.error(msg)

    # -- Expected columns (informational) -----------------------------------
    missing_expected = [
        c for c in _EXPECTED_COLUMNS_SAMPLE if c not in df.columns
    ]
    if missing_expected:
        msg = f"Missing expected (non-critical) columns: {missing_expected}"
        warnings_list.append(msg)
        logger.warning(msg)

    # -- Row count ----------------------------------------------------------
    if len(df) == 0:
        errors_list.append("DataFrame is empty")
    elif len(df) < 100:
        warnings_list.append(
            f"Very small dataset ({len(df)} rows); results may be unreliable"
        )

    # -- Price range --------------------------------------------------------
    if "price_numeric" in df.columns:
        prices = df["price_numeric"].dropna()
        if len(prices) == 0:
            errors_list.append("All price values are missing")
        else:
            if prices.min() <= 0:
                warnings_list.append(
                    f"Prices include non-positive values (min={prices.min():.2f})"
                )
            if prices.max() > 100_000:
                warnings_list.append(
                    f"Extreme price outliers present (max={prices.max():,.0f})"
                )

    # -- Coordinate ranges --------------------------------------------------
    if "lat" in df.columns and "long" in df.columns:
        lat_min, lat_max = QUEBEC_LAT_RANGE
        lon_min, lon_max = QUEBEC_LON_RANGE

        out_of_bounds = df[
            (df["lat"] < lat_min)
            | (df["lat"] > lat_max)
            | (df["long"] < lon_min)
            | (df["long"] > lon_max)
        ]
        if len(out_of_bounds) > 0:
            pct = 100.0 * len(out_of_bounds) / len(df)
            warnings_list.append(
                f"{len(out_of_bounds)} listings ({pct:.1f}%) have coordinates "
                f"outside Quebec bounding box"
            )

    # -- Duplicate coordinates (exact) --------------------------------------
    if "lat" in df.columns and "long" in df.columns:
        dup_coords = df.duplicated(subset=["lat", "long"], keep=False).sum()
        if dup_coords > 0:
            pct = 100.0 * dup_coords / len(df)
            warnings_list.append(
                f"{dup_coords} listings ({pct:.1f}%) share exact coordinates "
                f"(possible duplicates or jittered points)"
            )

    # -- Boolean columns ----------------------------------------------------
    for col in ("is_superhost", "is_guest_favorite", "pets_allowed"):
        if col in df.columns and df[col].dtype not in ("bool", "boolean"):
            warnings_list.append(
                f"Column '{col}' has dtype {df[col].dtype}, expected bool"
            )

    # -- Rating range -------------------------------------------------------
    if "rating" in df.columns:
        valid_ratings = df["rating"].dropna()
        if len(valid_ratings) > 0:
            if valid_ratings.min() < 0 or valid_ratings.max() > 5:
                warnings_list.append(
                    f"Ratings outside 0-5 range: "
                    f"min={valid_ratings.min():.2f}, max={valid_ratings.max():.2f}"
                )

    # -- Summary ------------------------------------------------------------
    is_valid = len(errors_list) == 0
    result = {
        "valid": is_valid,
        "warnings": warnings_list,
        "errors": errors_list,
        "n_rows": len(df),
        "n_cols": len(df.columns),
        "columns": list(df.columns),
    }

    if is_valid:
        logger.info(
            "Validation passed (%s warnings)", len(warnings_list)
        )
    else:
        logger.error("Validation FAILED with %s error(s)", len(errors_list))

    if strict and not is_valid:
        raise ValueError(
            "Data validation failed:\n" + "\n".join(errors_list)
        )

    return result


# ---------------------------------------------------------------------------
# Summary statistics
# ---------------------------------------------------------------------------


def get_data_summary(df: pd.DataFrame) -> Dict[str, Any]:
    """Return a rich summary-statistics dictionary for a listings DataFrame.

    Useful for dashboards, logging, and the cleaning report.
    """
    summary: Dict[str, Any] = {
        "n_rows": len(df),
        "n_cols": len(df.columns),
        "columns": list(df.columns),
        "dtypes": {col: str(dt) for col, dt in df.dtypes.items()},
        "missing_counts": df.isnull().sum().to_dict(),
        "missing_pct": (100.0 * df.isnull().mean()).round(2).to_dict(),
    }

    # Numeric column stats
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    if numeric_cols:
        desc = df[numeric_cols].describe(percentiles=[0.01, 0.05, 0.25, 0.5, 0.75, 0.95, 0.99])
        summary["numeric_stats"] = desc.to_dict()

    # Price specifics
    if "price_numeric" in df.columns:
        prices = df["price_numeric"].dropna()
        summary["price"] = {
            "count": int(len(prices)),
            "missing": int(df["price_numeric"].isnull().sum()),
            "mean": float(prices.mean()) if len(prices) > 0 else None,
            "median": float(prices.median()) if len(prices) > 0 else None,
            "std": float(prices.std()) if len(prices) > 0 else None,
            "min": float(prices.min()) if len(prices) > 0 else None,
            "max": float(prices.max()) if len(prices) > 0 else None,
            "p1": float(prices.quantile(0.01)) if len(prices) > 0 else None,
            "p5": float(prices.quantile(0.05)) if len(prices) > 0 else None,
            "p95": float(prices.quantile(0.95)) if len(prices) > 0 else None,
            "p99": float(prices.quantile(0.99)) if len(prices) > 0 else None,
            "skewness": float(prices.skew()) if len(prices) > 0 else None,
            "kurtosis": float(prices.kurtosis()) if len(prices) > 0 else None,
        }

    # Categorical breakdowns
    for cat_col in ("city", "property_type", "price_category", "rating_category", "room_type"):
        if cat_col in df.columns:
            vc = df[cat_col].value_counts(dropna=False)
            summary[f"{cat_col}_counts"] = vc.to_dict()

    # Boolean flags
    for bool_col in ("is_superhost", "is_guest_favorite", "pets_allowed"):
        if bool_col in df.columns:
            true_count = int(df[bool_col].sum())
            summary[f"{bool_col}_rate"] = round(
                100.0 * true_count / len(df), 2
            ) if len(df) > 0 else 0.0

    # Geo extent
    if "lat" in df.columns and "long" in df.columns:
        summary["geo_extent"] = {
            "lat_min": float(df["lat"].min()),
            "lat_max": float(df["lat"].max()),
            "lon_min": float(df["long"].min()),
            "lon_max": float(df["long"].max()),
        }

    # Top cities by count
    if "city" in df.columns:
        top_cities = df["city"].value_counts().head(15)
        summary["top_cities"] = top_cities.to_dict()

    logger.info(
        "Data summary: %s rows, %s cols, %s numeric features",
        f"{summary['n_rows']:,}",
        summary["n_cols"],
        len(numeric_cols),
    )
    return summary