"""
Side chain cleavage and fragment extraction for macrolactones.
"""
from rdkit import Chem
from rdkit.Chem import Descriptors, rdMolDescriptors
from typing import List, Tuple, Optional, Dict
from .ring_numbering import assign_ring_numbering, get_ring_position
from .fragment_dataclass import Fragment, MoleculeFragments


def identify_side_chains(mol: Chem.Mol, ring_atoms: List[int]) -> List[Tuple[int, int]]:
    """
    Identify side chains attached to the ring.
    
    Args:
        mol: RDKit molecule object
        ring_atoms: List of atom indices in the 16-membered ring
    
    Returns:
        List of tuples (ring_atom_idx, side_chain_first_atom_idx)
    """
    side_chains = []
    ring_atom_set = set(ring_atoms)
    
    for ring_atom_idx in ring_atoms:
        atom = mol.GetAtomWithIdx(ring_atom_idx)
        
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            
            # If neighbor is not in the ring, it's a side chain
            if neighbor_idx not in ring_atom_set:
                side_chains.append((ring_atom_idx, neighbor_idx))
    
    return side_chains


def extract_side_chain_fragment(mol: Chem.Mol, ring_atom_idx: int, side_chain_start_idx: int, 
                                 ring_atoms: List[int]) -> Optional[str]:
    """
    Extract a side chain fragment as a SMILES string.
    
    Args:
        mol: RDKit molecule object
        ring_atom_idx: Index of the ring atom where side chain is attached
        side_chain_start_idx: Index of the first atom in the side chain
        ring_atoms: List of all ring atom indices
    
    Returns:
        SMILES string of the fragment, or None if extraction failed
    """
    # Create an editable copy
    emol = Chem.RWMol(mol)
    
    # Find the bond between ring and side chain
    bond = mol.GetBondBetweenAtoms(ring_atom_idx, side_chain_start_idx)
    if bond is None:
        return None
    
    # Get all atoms in the side chain by BFS from side_chain_start_idx
    # but not crossing into the ring
    ring_atom_set = set(ring_atoms)
    visited = set()
    queue = [side_chain_start_idx]
    side_chain_atoms = []
    
    while queue:
        current_idx = queue.pop(0)
        if current_idx in visited:
            continue
        
        visited.add(current_idx)
        side_chain_atoms.append(current_idx)
        
        atom = mol.GetAtomWithIdx(current_idx)
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            
            # Only continue into non-ring atoms
            if neighbor_idx not in ring_atom_set and neighbor_idx not in visited:
                queue.append(neighbor_idx)
    
    if not side_chain_atoms:
        return None
    
    # Create a new molecule with only the side chain atoms
    # We'll use Chem.RWMol to build it
    fragment_mol = Chem.RWMol()
    
    # Map old atom indices to new ones
    old_to_new = {}
    
    # Add atoms
    for old_idx in side_chain_atoms:
        atom = mol.GetAtomWithIdx(old_idx)
        new_atom = Chem.Atom(atom.GetAtomicNum())
        new_atom.SetFormalCharge(atom.GetFormalCharge())
        new_atom.SetIsAromatic(atom.GetIsAromatic())
        new_idx = fragment_mol.AddAtom(new_atom)
        old_to_new[old_idx] = new_idx
    
    # Add bonds
    for old_idx in side_chain_atoms:
        atom = mol.GetAtomWithIdx(old_idx)
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            
            # Only add bonds within the side chain
            if neighbor_idx in old_to_new and old_idx < neighbor_idx:
                bond = mol.GetBondBetweenAtoms(old_idx, neighbor_idx)
                fragment_mol.AddBond(
                    old_to_new[old_idx],
                    old_to_new[neighbor_idx],
                    bond.GetBondType()
                )
    
    # Convert to molecule and get SMILES
    try:
        fragment_mol = fragment_mol.GetMol()
        Chem.SanitizeMol(fragment_mol)
        fragment_smiles = Chem.MolToSmiles(fragment_mol)
        return fragment_smiles
    except:
        # If sanitization fails, return None
        return None


def cleave_side_chains(mol: Chem.Mol, parent_smiles: str, parent_id: str, 
                       numbering: Dict[int, int]) -> MoleculeFragments:
    """
    Cleave all side chains from a macrolactone and create fragment records.
    
    Args:
        mol: RDKit molecule object
        parent_smiles: SMILES string of the parent molecule
        parent_id: Identifier for the parent molecule (e.g., "macro_0")
        numbering: Ring numbering dictionary from assign_ring_numbering
    
    Returns:
        MoleculeFragments object containing all fragments
    """
    mol_fragments = MoleculeFragments(
        parent_id=parent_id,
        parent_smiles=parent_smiles,
        ring_numbering={str(k): v for k, v in numbering.items()}  # Convert to str keys for JSON
    )
    
    ring_atoms = list(numbering.keys())
    side_chains = identify_side_chains(mol, ring_atoms)
    
    fragment_counter = 0
    
    for ring_atom_idx, side_chain_start_idx in side_chains:
        # Get the ring position
        cleavage_position = get_ring_position(numbering, ring_atom_idx)
        
        if cleavage_position is None:
            continue
        
        # Extract the fragment
        fragment_smiles = extract_side_chain_fragment(
            mol, ring_atom_idx, side_chain_start_idx, ring_atoms
        )
        
        if fragment_smiles is None:
            continue
        
        # Create fragment molecule for properties
        frag_mol = Chem.MolFromSmiles(fragment_smiles)
        if frag_mol is None:
            continue
        
        # Create Fragment object
        fragment = Fragment(
            fragment_smiles=fragment_smiles,
            parent_smiles=parent_smiles,
            cleavage_position=cleavage_position,
            fragment_id=f"{parent_id}_frag_{fragment_counter}",
            parent_id=parent_id,
            atom_count=frag_mol.GetNumAtoms(),
            molecular_weight=round(Descriptors.MolWt(frag_mol), 2)
        )
        
        mol_fragments.add_fragment(fragment)
        fragment_counter += 1
    
    return mol_fragments


def process_molecule(smiles: str, molecule_idx: int) -> Optional[MoleculeFragments]:
    """
    Process a single molecule: assign numbering and cleave side chains.
    
    Args:
        smiles: SMILES string of the molecule
        molecule_idx: Index number for naming (e.g., 0 -> "macro_0")
    
    Returns:
        MoleculeFragments object or None if processing failed
    """
    parent_id = f"macro_{molecule_idx}"
    
    # Parse molecule
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        print(f"Error: Could not parse SMILES for {parent_id}")
        return None
    
    # Assign ring numbering
    numbering = assign_ring_numbering(mol)
    if not numbering:
        print(f"Error: Could not assign ring numbering for {parent_id}")
        return None
    
    # Cleave side chains
    mol_fragments = cleave_side_chains(mol, smiles, parent_id, numbering)
    
    return mol_fragments

