"""
Molecule standardisation pipeline for Polars LazyFrames.

Design:
- Dedup on raw SMILES (string level) before any RDKit work
- Multiprocessing for mol parsing, standardisation, and representation
- Lazy Polars throughout except for the unavoidable collect at dedup boundary
- Re-join results back to original frame (preserving all rows)
"""

from __future__ import annotations

import logging as lg
from collections.abc import Callable
from concurrent.futures import ProcessPoolExecutor
from functools import partial
from multiprocessing import get_context as _get_mp_context
from typing import Any, Sequence, cast

import polars as pl
from rdkit.Chem import Mol

from fairfetched.utils._track import track

from ._types import MolToStrFn, RepFns, StdFn, StrToMolFn, T_in, T_out

# Spawn context: safe with JAX (no fork)
_SPAWN_CONTEXT_MP = _get_mp_context("spawn")


def _apply_str_to_mol(molstring: str, read_fn: StrToMolFn) -> tuple[str, Mol | None]:
    try:
        return molstring, read_fn(molstring)
    except Exception:
        return molstring, None


def _apply_standardisation_fns(
    str_mol_tuple: tuple[str, Mol | None], std_fns: list[StdFn]
) -> tuple[str, Mol | None]:
    molstring, mol = str_mol_tuple
    if mol is None:
        return molstring, None
    for fn in std_fns:
        try:
            mol = fn(mol)
        except Exception as e:
            lg.error(f"failed standardisation of {molstring} with exception: {e}")
            mol = None
        if mol is None:
            lg.error(f"standardisation of {molstring} returned None")
            return molstring, None
    return molstring, mol


def _apply_mol_to_reprs(
    str_mol_tuple: tuple[str, Mol | None], rep_fns: dict[str, MolToStrFn]
) -> tuple[str, dict[str, Any | None]]:
    og_string, mol = str_mol_tuple
    if mol is None:
        return og_string, {col: None for col in rep_fns}
    results = {}
    for col, fn in rep_fns.items():
        try:
            results[col] = fn(mol)
        except Exception as e:
            lg.error(f"failed to generate {col} for {og_string} with exception {e}")
            results[col] = None
    return og_string, results


def _imap_mp(
    fn: Callable[[T_in], T_out],
    items: Sequence[T_in],
    n_workers: int,
    chunksize: int,
    desc: str | None = None,
) -> list[T_out]:
    """
    Spawn-safe parallel map via ProcessPoolExecutor.
    Submits chunks as single futures to reduce IPC overhead.
    Results are returned in original order.
    tqdm tracks chunk completion (not individual mols, to avoid IPC per-item).
    """
    if not items:
        return []

    with ProcessPoolExecutor(
        max_workers=n_workers, mp_context=_SPAWN_CONTEXT_MP
    ) as executor:
        result = track(
            executor.map(fn, items, chunksize=chunksize), total=len(items), desc=desc
        )
        return list(result)


# ── steps ─────────────────────────────────────────────────────────────


def _imap_mp_standardisation_fns(
    parsed: list[tuple[str, Mol | None]],
    std_fns: list[StdFn],
    n_workers: int,
    chunksize: int,
) -> list[tuple[str, Mol | None]]:
    """Apply standardisation functions sequentially per mol, in parallel across mols."""
    if not std_fns:
        return parsed
    fn = partial(_apply_standardisation_fns, std_fns=std_fns)
    return _imap_mp(fn, parsed, n_workers, chunksize, desc="standardising")


def build_result_frame(
    results: list[tuple[str, dict[str, Any | None]]],
    molstring_col: str,
) -> pl.LazyFrame:
    if not results:
        return pl.LazyFrame({molstring_col: []})
    og_molstring_list, repr_dicts = zip(*results)
    data: dict[str, list] = {molstring_col: list(og_molstring_list)}
    for col in repr_dicts[0]:
        data[col] = [d[col] for d in repr_dicts]
    return pl.LazyFrame(data)


def add_results(
    lf: pl.LazyFrame,
    results: list[tuple[str, dict[str, Any | None]]],
    molstring_col: str,
) -> pl.LazyFrame:
    if not results:
        return lf
    og_molstring_list, repr_dicts = zip(*results)
    data: dict[str, list] = {molstring_col: list(og_molstring_list)}
    for col in repr_dicts[0]:
        data[col] = [d[col] for d in repr_dicts]
    return (
        lf.with_row_index()
        .join(pl.LazyFrame(data), on=molstring_col, how="left")
        .sort("index")
        .drop("index")
    )


