"""
PROTAC Builder — Abstract Base Class

Three-component assembly: warhead + linker + E3 ligand.

All methods accept Chem.Mol objects. The DeFoG graph → Chem.Mol conversion
happens once upstream; this class never parses SMILES on method inputs.
Assembly output (assembled PROTAC) is returned as a SMILES string — this is
RDKit's native output from reaction SMARTS and is fine as an output format.
"""

from abc import ABC, abstractmethod
from typing import Optional
from rdkit import Chem
from rdkit.Chem import AllChem, Descriptors


class PROTAC_Builder(ABC):
    """
    Abstract base class for PROTAC builders.

    Subclasses encode one specific warhead + E3 ligand + linker chemistry
    combination (e.g. erlotinib-COOH + VH032-COOH + diamine linker).

    Parameters
    ----------
    warhead : Chem.Mol
        Fixed warhead molecule.
    e3 : Chem.Mol
        Fixed E3 ligand molecule.
    warhead_name / e3_name : str
        Display names for logging.
    """

    def __init__(self, warhead: Chem.Mol, e3: Chem.Mol,
                 warhead_name: str = "Warhead", e3_name: str = "E3 Ligand"):
        if warhead is None:
            raise ValueError("warhead must be a Chem.Mol object")
        if e3 is None:
            raise ValueError("e3 must be a Chem.Mol object")

        self.warhead_mol = warhead
        self.e3_mol = e3
        self.warhead_name = warhead_name
        self.e3_name = e3_name

        # Canonical SMILES stored for display / serialization only
        self.warhead_smiles = Chem.MolToSmiles(warhead)
        self.e3_smiles = Chem.MolToSmiles(e3)

        self._verify_warhead_reactive_group()
        self._verify_e3_reactive_group()

    # ------------------------------------------------------------------
    # Abstract interface
    # ------------------------------------------------------------------

    @abstractmethod
    def _verify_warhead_reactive_group(self):
        """Confirm warhead has the expected reactive group for linker attachment."""
        pass

    @abstractmethod
    def _verify_e3_reactive_group(self):
        """Confirm E3 ligand has the expected reactive group for linker attachment."""
        pass

    @abstractmethod
    def find_linker_attachment_sites(self, linker_mol: Chem.Mol) -> list:
        """
        Find attachment sites on the linker for both warhead and E3 ends.

        Returns list of dicts:
        [{'warhead_site': atom_idx, 'e3_site': atom_idx, ...}]
        """
        pass

    @abstractmethod
    def _assemble_protac(self, linker_mol: Chem.Mol, warhead_site_idx: int,
                         e3_site_idx: int) -> Optional[str]:
        """
        Perform the two bond-forming reactions.
        Returns SMILES of the assembled PROTAC, or None on failure.
        """
        pass

    @abstractmethod
    def check_attachment_site_compatibility(self, protac_mol: Chem.Mol) -> bool:
        """Verify assembled PROTAC has expected bond signatures at both ends."""
        pass

    @abstractmethod
    def _basic_validation(self, protac_mol: Chem.Mol) -> dict:
        """
        Run PROTAC-specific validation checks.
        Must include '_passed': bool in the returned dict.
        """
        pass

    # ------------------------------------------------------------------
    # Shared public interface
    # ------------------------------------------------------------------

    def assemble(self, linker_mol: Chem.Mol, return_all: bool = True,
                 amide_only: bool = False):
        """
        Assemble warhead + linker + E3 into a full PROTAC.

        Parameters
        ----------
        linker_mol : Chem.Mol
            Generated linker as a molecular graph.
        return_all : bool
            True → return all unique assembled PROTACs (list[str]).
            False → return first valid PROTAC (str).
        amide_only : bool
            If True, only attempt the amide (COOH/CHO) assembly path and
            skip the CH3 alkylation path. Use for strict final validation.

        Returns
        -------
        list[str] if return_all else str | None
        """
        sites = self.find_linker_attachment_sites(linker_mol, amide_only=amide_only)
        if not sites:
            raise ValueError("No attachment sites found on linker")

        protacs = []
        seen = set()
        for site in sites:
            try:
                smiles = self._assemble_protac(
                    linker_mol, site["warhead_site"], site["e3_site"],
                    warhead_type=site.get("warhead_type", "nh2"),
                    e3_type=site.get("e3_type", "nh2"),
                    warhead_restored=site.get("warhead_restored", False),
                    e3_restored=site.get("e3_restored", False),
                )
                if smiles and smiles not in seen:
                    seen.add(smiles)
                    protacs.append(smiles)
            except Exception:
                continue

        if not protacs:
            raise ValueError("Assembly failed at all linker sites")

        return protacs if return_all else protacs[0]

    def validate_protac(self, protac_mol: Chem.Mol):
        """
        Returns (passed: bool, checks: dict) for an assembled PROTAC Mol.
        """
        checks = self._basic_validation(protac_mol)
        passed = checks.pop("_passed", False)
        return passed, checks

    def get_protac_properties(self, protac_mol: Chem.Mol) -> dict:
        """Compute molecular properties for an assembled PROTAC."""
        return {
            "mw":          round(Descriptors.MolWt(protac_mol), 2),
            "logp":        round(Descriptors.MolLogP(protac_mol), 2),
            "tpsa":        round(Descriptors.TPSA(protac_mol), 2),
            "heavy_atoms": protac_mol.GetNumHeavyAtoms(),
            "hbd":         Descriptors.NumHDonors(protac_mol),
            "hba":         Descriptors.NumHAcceptors(protac_mol),
            "rotbonds":    Descriptors.NumRotatableBonds(protac_mol),
        }

    def batch_assemble(self, linker_mols: list[Chem.Mol],
                       return_all: bool = True) -> dict:
        """Assemble PROTACs for a list of linker Mol objects."""
        results = {"protacs": [], "successful": [], "failed": [], "total": 0}
        for i, mol in enumerate(linker_mols):
            try:
                p = self.assemble(mol, return_all=return_all)
                results["protacs"].append(p)
                results["successful"].append(i)
                results["total"] += len(p) if return_all else 1
            except Exception as e:
                results["protacs"].append(None)
                results["failed"].append((i, str(e)))
        print(f"Assembled {results['total']} PROTAC(s) from "
              f"{len(results['successful'])}/{len(linker_mols)} linkers")
        return results