from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from typing import Iterable

from rdkit import Chem

from .errors import MacrolactoneDetectionError, RingNumberingError
from .models import RingNumberingResult


VALID_RING_SIZES = tuple(range(12, 21))


@dataclass(frozen=True)
class DetectedMacrolactone:
    ring_size: int
    ring_atoms: tuple[int, ...]
    carbonyl_carbon_idx: int
    ester_oxygen_idx: int


def ensure_mol(mol_input: str | Chem.Mol) -> tuple[Chem.Mol, str]:
    if isinstance(mol_input, Chem.Mol):
        mol = Chem.Mol(mol_input)
    else:
        mol = Chem.MolFromSmiles(mol_input)
        if mol is None:
            raise MacrolactoneDetectionError(f"Unable to parse SMILES: {mol_input}")
    return mol, Chem.MolToSmiles(mol, isomericSmiles=True)


def find_macrolactone_candidates(
    mol: Chem.Mol,
    ring_size: int | None = None,
) -> list[DetectedMacrolactone]:
    candidates: list[DetectedMacrolactone] = []
    seen: set[tuple[tuple[int, ...], int, int]] = set()

    for ring in mol.GetRingInfo().AtomRings():
        if len(ring) not in VALID_RING_SIZES:
            continue
        if ring_size is not None and len(ring) != ring_size:
            continue

        ester_match = _find_ring_ester_atoms(mol, ring)
        if ester_match is None:
            continue

        carbonyl_carbon_idx, ester_oxygen_idx = ester_match
        key = (tuple(sorted(ring)), carbonyl_carbon_idx, ester_oxygen_idx)
        if key in seen:
            continue
        seen.add(key)
        candidates.append(
            DetectedMacrolactone(
                ring_size=len(ring),
                ring_atoms=tuple(ring),
                carbonyl_carbon_idx=carbonyl_carbon_idx,
                ester_oxygen_idx=ester_oxygen_idx,
            )
        )

    return sorted(
        candidates,
        key=lambda candidate: (
            candidate.ring_size,
            tuple(sorted(candidate.ring_atoms)),
            candidate.carbonyl_carbon_idx,
            candidate.ester_oxygen_idx,
        ),
    )


def build_numbering_result(mol: Chem.Mol, candidate: DetectedMacrolactone) -> RingNumberingResult:
    ring_atoms = list(candidate.ring_atoms)
    ring_atom_set = set(ring_atoms)
    adjacency = {
        atom_idx: sorted(
            neighbor.GetIdx()
            for neighbor in mol.GetAtomWithIdx(atom_idx).GetNeighbors()
            if neighbor.GetIdx() in ring_atom_set
        )
        for atom_idx in ring_atoms
    }

    ordered_atoms = [candidate.carbonyl_carbon_idx, candidate.ester_oxygen_idx]
    visited = set(ordered_atoms)
    previous_atom = candidate.carbonyl_carbon_idx
    current_atom = candidate.ester_oxygen_idx

    while len(ordered_atoms) < candidate.ring_size:
        next_atoms = [
            atom_idx
            for atom_idx in adjacency[current_atom]
            if atom_idx != previous_atom and atom_idx not in visited
        ]
        if len(next_atoms) != 1:
            raise RingNumberingError("Unable to determine a unique ring traversal order.")
        next_atom = next_atoms[0]
        ordered_atoms.append(next_atom)
        visited.add(next_atom)
        previous_atom, current_atom = current_atom, next_atom

    if candidate.carbonyl_carbon_idx not in adjacency[current_atom]:
        raise RingNumberingError("Ring traversal did not close on the lactone carbonyl carbon.")

    position_to_atom = {position: atom_idx for position, atom_idx in enumerate(ordered_atoms, start=1)}
    atom_to_position = {atom_idx: position for position, atom_idx in position_to_atom.items()}

    return RingNumberingResult(
        ring_size=candidate.ring_size,
        ring_atoms=sorted(ring_atoms),
        ordered_atoms=ordered_atoms,
        carbonyl_carbon_idx=candidate.carbonyl_carbon_idx,
        ester_oxygen_idx=candidate.ester_oxygen_idx,
        atom_to_position=atom_to_position,
        position_to_atom=position_to_atom,
    )