def standardise_mol(
    lf: pl.LazyFrame,
    str_to_mol_fn: StrToMolFn,
    std_fns: list[StdFn],
    mol_to_repr_fns: RepFns,
    molstring_col: str = "smiles",
    output_col: str = "representation",  # ignored when dict passed
    suffix: str = "_original",
    workers: int | None = None,
    chunksize: int = 256,
) -> pl.LazyFrame:
    """
    Standardise a lazyframe

    Args:
        lf:         pl.LazyFrame
        read_fn:    str -> Mol | None
        std_fns:    list of Mol -> Mol | None, applied in order
        rep_fn:     Mol -> Any | None (e.g. canonical SMILES, fingerprint)
        molstring_col: input column name
        output_col: output column name
        n_workers:  number of processes (default: cpu_count)
        chunksize:  imap chunksize (tune for mol size vs overhead)

    Returns:
        pl.LazyFrame, collects only unique smiles

    Use:
        either as a function with lf as the first argument,
            ```
            lf = pl.scan_parquet("my_parquet.parquet")
            lf = standardise_mol(lf, *args, **kwargs)
            ```
        or through polar's `pipe`:
            ```
            lf = lf.pipe(standardise_mol, *args, **kwargs)
            ```
    Note:
        str_to_mol_fn, std_fns, and mol_to_repr_fns must be picklable.
        Use module-level Python wrapper functions, not bare RDKit/Boost.Python
        builtins (e.g. use `smiles_to_mol` not `MolFromSmiles` directly, or make a
        ```
        def mol_from_smiles(smi):
            return MolFromSmiles(smi)
        ```).
    """
    if not isinstance(mol_to_repr_fns, dict):
        mol_to_repr_fns: dict[str, MolToStrFn] = {output_col: mol_to_repr_fns}

    workers = workers or _SPAWN_CONTEXT_MP.cpu_count()

    unique_molstrings: list[str] = (
        lf.select(molstring_col)
        .unique()
        .collect()
        .get_column(molstring_col)  # ty: ignore[unresolved-attribute]
        .to_list()
    )
    read_fn = cast(
        Callable[[str], tuple[str, Mol | None]],
        partial(_apply_str_to_mol, read_fn=str_to_mol_fn),
    )
    to_representations_fn = cast(
        Callable[[tuple[str, Mol | None]], tuple[str, dict[str, Any | None]]],
        partial(_apply_mol_to_reprs, rep_fns=mol_to_repr_fns),
    )

    parsed = _imap_mp(read_fn, unique_molstrings, workers, chunksize, desc="parsing")
    standardised = _imap_mp_standardisation_fns(parsed, std_fns, workers, chunksize)
    results = _imap_mp(
        to_representations_fn,
        standardised,
        workers,
        chunksize,
        desc="converting to string representations",
    )
    return lf.pipe(add_results, results, molstring_col)


def standardisation_function(
    str_to_mol_fn: StrToMolFn,
    std_fns: list[StdFn],
    mol_to_repr_fns: RepFns,
    molstring_col: str = "smiles",
    output_col: str = "representation",
    suffix: str = "_original",
    workers: int | None = None,
    chunksize: int = 256,
) -> Callable[[pl.LazyFrame], pl.LazyFrame]:
    """
    Build a standardise() function from composable pieces.

    Args:
        read_fn:    str -> Mol | None
        std_fns:    list of Mol -> Mol | None, applied in order
        rep_fn:     Mol -> Any | None (e.g. canonical SMILES, fingerprint)
        molstring_col: input column name
        output_col: output column name
        n_workers:  number of processes (default: cpu_count)
        chunksize:  imap chunksize (tune for mol size vs overhead)

    Returns:
        standardise(lf: pl.LazyFrame) -> pl.LazyFrame
    """
    workers = workers or _SPAWN_CONTEXT_MP.cpu_count()

    return partial(
        standardise_mol,
        str_to_mol_fn=str_to_mol_fn,
        std_fns=std_fns,
        mol_to_repr_fns=mol_to_repr_fns,
        molstring_col=molstring_col,
        output_col=output_col,
        suffix=suffix,
        workers=workers,
        chunksize=chunksize,
    )
