from __future__ import annotations

from pathlib import Path
from typing import Iterable

from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D

from ._core import build_numbering_result, ensure_mol, find_macrolactone_candidates
from .errors import AmbiguousMacrolactoneError, MacrolactoneDetectionError


def numbered_molecule_svg(
    mol_input: str | Chem.Mol,
    ring_size: int | None = None,
    size: tuple[int, int] = (800, 800),
    allowed_ring_atom_types: list[str] | None = None,
    show_atom_labels: bool = True,
) -> str:
    mol, _ = ensure_mol(mol_input)
    numbering = _get_visualization_numbering(
        mol,
        ring_size=ring_size,
        allowed_ring_atom_types=allowed_ring_atom_types,
    )
    drawer = rdMolDraw2D.MolDraw2DSVG(*size)
    _draw_numbered_molecule(
        mol=mol,
        drawer=drawer,
        position_to_atom=numbering.position_to_atom,
        show_atom_labels=show_atom_labels,
    )
    return drawer.GetDrawingText()


def save_numbered_molecule_png(
    mol_input: str | Chem.Mol,
    output_path: str | Path,
    ring_size: int | None = None,
    size: tuple[int, int] = (800, 800),
    allowed_ring_atom_types: list[str] | None = None,
    show_atom_labels: bool = True,
    dpi: int = 600,
) -> Path:
    del dpi
    mol, _ = ensure_mol(mol_input)
    numbering = _get_visualization_numbering(
        mol,
        ring_size=ring_size,
        allowed_ring_atom_types=allowed_ring_atom_types,
    )
    drawer = rdMolDraw2D.MolDraw2DCairo(*size)
    _draw_numbered_molecule(
        mol=mol,
        drawer=drawer,
        position_to_atom=numbering.position_to_atom,
        show_atom_labels=show_atom_labels,
    )
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(drawer.GetDrawingText())
    return path


def fragment_svg(fragment_or_smiles: str | Chem.Mol, size: tuple[int, int] = (400, 400)) -> str:
    mol, _ = ensure_mol(fragment_or_smiles)
    drawer = rdMolDraw2D.MolDraw2DSVG(*size)
    _draw_plain_molecule(mol, drawer)
    return drawer.GetDrawingText()


def save_fragment_png(
    fragment_or_smiles: str | Chem.Mol,
    output_path: str | Path,
    size: tuple[int, int] = (400, 400),
    dpi: int = 600,
) -> Path:
    del dpi
    mol, _ = ensure_mol(fragment_or_smiles)
    drawer = rdMolDraw2D.MolDraw2DCairo(*size)
    _draw_plain_molecule(mol, drawer)
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(drawer.GetDrawingText())
    return path


def _draw_numbered_molecule(
    mol: Chem.Mol,
    drawer: rdMolDraw2D.MolDraw2DSVG | rdMolDraw2D.MolDraw2DCairo,
    position_to_atom: dict[int, int],
    show_atom_labels: bool,
) -> None:
    draw_mol = Chem.Mol(mol)
    draw_options = drawer.drawOptions()
    if show_atom_labels:
        for position, atom_idx in position_to_atom.items():
            draw_options.atomLabels[atom_idx] = str(position)

    highlight_atoms = list(position_to_atom.values())
    highlight_colors = {
        atom_idx: (0.96, 0.84, 0.48)
        for atom_idx in highlight_atoms
    }
    rdMolDraw2D.PrepareAndDrawMolecule(
        drawer,
        draw_mol,
        highlightAtoms=highlight_atoms,
        highlightAtomColors=highlight_colors,
    )
    drawer.FinishDrawing()


def _draw_plain_molecule(
    mol: Chem.Mol,
    drawer: rdMolDraw2D.MolDraw2DSVG | rdMolDraw2D.MolDraw2DCairo,
) -> None:
    rdMolDraw2D.PrepareAndDrawMolecule(drawer, Chem.Mol(mol))
    drawer.FinishDrawing()


def _get_visualization_numbering(
    mol: Chem.Mol,
    ring_size: int | None,
    allowed_ring_atom_types: list[str] | None,
):
    candidates = find_macrolactone_candidates(mol, ring_size=ring_size)
    if not candidates:
        requested = f"{ring_size}-membered " if ring_size is not None else ""
        raise MacrolactoneDetectionError(f"No valid {requested}macrolactone was detected.")

    if allowed_ring_atom_types is not None:
        allowed_atomic_numbers = _normalize_allowed_ring_atom_types(allowed_ring_atom_types)
        candidates = [
            candidate
            for candidate in candidates
            if _candidate_matches_allowed_ring_atom_types(
                mol,
                candidate,
                allowed_atomic_numbers,
            )
        ]
        if not candidates:
            raise ValueError("No macrolactone candidate matched the allowed ring atom types.")

    valid_ring_sizes = sorted({candidate.ring_size for candidate in candidates})
    if len(candidates) > 1 or len(valid_ring_sizes) > 1:
        raise AmbiguousMacrolactoneError(
            "Multiple valid macrolactone candidates were detected. Pass ring_size explicitly."
        )

    return build_numbering_result(mol, candidates[0])


def _normalize_allowed_ring_atom_types(atom_types: Iterable[str]) -> set[int]:
    periodic_table = Chem.GetPeriodicTable()
    normalized: set[int] = set()
    for atom_type in atom_types:
        if atom_type.isdigit():
            normalized.add(int(atom_type))
            continue
        atomic_number = periodic_table.GetAtomicNumber(atom_type.capitalize())
        if atomic_number <= 0:
            raise ValueError(f"Unsupported atom type: {atom_type}")
        normalized.add(atomic_number)
    return normalized


def _candidate_matches_allowed_ring_atom_types(
    mol: Chem.Mol,
    candidate,
    allowed_atomic_numbers: set[int],
) -> bool:
    numbering = build_numbering_result(mol, candidate)
    for position, atom_idx in numbering.position_to_atom.items():
        if position == 2:
            continue
        if mol.GetAtomWithIdx(atom_idx).GetAtomicNum() not in allowed_atomic_numbers:
            return False
    return True
