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 _register_extra_xml_template_matcher(forcefield, extra_xml):
    """Prefer the template whose NAME matches the topology residue's name,
    whenever such a same-named template exists and actually matches the
    residue's atoms.

    Ported from ``htmd.builder.openmm``. OpenMM's default matcher selects
    templates by atom-element + bond-graph signature, which returns multiple
    matches for topologically-indistinguishable templates (e.g. a renamed CYS
    scaffold ``XX1`` (S-C external bond) vs ff14SB ``NCYX`` (S-S disulfide):
    identical atom set, different bonded partner -> "Multiple non-identical
    matching templates"). Resolving by NAME is correct because the builder
    emitted the residue with a specific name on purpose. The same-named
    template is only returned when it actually matches the residue's atoms;
    otherwise we fall through to OpenMM's default logic.
    """
    import xml.etree.ElementTree as ET
    from openmm.app.internal import compiled

    if isinstance(extra_xml, str):
        extra_xml = [extra_xml]

    # Only arm the matcher when extra_xml defines residue templates; without
    # them the default OpenMM logic already produces the right answer.
    extra_resnames = set()
    for path in extra_xml:
        try:
            root = ET.parse(path).getroot()
        except (ET.ParseError, FileNotFoundError):
            continue
        for res in root.findall("./Residues/Residue"):
            name = res.attrib.get("name")
            if name:
                extra_resnames.add(name)

    if not extra_resnames:
        return

    def _matcher(forcefield, residue, bonded_to_atom, ignore_external, ignore_extra):
        template = forcefield._templates.get(residue.name)
        if template is None:
            return None
        # Verify the same-named template actually matches before returning it -
        # registerTemplateMatcher raises if a returned template doesn't match.
        match = compiled.matchResidueToTemplate(
            residue, template, bonded_to_atom, ignore_external, ignore_extra,
        )
        if match is None:
            return None
        return template

    forcefield.registerTemplateMatcher(_matcher)


def _mol_positions(mol):
    """Atom positions (frame 0) as an OpenMM Quantity of Vec3 in angstrom.

    OpenMM's ``Modeller`` / ``createDisulfideBonds`` need Vec3 objects (which
    support vector arithmetic); a list of plain ``[x, y, z]`` lists does not.
    """
    from openmm import Vec3, unit

    return unit.Quantity(
        [Vec3(float(x), float(y), float(z)) for x, y, z in mol.coords[:, :, 0]],
        unit.angstrom,
    )


def _molecule_to_openmm_topology(mol):
    """Build an OpenMM Topology directly from a moleculekit Molecule's
    explicit atoms/residues/bonds.

    Avoids the PDB round-trip (and its 99,999-atom CONECT ceiling) so that
    inter-residue / external bonds transfer verbatim.
    """
    from openmm.app import Topology
    from openmm.app import element as elem

    top = Topology()
    chains, residues, atoms = {}, {}, []
    for i in range(mol.numAtoms):
        ckey = (mol.chain[i], mol.segid[i])
        if ckey not in chains:
            chains[ckey] = top.addChain(id=(str(mol.chain[i]) or None))
        rkey = (mol.chain[i], mol.segid[i], int(mol.resid[i]), mol.insertion[i])
        if rkey not in residues:
            residues[rkey] = top.addResidue(
                mol.resname[i], chains[ckey], id=str(mol.resid[i])
            )
        try:
            e = elem.get_by_symbol(mol.element[i]) if mol.element[i] else None
        except KeyError:
            e = None
        atoms.append(top.addAtom(mol.name[i], e, residues[rkey]))

    # Standard-residue + disulfide connectivity from OpenMM's residue database
    # (this is what PDBFile does internally, so a plain PDB without CONECT still
    # matches templates), then layer on any explicit bonds the Molecule carries
    # (ligands, non-standard residues, external bonds, mmCIF connectivity) that
    # are not already present.
    top.createStandardBonds()
    top.createDisulfideBonds(_mol_positions(mol))
    # Skip metal-coordination bonds (moleculekit bondtype "mc", e.g. trypsin's
    # Ca2+ to GLU/ASP carboxylate O or to waters in a crystal PDB). These are
    # dative, not covalent: MM force fields treat ions as non-bonded point
    # charges. Adding them would give a standard residue a spurious external
    # bond and break its template match ("...has 1 O atom too many").
    bondtypes = (
        mol.bondtype
        if len(mol.bondtype) == len(mol.bonds)
        else [None] * len(mol.bonds)
    )
    existing = {frozenset((b.atom1.index, b.atom2.index)) for b in top.bonds()}
    for (a1, a2), btype in zip(mol.bonds, bondtypes):
        if btype == "mc":
            continue
        a1, a2 = int(a1), int(a2)
        key = frozenset((a1, a2))
        if key not in existing:
            top.addBond(atoms[a1], atoms[a2])
            existing.add(key)
    return top, atoms


def create_forcefield(mol, parameters, box_vectors, molecules=None, protonate=False):
    from openmm.app.forcefield import ForceField
    from openmm import app

    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)

    # Local ffxml fragments (e.g. emitted by openmm.build for non-standard
    # residues) may define names that clash with stock templates; arm the
    # name-priority matcher so they resolve. Stock-by-name entries are not
    # files, so filtering to existing paths is correct here.
    local_xml = [p for p in parameters if os.path.isfile(p)]
    if local_xml:
        _register_extra_xml_template_matcher(forcefield, local_xml)

    # 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)

    # Build the OpenMM Topology directly from the moleculekit Molecule's
    # explicit bonds (no PDB round-trip, no 99,999-atom CONECT ceiling).
    topology, _ = _molecule_to_openmm_topology(mol)
    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:
        pH = 7.4 if isinstance(protonate, bool) else float(protonate)
        positions = _mol_positions(mol)
        modeller = app.Modeller(topology, positions)
        modeller.addHydrogens(forcefield, pH=pH)
        topology = modeller.topology
        # Hydrogens changed the atom set; rebuild mol from the new topology.
        mol = openmm_to_molecule(topology, modeller.positions)

    return forcefield, [topology], topology, mol
