from rdkit import Chem
from rdkit.Chem import Crippen, Descriptors, QED


def _calc_props(smiles: str) -> dict | None:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None

    mw = Descriptors.MolWt(mol)
    logp = Crippen.MolLogP(mol)
    hbd = Descriptors.NumHDonors(mol)
    hba = Descriptors.NumHAcceptors(mol)
    rot = Descriptors.NumRotatableBonds(mol)
    tpsa = Descriptors.TPSA(mol)
    qed = float(QED.qed(mol))

    return {
        "mw": float(mw),
        "logp": float(logp),
        "hbd": int(hbd),
        "hba": int(hba),
        "rotatable_bonds": int(rot),
        "tpsa": float(tpsa),
        "qed": float(qed),
    }


def filter_molecules(candidates: list[dict]) -> list[dict]:
    """Apply basic drug-likeness heuristics (Lipinski-ish + QED)."""

    filtered: list[dict] = []
    for c in candidates:
        smiles = (c or {}).get("smiles")
        if not smiles or not isinstance(smiles, str):
            continue

        props = _calc_props(smiles)
        if props is None:
            out = dict(c)
            out["passes_filter"] = False
            out["parse_error"] = True
            filtered.append(out)
            continue

        passes = True
        if not (150 <= props["mw"] <= 500):
            passes = False
        if props["logp"] > 5:
            passes = False
        if props["hbd"] > 5:
            passes = False
        if props["hba"] > 10:
            passes = False
        if props["qed"] < 0.3:
            passes = False

        out = dict(c)
        out.update(props)
        out["passes_filter"] = passes
        filtered.append(out)

    filtered.sort(key=lambda x: (x.get("score") is None, -(x.get("score") or 0), -x.get("qed", 0)))
    return filtered