"""
GDPR compliance utilities for LakeLogic.

Provides right-to-forget (erasure) and PII masking/nullification capabilities
driven by the contract's ``pii: true`` field annotations.

Usage (Python API):
    from lakelogic.core.gdpr import forget_subjects, mask_pii_columns

    # Erase specific subjects from a dataframe
    cleaned_df = forget_subjects(df, contract, subject_column="customer_id",
                                  subject_ids=["cust_123", "cust_456"])

    # Mask all PII columns (nullify, hash, or redact)
    masked_df = mask_pii_columns(df, contract, strategy="nullify")

Usage via DataProcessor:
    proc = DataProcessor("contract.yaml")
    cleaned = proc.forget(subject_column="customer_id", subject_ids=["cust_123"])
    masked = proc.mask_pii(strategy="hash")
"""

from __future__ import annotations

import hashlib
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from loguru import logger

from lakelogic.core.models import DataContract, FieldDefinition


def _get_pii_fields(contract: DataContract) -> List[FieldDefinition]:
    """Extract field definitions marked as PII from the contract."""
    if not contract.model or not contract.model.fields:
        return []
    return [f for f in contract.model.fields if f.pii]


def _get_pii_column_names(contract: DataContract) -> List[str]:
    """Get the names of PII-marked columns."""
    return [f.name for f in _get_pii_fields(contract)]


def _hash_value(value: Any, salt: str = "") -> Optional[str]:
    """
    One-way SHA-256 hash of a value. Returns None for null values.

    Args:
        value: Value to hash.
        salt: Optional salt for the hash.

    Returns:
        Hex digest string or None.
    """
    if value is None or (isinstance(value, float) and value != value):  # NaN check
        return None
    raw = f"{salt}{value}".encode("utf-8")
    return hashlib.sha256(raw).hexdigest()


def _redact_value(value: Any, replacement: str = "***REDACTED***") -> Optional[str]:
    """Replace a value with a redaction marker. Returns None for null values."""
    if value is None or (isinstance(value, float) and value != value):
        return None
    return replacement


# ── Polars implementations ───────────────────────────────────────────────────


def _forget_polars(
    df,
    pii_columns: List[str],
    subject_column: str,
    subject_ids: List[str],
    erasure_strategy: str,
    hash_salt: str,
):
    """Erase/mask PII for specific subjects in a Polars DataFrame."""
    import polars as pl

    if isinstance(df, pl.LazyFrame):
        df = df.collect()

    present_pii = [c for c in pii_columns if c in df.columns]
    if not present_pii:
        logger.info("No PII columns found in dataframe; nothing to erase.")
        return df

    if subject_column not in df.columns:
        raise ValueError(f"Subject column '{subject_column}' not found in dataframe.")

    # Cast subject_ids to match column type
    subject_ids_set = set(str(s) for s in subject_ids)
    mask = df[subject_column].cast(pl.Utf8).is_in(list(subject_ids_set))

    for col in present_pii:
        if col == subject_column and erasure_strategy == "nullify":
            # Don't nullify the subject column itself in nullify mode
            # (would lose the key needed for audit)
            continue

        if erasure_strategy == "nullify":
            df = df.with_columns(pl.when(mask).then(pl.lit(None)).otherwise(pl.col(col)).alias(col))
        elif erasure_strategy == "hash":
            # Hash values for matching subjects
            df = df.with_columns(
                pl.when(mask)
                .then(
                    pl.col(col)
                    .cast(pl.Utf8)
                    .map_elements(
                        lambda v, _salt=hash_salt: _hash_value(v, _salt),
                        return_dtype=pl.Utf8,
                    )
                )
                .otherwise(pl.col(col))
                .alias(col)
            )
        elif erasure_strategy == "redact":
            df = df.with_columns(pl.when(mask).then(pl.lit("***REDACTED***")).otherwise(pl.col(col)).alias(col))

    return df


