#!/usr/bin/env python3

from argparse import ArgumentParser
from pathlib import Path
from typing import TextIO

from ase import io, Atoms


def assign_connectivity(atoms: Atoms) -> tuple[list[tuple[int, int, int]], list[int]]:
    """Return bonds as (i, j, order) with 0-based atom indices and per-atom formal charges.

    Bond orders follow V2000 bond types: 1 single, 2 double, 3 triple, 4 aromatic.
    Dummy placeholder: no bonds, all atoms neutral.
    """
    return [], [0] * len(atoms)


def write_sdf_record(
    fh: TextIO, atoms: Atoms, bonds: list[tuple[int, int, int]], charges: list[int]
) -> None:
    """Append one V2000 molfile record."""
    if len(atoms) > 999 or len(bonds) > 999:
        raise ValueError(f"{atoms.info.get('name')}: too large for V2000")

    fh.write(f"{atoms.info.get('name', '')}\n     assign         3D\n\n")
    fh.write(f"{len(atoms):3d}{len(bonds):3d}  0  0  0  0  0  0  0  0999 V2000\n")
    for sym, (x, y, z) in zip(atoms.get_chemical_symbols(), atoms.positions):
        fh.write(f"{x:10.4f}{y:10.4f}{z:10.4f} {sym:<3} 0" + "  0" * 11 + "\n")
    for i, j, order in bonds:
        fh.write(f"{i + 1:3d}{j + 1:3d}{order:3d}  0\n")

    charged = [(i + 1, q) for i, q in enumerate(charges) if q]
    for k in range(0, len(charged), 8):
        chunk = charged[k : k + 8]
        fh.write(
            f"M  CHG{len(chunk):3d}"
            + "".join(f" {i:3d} {q:3d}" for i, q in chunk)
            + "\n"
        )
    fh.write("M  END\n")
    fh.write("$$$$\n")


def main(input: Path, output: Path):
    with open(output, "w") as fh:
        for atoms in io.iread(input, index=slice(None)):
            bonds, charges = assign_connectivity(atoms)
            write_sdf_record(fh, atoms, bonds, charges)


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Assign connectivity to xyz frames and write SDF."
    )
    parser.add_argument(
        "-i", "--input", type=Path, required=True, help="multi-frame (ext)xyz file"
    )
    parser.add_argument(
        "-o", "--output", type=Path, required=True, help="output SDF file"
    )
    args = parser.parse_args()
    main(args.input, args.output)