import os


def _get_small_molecule_template_generator(mol, molecules):
    from openff.toolkit.topology import Molecule

    charge_model = molecules.get("charge_model", "am1bcc")
    ff_path = molecules.get("forcefield", "gaff-2.2.20")
    ff_type = molecules.get("forcefield_type", None)
    if ff_type is None:
        ff_name = os.path.basename(ff_path)
        if ff_name.startswith("gaff"):
            ff_type = "gaff"
        elif ff_name.startswith("openff"):
            ff_type = "openff"
        elif ff_name.startswith("espaloma"):
            ff_type = "espaloma"
        else:
            raise ValueError(f"Invalid small molecule force field: {ff_name}")
    else:
        ff_type = ff_type.lower()

    ffmols = []
    for resname, smiles in molecules["smiles"].items():
        mol.templateResidueFromSmiles(f"resname {resname}", smiles, addHs=True)
        ffmol = Molecule.from_smiles(smiles)
        use_conformers = None
        if charge_model != "gasteiger":
            # TODO: Would be nice if we could use the conformer from the input structure
            ffmol.generate_conformers(n_conformers=1)
            use_conformers = ffmol.conformers[0]
        ffmol.assign_partial_charges(
            partial_charge_method=charge_model, use_conformers=use_conformers
        )
        ffmols.append(ffmol)

    if ff_type == "gaff":
        from openmmforcefields.generators import GAFFTemplateGenerator

        if ff_path not in GAFFTemplateGenerator.INSTALLED_FORCEFIELDS:
            raise ValueError(
                f"Invalid GAFF force field: {ff_path}. Available force fields: {GAFFTemplateGenerator.INSTALLED_FORCEFIELDS}"
            )

        template_gen = GAFFTemplateGenerator(molecules=ffmols, forcefield=ff_path)
    elif ff_type == "openff":
        from openmmforcefields.generators import SMIRNOFFTemplateGenerator

        if (
            not os.path.exists(ff_path)
            and ff_path not in SMIRNOFFTemplateGenerator.INSTALLED_FORCEFIELDS
        ):
            raise ValueError(
                f"Invalid OpenFF force field: {ff_path}. Available force fields: {SMIRNOFFTemplateGenerator.INSTALLED_FORCEFIELDS}"
            )

        template_gen = SMIRNOFFTemplateGenerator(molecules=ffmols, forcefield=ff_path)
    elif ff_type == "espaloma":
        from openmmforcefields.generators import EspalomaTemplateGenerator

        template_gen = EspalomaTemplateGenerator(molecules=ffmols, forcefield=ff_path)
    else:
        raise ValueError(
            f"Invalid small molecule force field: {ff_type}. Available force fields: gaff, openff, espaloma"
        )

    return template_gen


def openmm_to_molecule(topology, positions):
    # TODO: Deprecate in next moleculekit release
    from moleculekit.molecule import Molecule
    from openmm import unit
    import numpy as np

    positions = np.array(positions.value_in_unit(unit.angstrom))

    mol = Molecule()
    mol.empty(topology.getNumAtoms(), numFrames=1)

    for i, atom in enumerate(topology.atoms()):
        mol.name[i] = atom.name
        mol.element[i] = atom.element.symbol
        mol.resname[i] = atom.residue.name
        mol.resid[i] = int(atom.residue.id)
        mol.chain[i] = atom.residue.chain.id
        mol.insertion[i] = atom.residue.insertionCode.strip()
        mol.formalcharge[i] = atom.formalCharge if atom.formalCharge is not None else 0
        mol.coords[i, :, 0] = positions[i]

    bonds = []
    bondtype = []
    for bond in topology.bonds():
        bonds.append([bond.atom1.index, bond.atom2.index])
        bondtype.append(bond.order if bond.order is not None else "un")
    mol.bonds = np.array(bonds, dtype=np.uint32)
    mol.bondtype = np.array(bondtype, dtype=np.object_)

    return mol


def create_forcefield(mol, parameters, box_vectors, molecules=None, protonate=False):
    from openmm.app.pdbfile import PDBFile
    from openmm.app.forcefield import ForceField
    from openmm import app
    from moleculekit.molecule import Molecule
    import tempfile

    for prm in parameters:
        if not prm.lower().endswith(".xml"):
            raise ValueError(
                f"Only OpenMM XML force field files are supported for PDB files: {prm}"
            )
    forcefield = ForceField(*parameters)

    # If molecule SMILES are defined in the input file we need to create parameters for them
    if molecules:
        template_gen = _get_small_molecule_template_generator(mol, molecules)
        forcefield.registerTemplateGenerator(template_gen.generator)

    with tempfile.TemporaryDirectory() as tmpdir:
        mol.write(os.path.join(tmpdir, "mol.pdb"), writebonds=True)
        pdbfile = PDBFile(os.path.join(tmpdir, "mol.pdb"))
        topology = pdbfile.topology
        positions = pdbfile.positions
        topology.setPeriodicBoxVectors(box_vectors)

    # If the user has specified to protonate the system use modeller. `protonate`
    # accepts a bool (True → pH 7.4) or a float pH value.
    if protonate:
        if isinstance(protonate, bool):
            pH = 7.4
        else:
            pH = float(protonate)
        modeller = app.Modeller(pdbfile.topology, pdbfile.positions)
        modeller.addHydrogens(forcefield, pH=pH)
        topology = modeller.topology
        positions = modeller.positions

    # mol = Molecule.fromOpenMMTopology(topology, positions)
    mol = openmm_to_molecule(topology, positions)

    return forcefield, [topology], topology, mol