def _forget_pandas(
    df,
    pii_columns: List[str],
    subject_column: str,
    subject_ids: List[str],
    erasure_strategy: str,
    hash_salt: str,
):
    """Erase/mask PII for specific subjects in a Pandas DataFrame."""

    present_pii = [c for c in pii_columns if c in df.columns]
    if not present_pii:
        logger.info("No PII columns found in dataframe; nothing to erase.")
        return df

    if subject_column not in df.columns:
        raise ValueError(f"Subject column '{subject_column}' not found in dataframe.")

    df = df.copy()
    subject_ids_set = set(str(s) for s in subject_ids)
    mask = df[subject_column].astype(str).isin(subject_ids_set)
    affected_count = mask.sum()

    for col in present_pii:
        if col == subject_column and erasure_strategy == "nullify":
            continue

        if erasure_strategy == "nullify":
            df.loc[mask, col] = None
        elif erasure_strategy == "hash":
            df.loc[mask, col] = df.loc[mask, col].apply(lambda v: _hash_value(v, hash_salt))
        elif erasure_strategy == "redact":
            df.loc[mask, col] = "***REDACTED***"

    logger.info(f"GDPR erasure ({erasure_strategy}): affected {affected_count} rows, {len(present_pii)} PII columns.")
    return df


def _mask_polars(df, pii_columns: List[str], strategy: str, hash_salt: str):
    """Mask all PII columns in a Polars DataFrame (all rows)."""
    import polars as pl

    if isinstance(df, pl.LazyFrame):
        df = df.collect()

    present_pii = [c for c in pii_columns if c in df.columns]
    if not present_pii:
        logger.info("No PII columns found in dataframe; nothing to mask.")
        return df

    for col in present_pii:
        if strategy == "nullify":
            df = df.with_columns(pl.lit(None).alias(col))
        elif strategy == "hash":
            df = df.with_columns(
                pl.col(col)
                .cast(pl.Utf8)
                .map_elements(
                    lambda v, _salt=hash_salt: _hash_value(v, _salt),
                    return_dtype=pl.Utf8,
                )
                .alias(col)
            )
        elif strategy == "redact":
            df = df.with_columns(pl.lit("***REDACTED***").alias(col))

    logger.info(f"PII masking ({strategy}): masked {len(present_pii)} columns across all rows.")
    return df


def _mask_pandas(df, pii_columns: List[str], strategy: str, hash_salt: str):
    """Mask all PII columns in a Pandas DataFrame (all rows)."""
    df = df.copy()
    present_pii = [c for c in pii_columns if c in df.columns]
    if not present_pii:
        logger.info("No PII columns found in dataframe; nothing to mask.")
        return df

    for col in present_pii:
        if strategy == "nullify":
            df[col] = None
        elif strategy == "hash":
            df[col] = df[col].apply(lambda v: _hash_value(v, hash_salt))
        elif strategy == "redact":
            df[col] = "***REDACTED***"

    logger.info(f"PII masking ({strategy}): masked {len(present_pii)} columns across {len(df)} rows.")
    return df


# ── Public API ───────────────────────────────────────────────────────────────


def forget_subjects(
    df: Any,
    contract: DataContract,
    subject_column: str,
    subject_ids: List[str],
    *,
    erasure_strategy: str = "nullify",
    hash_salt: str = "",
    audit: bool = True,
) -> Any:
    """
    GDPR Right-to-be-Forgotten: erase PII for specific data subjects.

    Identifies PII columns from the contract's ``pii: true`` field annotations
    and nullifies, hashes, or redacts values for matching subjects.

    Args:
        df: Input dataframe (Polars, Pandas, or DuckDB relation).
        contract: DataContract with PII-annotated fields.
        subject_column: Column containing the data subject identifier
                        (e.g., "customer_id", "email").
        subject_ids: List of subject identifiers to erase.
        erasure_strategy: How to erase PII values:
            - "nullify" (default): Set PII fields to NULL.
            - "hash": Replace with one-way SHA-256 hash.
            - "redact": Replace with "***REDACTED***".
        hash_salt: Salt for hashing (only used with "hash" strategy).
        audit: If True, log an audit entry about the erasure.

    Returns:
        DataFrame with PII erased for the specified subjects.

    Raises:
        ValueError: If subject_column not in dataframe or no PII fields defined.
    """
    if erasure_strategy not in ("nullify", "hash", "redact"):
        raise ValueError(f"Invalid erasure_strategy: {erasure_strategy}. Must be 'nullify', 'hash', or 'redact'.")

    pii_columns = _get_pii_column_names(contract)
    if not pii_columns:
        logger.warning(
            "No PII fields defined in contract. "
            "Mark fields with 'pii: true' in the model section to enable GDPR erasure."
        )
        return df

    if audit:
        logger.info(
            f"GDPR erasure request: strategy={erasure_strategy}, "
            f"subjects={len(subject_ids)}, pii_columns={pii_columns}, "
            f"timestamp={datetime.now(timezone.utc).isoformat()}"
        )

    # Detect frame type and dispatch
    try:
        import polars as pl

        if isinstance(df, (pl.DataFrame, pl.LazyFrame)):
            return _forget_polars(
                df,
                pii_columns,
                subject_column,
                subject_ids,
                erasure_strategy,
                hash_salt,
            )
    except ImportError:
        pass

    try:
        import pandas as pd

        if isinstance(df, pd.DataFrame):
            return _forget_pandas(
                df,
                pii_columns,
                subject_column,
                subject_ids,
                erasure_strategy,
                hash_salt,
            )
    except ImportError:
        pass

    # DuckDB relation → convert to pandas, process, convert back
    if hasattr(df, "fetchdf"):
        import pandas as pd

        pdf = df.fetchdf()
        result = _forget_pandas(pdf, pii_columns, subject_column, subject_ids, erasure_strategy, hash_salt)
        import duckdb

        return duckdb.from_df(result)

    raise TypeError(f"Unsupported dataframe type: {type(df)}")


