from __future__ import annotations

from pathlib import Path

from rdkit import Chem

from macro_lactone_toolkit.visualization import save_fragment_png, save_numbered_molecule_png


def get_output_paths(output_dir: Path, source_id: str, ring_size: int, classification: str) -> dict:
    """Get organized output paths for a molecule."""
    ring_dir = output_dir / f"ring_size_{ring_size}"

    if classification == "standard_macrolactone":
        base_dir = ring_dir / "standard"
        numbered_dir = base_dir / "numbered"
        sidechains_dir = base_dir / "sidechains" / source_id
    elif classification == "non_standard_macrocycle":
        base_dir = ring_dir / "non_standard" / "original"
        numbered_dir = base_dir
        sidechains_dir = None
    else:
        base_dir = ring_dir / "rejected" / "original"
        numbered_dir = base_dir
        sidechains_dir = None

    numbered_dir.mkdir(parents=True, exist_ok=True)
    if sidechains_dir:
        sidechains_dir.mkdir(parents=True, exist_ok=True)

    return {
        "numbered_image": numbered_dir / f"{source_id}_numbered.png",
        "sidechains_dir": sidechains_dir,
    }


def save_numbered_molecule(
    smiles: str,
    output_path: Path,
    ring_size: int | None = None,
    size: tuple[int, int] = (800, 800),
) -> Path | None:
    """Save numbered molecule image."""
    try:
        return save_numbered_molecule_png(
            smiles,
            output_path,
            ring_size=ring_size,
            size=size,
        )
    except Exception as e:
        print(f"Failed to save numbered image: {e}")
        return None


def save_fragment_images(
    fragments: list,
    output_dir: Path,
    source_id: str,
    size: tuple[int, int] = (400, 400),
) -> list[str]:
    """Save fragment images and return paths."""
    paths = []

    for i, fragment in enumerate(fragments):
        try:
            output_path = output_dir / f"{source_id}_frag_{i}_pos{fragment.cleavage_position}.png"
            save_fragment_png(fragment.fragment_smiles_plain, output_path, size=size)
            paths.append(str(output_path.relative_to(output_dir.parent.parent)))
        except Exception as e:
            print(f"Failed to save fragment {i}: {e}")
            paths.append(None)

    return paths
