from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import hashlib
import math
import numpy as np
import pandas as pd


class InputValidationError(ValueError):
    pass


@dataclass
class PreparedData:
    joined: pd.DataFrame
    radial: pd.DataFrame
    signatures: pd.DataFrame
    report: dict[str, Any]


def sha256_file(path: str | Path) -> str:
    h = hashlib.sha256()
    with Path(path).open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _as_bool(series: pd.Series) -> pd.Series:
    if pd.api.types.is_bool_dtype(series):
        return series.fillna(False)
    return series.astype(str).str.strip().str.lower().isin({"true", "1", "yes", "y"})


def load_and_validate(radial_path: str | Path, signature_path: str | Path, config: dict[str, Any]) -> PreparedData:
    radial_path = Path(radial_path)
    signature_path = Path(signature_path)
    radial = pd.read_csv(radial_path)
    signatures = pd.read_csv(signature_path)

    group = config["columns"]["group"]
    target = config["columns"]["target"]
    local_gbar = config["columns"]["local_gbar"]
    radius = config["columns"]["radius"]
    filter_col = config["sample"]["radial_filter_column"]

    radial_required = {group, target, local_gbar, radius, filter_col, "valid_point"}
    signature_required = {group}
    for feature_group_name, values in config["feature_groups"].items():
        if feature_group_name == "local_spline":
            continue
        signature_required.update(values)
    # Derived in this module from the native exported threshold.
    signature_required.discard("perc_log10_kappa_connect")
    signature_required.add("perc_kappa_connect_exported")

    missing_radial = sorted(radial_required - set(radial.columns))
    missing_signature = sorted(signature_required - set(signatures.columns))
    if missing_radial:
        raise InputValidationError(f"Radial input is missing columns: {missing_radial}")
    if missing_signature:
        raise InputValidationError(f"Signature input is missing columns: {missing_signature}")

    radial = radial[_as_bool(radial[filter_col]) & _as_bool(radial["valid_point"])].copy()
    if radial.empty:
        raise InputValidationError("No registered primary radial rows remain after filtering.")

    # One row per signature galaxy is mandatory.
    duplicate_signatures = signatures[group][signatures[group].duplicated()].unique().tolist()
    if duplicate_signatures:
        raise InputValidationError(f"Duplicate signature rows: {duplicate_signatures[:10]}")

    expected = int(config["sample"]["expected_galaxies"])
    radial_groups = set(radial[group].astype(str))
    signature_groups = set(signatures[group].astype(str))
    if len(radial_groups) != expected:
        raise InputValidationError(f"Expected {expected} radial galaxies; found {len(radial_groups)}.")
    if len(signature_groups) != expected:
        raise InputValidationError(f"Expected {expected} signature galaxies; found {len(signature_groups)}.")
    if radial_groups != signature_groups:
        missing_from_signatures = sorted(radial_groups - signature_groups)
        extra_signatures = sorted(signature_groups - radial_groups)
        raise InputValidationError(
            f"Galaxy mismatch. Missing signatures={missing_from_signatures}; extra signatures={extra_signatures}"
        )

    counts = radial.groupby(group).size()
    minimum = int(config["sample"]["minimum_radial_points_per_galaxy"])
    too_short = counts[counts < minimum]
    if len(too_short):
        raise InputValidationError(f"Galaxies below {minimum} radial points: {too_short.to_dict()}")

    for col in [target, local_gbar, radius]:
        radial[col] = pd.to_numeric(radial[col], errors="coerce")
        if not np.isfinite(radial[col].to_numpy(dtype=float)).all():
            raise InputValidationError(f"Non-finite values in radial column {col}.")

    if ((radial[radius] < -1e-12) | (radial[radius] > 1 + 1e-12)).any():
        raise InputValidationError("radius_normalized must lie in [0, 1].")

    signatures = signatures.copy()
    threshold = pd.to_numeric(signatures["perc_kappa_connect_exported"], errors="coerce")
    signatures["perc_log10_kappa_connect"] = np.log10(1.0 + threshold)

    # Prevent accidental target leakage from signature tables.
    forbidden_signature_columns = {
        "g_obs", "g_obs_ms2", "log10_g_obs", "log10_discrepancy", "target", "y"
    }
    leaked = sorted(forbidden_signature_columns & set(signatures.columns))
    if leaked:
        raise InputValidationError(f"Potential target leakage columns in signatures: {leaked}")

    joined = radial.merge(signatures, on=group, how="left", validate="many_to_one", suffixes=("", "_signature"))
    if joined.isna().all(axis=1).any():
        raise InputValidationError("Join produced empty rows.")
    if len(joined) != len(radial):
        raise InputValidationError("Join changed radial row count.")

    joined.insert(0, "row_id", np.arange(len(joined), dtype=int))

    report = {
        "status": "PASS",
        "radial_rows": int(len(radial)),
        "joined_rows": int(len(joined)),
        "galaxies": int(joined[group].nunique()),
        "minimum_rows_per_galaxy": int(counts.min()),
        "maximum_rows_per_galaxy": int(counts.max()),
        "target_min": float(joined[target].min()),
        "target_max": float(joined[target].max()),
        "radial_sha256": sha256_file(radial_path),
        "signatures_sha256": sha256_file(signature_path),
    }
    return PreparedData(joined=joined, radial=radial, signatures=signatures, report=report)