def mask_pii_columns(
    df: Any,
    contract: DataContract,
    *,
    strategy: str = "nullify",
    hash_salt: str = "",
    columns: Optional[List[str]] = None,
) -> Any:
    """
    Mask all PII columns across all rows in a dataframe.

    Useful for creating anonymised datasets for development, testing,
    or analytics where PII is not needed.

    Args:
        df: Input dataframe (Polars, Pandas, or DuckDB relation).
        contract: DataContract with PII-annotated fields.
        strategy: Masking strategy:
            - "nullify" (default): Set to NULL.
            - "hash": One-way SHA-256 hash (preserves referential integrity).
            - "redact": Replace with "***REDACTED***".
        hash_salt: Salt for hashing.
        columns: Optional list of specific columns to mask (overrides
                 contract PII annotations).

    Returns:
        DataFrame with PII columns masked.
    """
    if strategy not in ("nullify", "hash", "redact"):
        raise ValueError(f"Invalid strategy: {strategy}. Must be 'nullify', 'hash', or 'redact'.")

    pii_columns = columns or _get_pii_column_names(contract)
    if not pii_columns:
        logger.warning("No PII columns to mask.")
        return df

    logger.info(f"Masking PII columns: {pii_columns} with strategy '{strategy}'")

    try:
        import polars as pl

        if isinstance(df, (pl.DataFrame, pl.LazyFrame)):
            return _mask_polars(df, pii_columns, strategy, hash_salt)
    except ImportError:
        pass

    try:
        import pandas as pd

        if isinstance(df, pd.DataFrame):
            return _mask_pandas(df, pii_columns, strategy, hash_salt)
    except ImportError:
        pass

    if hasattr(df, "fetchdf"):
        import pandas as pd

        pdf = df.fetchdf()
        result = _mask_pandas(pdf, pii_columns, strategy, hash_salt)
        import duckdb

        return duckdb.from_df(result)

    raise TypeError(f"Unsupported dataframe type: {type(df)}")


def generate_erasure_report(
    contract: DataContract,
    subject_column: str,
    subject_ids: List[str],
    erasure_strategy: str = "nullify",
    affected_rows: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Generate a GDPR-compliant audit report for an erasure operation.

    Args:
        contract: DataContract.
        subject_column: Column used for subject identification.
        subject_ids: List of erased subject IDs.
        erasure_strategy: Strategy used.
        affected_rows: Number of rows affected (if known).

    Returns:
        Audit report dictionary.
    """
    pii_columns = _get_pii_column_names(contract)
    contract_name = contract.info.title if contract.info else "unknown"

    return {
        "report_type": "gdpr_erasure",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "contract": contract_name,
        "subject_column": subject_column,
        "subjects_erased": len(subject_ids),
        "erasure_strategy": erasure_strategy,
        "pii_columns_affected": pii_columns,
        "affected_rows": affected_rows,
        "compliance_note": (
            "PII data has been erased in accordance with GDPR Article 17 "
            "(Right to Erasure). This report should be retained for audit purposes."
        ),
    }
