"""
Mol Writer
==========

"""

from __future__ import annotations

import pathlib
import typing

from stk._internal.molecule import Molecule
from stk._internal.utilities.utilities import OneOrMany


class MolWriter:
    """
    A writer class for V3000 ``.mol`` files.

    Examples:

        *Writing to a File*

        .. testcode:: writing-to-a-file

            import stk

            bb1 = stk.BuildingBlock('BrCCBr', [stk.BromoFactory()])

            writer = stk.MolWriter()
            writer.write(molecule=bb1, path='bb1.mol')

        .. testcode:: writing-to-a-file
            :hide:

            import os

            assert os.path.exists('bb1.mol')

        .. testcleanup:: writing-to-a-file

            os.remove('bb1.mol')

    """

    def _write_content(
        self,
        molecule: Molecule,
        atom_ids: typing.Optional[OneOrMany[int]],
    ) -> str:
        if atom_ids is None:
            atom_ids = range(molecule.get_num_atoms())
        elif isinstance(atom_ids, int):
            atom_ids = (atom_ids,)

        atom_lines = []
        coords = molecule.get_position_matrix()
        # This set gets used by bonds.
        id_map = {}
        for new_atom_id, old_atom_id in enumerate(atom_ids, 1):
            id_map[old_atom_id] = new_atom_id

            x, y, z = (i for i in coords[old_atom_id])
            (atom,) = molecule.get_atoms(atom_ids=old_atom_id)
            symbol = atom.__class__.__name__
            charge_ = atom.get_charge()
            charge = f" CHG={charge_}" if charge_ else ""
            atom_lines.append(
                f"M  V30 {new_atom_id} {symbol} {x:.4f} "
                f"{y:.4f} {z:.4f} 0{charge}\n"
            )
        atom_block = "".join(atom_lines)

        bond_lines: list[str] = []
        for bond in molecule.get_bonds():
            a1 = bond.get_atom1().get_id()
            a2 = bond.get_atom2().get_id()
            if a1 in id_map and a2 in id_map:
                bond_lines.append(
                    f"M  V30 {len(bond_lines) + 1} "
                    f"{int(bond.get_order())} "
                    f"{id_map[a1]} {id_map[a2]}\n"
                )

        num_bonds = len(bond_lines)
        bond_block = "".join(bond_lines)
        return (
            "\n"
            "     RDKit          3D\n"
            "\n"
            "  0  0  0  0  0  0  0  0  0  0999 V3000\n"
            "M  V30 BEGIN CTAB\n"
            f"M  V30 COUNTS {len(id_map)} {num_bonds} 0 0 0\n"
            "M  V30 BEGIN ATOM\n"
            f"{atom_block}"
            "M  V30 END ATOM\n"
            "M  V30 BEGIN BOND\n"
            f"{bond_block}"
            "M  V30 END BOND\n"
            "M  V30 END CTAB\n"
            "M  END\n"
            "\n"
            "$$$$\n"
        )

    def to_string(
        self,
        molecule: Molecule,
        atom_ids: typing.Optional[OneOrMany[int]] = None,
    ) -> str:
        """
        Get a V3000 ``.mol`` file format string of `molecule`.

        Parameters:

            molecule:
                Molecule to write to V3000 ``.mol`` format.

            atom_ids:
                The atom ids of atoms to write. Can be a single
                :class:`int`, if a single atom is to be used, or
                ``None``, if all atoms are to be used.

        Returns:

            String in V3000 ``.mol`` file format.

        """

        return self._write_content(molecule, atom_ids)

    def write(
        self,
        molecule: Molecule,
        path: pathlib.Path | str,
        atom_ids: typing.Optional[OneOrMany[int]] = None,
    ) -> None:
        """
        Write `molecule` to V3000 ``.mol`` file format.

        Parameters:

            molecule:
                Molecule to write to V3000 ``.mol`` format.

            path:
                The full path to the file being written.

            atom_ids:
                The atom ids of atoms to write. Can be a single
                :class:`int`, if a single atom is to be used, or
                ``None``, if all atoms are to be used.

        """

        with open(path, "w") as f:
            f.write(self._write_content(molecule, atom_ids))