def collect_side_chain_atoms(
    mol: Chem.Mol,
    start_atom_idx: int,
    ring_atom_indices: Iterable[int],
) -> list[int]:
    ring_atom_set = set(ring_atom_indices)
    queue: deque[int] = deque([start_atom_idx])
    visited: set[int] = set()
    side_chain_atoms: list[int] = []

    while queue:
        atom_idx = queue.popleft()
        if atom_idx in visited or atom_idx in ring_atom_set:
            continue
        visited.add(atom_idx)
        side_chain_atoms.append(atom_idx)

        atom = mol.GetAtomWithIdx(atom_idx)
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            if neighbor_idx not in visited and neighbor_idx not in ring_atom_set:
                queue.append(neighbor_idx)

    return side_chain_atoms


def is_intrinsic_lactone_neighbor(
    mol: Chem.Mol,
    candidate: DetectedMacrolactone,
    ring_atom_idx: int,
    neighbor_idx: int,
) -> bool:
    if ring_atom_idx != candidate.carbonyl_carbon_idx:
        return False
    bond = mol.GetBondBetweenAtoms(ring_atom_idx, neighbor_idx)
    neighbor = mol.GetAtomWithIdx(neighbor_idx)
    return bond is not None and bond.GetBondType() == Chem.BondType.DOUBLE and neighbor.GetAtomicNum() == 8


def build_fragment_smiles(
    mol: Chem.Mol,
    side_chain_atoms: list[int],
    side_chain_start_idx: int,
    ring_atom_idx: int,
    dummy_isotope: int,
) -> str:
    fragment_mol = Chem.RWMol()
    old_to_new: dict[int, int] = {}

    for old_idx in side_chain_atoms:
        old_atom = mol.GetAtomWithIdx(old_idx)
        new_atom = Chem.Atom(old_atom)
        new_idx = fragment_mol.AddAtom(new_atom)
        old_to_new[old_idx] = new_idx

    for old_idx in side_chain_atoms:
        atom = mol.GetAtomWithIdx(old_idx)
        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            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(),
                )

    bond_to_ring = mol.GetBondBetweenAtoms(ring_atom_idx, side_chain_start_idx)
    if bond_to_ring is None:
        raise RingNumberingError("Missing ring-to-side-chain bond during fragmentation.")

    dummy_atom = Chem.Atom(0)
    dummy_atom.SetIsotope(dummy_isotope)
    dummy_idx = fragment_mol.AddAtom(dummy_atom)
    fragment_mol.AddBond(
        old_to_new[side_chain_start_idx],
        dummy_idx,
        bond_to_ring.GetBondType(),
    )

    fragment = fragment_mol.GetMol()
    Chem.SanitizeMol(fragment)
    return Chem.MolToSmiles(fragment, isomericSmiles=True)


def _find_ring_ester_atoms(mol: Chem.Mol, ring_atoms: Iterable[int]) -> tuple[int, int] | None:
    ring_atom_set = set(ring_atoms)
    matches: list[tuple[int, int]] = []

    for atom_idx in ring_atom_set:
        atom = mol.GetAtomWithIdx(atom_idx)
        if atom.GetAtomicNum() != 6:
            continue

        has_carbonyl_oxygen = any(
            bond.GetBondType() == Chem.BondType.DOUBLE and bond.GetOtherAtom(atom).GetAtomicNum() == 8
            for bond in atom.GetBonds()
        )
        if not has_carbonyl_oxygen:
            continue

        for neighbor in atom.GetNeighbors():
            neighbor_idx = neighbor.GetIdx()
            bond = mol.GetBondBetweenAtoms(atom_idx, neighbor_idx)
            if (
                neighbor_idx in ring_atom_set
                and neighbor.GetAtomicNum() == 8
                and bond is not None
                and bond.GetBondType() == Chem.BondType.SINGLE
            ):
                ring_neighbors = [
                    ring_neighbor.GetIdx()
                    for ring_neighbor in neighbor.GetNeighbors()
                    if ring_neighbor.GetIdx() in ring_atom_set
                ]
                if len(ring_neighbors) == 2:
                    matches.append((atom_idx, neighbor_idx))

    if not matches:
        return None
    return sorted(matches)[0]
