from __future__ import annotations

import hashlib
from typing import Iterable

import pandas as pd


def _normalise_for_hashing(df: pd.DataFrame) -> pd.DataFrame:
    """
    Convert values into a stable string representation before hashing.

    This avoids dtype-related surprises and keeps hashing deterministic.
    """
    out = df.copy()

    for col in out.columns:
        series = out[col]

        try:
            if pd.api.types.is_datetime64_any_dtype(series):
                out[col] = series.dt.strftime("%Y-%m-%d %H:%M:%S.%f").fillna("__NaT__")
            else:
                out[col] = series.astype("object").where(series.notna(), "__NaN__").astype(str)
        except Exception:
            out[col] = series.astype(str)

    return out


def stable_row_hashes(
    df: pd.DataFrame,
    exclude_cols: Iterable[str] | None = None,
    hash_name: str = "sha256",
) -> pd.Series:
    """
    Return a stable hash per row, optionally excluding some columns.

    Parameters
    ----------
    df:
        Input dataframe.
    exclude_cols:
        Columns to exclude before hashing.
    hash_name:
        Hash algorithm name accepted by hashlib.
    """
    exclude_set = set(exclude_cols or [])
    cols = [col for col in df.columns if col not in exclude_set]

    if not cols:
        return pd.Series([""] * len(df), index=df.index, dtype="object")

    normalised = _normalise_for_hashing(df[cols])
    joined = normalised.astype(str).agg("||".join, axis=1)

    def _hash_text(text: str) -> str:
        h = hashlib.new(hash_name)
        h.update(text.encode("utf-8", errors="replace"))
        return h.hexdigest()

    return joined.map(_hash_text)


def overlapping_row_hashes(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
    exclude_cols: Iterable[str] | None = None,
    hash_name: str = "sha256",
) -> set[str]:
    """
    Return the set of row hashes appearing in both train and test.
    """
    train_hashes = set(stable_row_hashes(train_df, exclude_cols=exclude_cols, hash_name=hash_name))
    test_hashes = set(stable_row_hashes(test_df, exclude_cols=exclude_cols, hash_name=hash_name))
    return train_hashes & test_hashes


def duplicate_row_mask(
    df: pd.DataFrame,
    exclude_cols: Iterable[str] | None = None,
    keep: str | bool = False,
) -> pd.Series:
    """
    Duplicate mask based on stable row hashing.

    Parameters match pandas duplicate semantics loosely:
    - keep=False marks all duplicates
    - keep='first' keeps first
    - keep='last' keeps last
    """
    hashes = stable_row_hashes(df, exclude_cols=exclude_cols)

    if keep is False:
        return hashes.duplicated(keep=False)

    if keep in {"first", "last"}:
        return hashes.duplicated(keep=keep)

    raise ValueError("keep must be False, 'first', or 'last'.")