# -*- coding: utf-8 -*-
"""
Mocked unit tests for AmorphousBuilder.build().

Replaces the heavy scientific dependencies (OpenFF / Packmol /
Interchange / RDKit / AmberTools) with lightweight stubs so the build
pipeline can be exercised in CI without installing those stacks.

What is covered:
- The 6-stage build() method wires the stages in the right order.
- The returned dict has every documented key, including the new
  ``wrap_script`` key added in 1.15.2.
- ``input/config.json`` is written before any stage starts.
- ``md/run_all.sh`` and ``md/wrap_pbc.sh`` are generated alongside the
  MDP files.
- counts / box_nm propagate from the stub stages.
"""
from __future__ import annotations

import json
import os
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import pytest

from abmptools.amorphous import BuildConfig, ComponentSpec
from abmptools.amorphous.builder import AmorphousBuilder


class _FakeMol:
    """Minimal stand-in for an openff.toolkit.Molecule."""

    def __init__(self, name: str, n_atoms: int = 5, mw: float = 16.0):
        self.name = name
        self.n_atoms = n_atoms
        self._mw = mw


@pytest.fixture
def _stubbed_builder(tmp_path, monkeypatch):
    """Build a ready-to-use AmorphousBuilder with every external call stubbed."""
    cfg = BuildConfig(
        components=[
            ComponentSpec(smiles="C", name="methane", n_mol=10),
        ],
        density_g_cm3=0.4,
        temperature=300.0,
        seed=42,
        output_dir=str(tmp_path / "run"),
    )

    # prepare_molecule -> returns a FakeMol; no SMILES -> 3D work happens.
    def _prepare_molecule(smiles="", sdf_path="", name="MOL"):
        return _FakeMol(name=name, n_atoms=5, mw=16.0)

    def _get_molecular_weight(mol):
        return mol._mw

    def _write_single_mol_pdb(mol, output_path):
        # Write an empty placeholder so downstream code paths see a file.
        Path(output_path).write_text(f"REMARK FakePDB for {mol.name}\nEND\n")
        return str(Path(output_path).resolve())

    # packmol stub: create an empty mixture.pdb where expected.
    def _run_packmol(pdb_paths, counts, box_size_nm, output_pdb,
                     build_dir, tolerance, seed, packmol_path):
        Path(output_pdb).write_text("REMARK fake mixture\nEND\n")
        return str(Path(output_pdb).resolve())

    # parameterizer stubs
    def _create_interchange(molecules, counts, box_size_nm,
                            mixture_pdb, forcefield_name):
        return SimpleNamespace(_fake="interchange")

    def _export_gromacs(interchange, gro_path, top_path):
        Path(gro_path).write_text("Generated by stub\n0\n 1 1 1\n")
        Path(top_path).write_text("; stub top\n[ system ]\nFake\n")
        return {"gro": str(Path(gro_path).resolve()),
                "top": str(Path(top_path).resolve())}

    # ndx_writer stub
    def _write_ndx(component_names, atom_counts, counts, ndx_path):
        lines = ["[ System ]"] + [f"[ {n} ]" for n in component_names]
        Path(ndx_path).write_text("\n".join(lines) + "\n")
        return str(Path(ndx_path).resolve())

    monkeypatch.setattr(
        "abmptools.amorphous.builder.prepare_molecule", _prepare_molecule)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.get_molecular_weight", _get_molecular_weight)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.write_single_mol_pdb", _write_single_mol_pdb)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.run_packmol", _run_packmol)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.create_interchange", _create_interchange)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.export_gromacs", _export_gromacs)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.write_ndx", _write_ndx)

    return cfg, AmorphousBuilder(cfg)


# ---------------------------------------------------------------- happy path


def test_build_produces_expected_result_keys(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()

    expected = {
        "output_dir", "gro", "top", "ndx", "mdp_files",
        "run_script", "wrap_script", "config_json",
        "box_nm", "counts",
    }
    assert expected.issubset(result.keys())


def test_build_writes_config_json_with_fields(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()
    cfg_path = result["config_json"]
    assert os.path.isfile(cfg_path)
    data = json.loads(Path(cfg_path).read_text())
    # The saved config reflects the BuildConfig we passed in.
    assert data["density_g_cm3"] == 0.4
    assert data["temperature"] == 300.0
    assert data["seed"] == 42
    assert data["components"][0]["name"] == "methane"
    assert data["components"][0]["n_mol"] == 10


def test_build_generates_run_and_wrap_scripts(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()
    assert Path(result["run_script"]).name == "run_all.sh"
    assert Path(result["wrap_script"]).name == "wrap_pbc.sh"
    assert Path(result["run_script"]).is_file()
    assert Path(result["wrap_script"]).is_file()
    # And they should be executable.
    assert os.access(result["run_script"], os.X_OK)
    assert os.access(result["wrap_script"], os.X_OK)


def test_build_generates_all_five_mdp_files(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()
    names = [Path(p).name for p in result["mdp_files"]]
    assert names == [
        "01_em.mdp",
        "02_nvt_highT.mdp",
        "03_npt_highT.mdp",
        "04_anneal.mdp",
        "05_npt_final.mdp",
    ]
    for p in result["mdp_files"]:
        assert Path(p).is_file()


def test_build_box_and_counts_propagate(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()
    # The component declared n_mol=10 and no explicit box -> density-driven.
    assert result["counts"] == [10]
    assert result["box_nm"] > 0.0


def test_build_writes_single_component_ndx(_stubbed_builder):
    cfg, builder = _stubbed_builder
    result = builder.build()
    ndx_text = Path(result["ndx"]).read_text()
    assert "[ System ]" in ndx_text
    assert "[ methane ]" in ndx_text


# ---------------------------------------------------------------- ordering


def test_build_invokes_stages_in_order(_stubbed_builder, monkeypatch):
    cfg, builder = _stubbed_builder

    calls: list = []

    orig_pack = builder._run_packing
    orig_param = builder._parameterize
    orig_ndx = builder._write_ndx

    def _pack(*a, **kw):
        calls.append("packing")
        return orig_pack(*a, **kw)

    def _param(*a, **kw):
        calls.append("parameterize")
        return orig_param(*a, **kw)

    def _write(*a, **kw):
        calls.append("ndx")
        return orig_ndx(*a, **kw)

    monkeypatch.setattr(builder, "_run_packing", _pack)
    monkeypatch.setattr(builder, "_parameterize", _param)
    monkeypatch.setattr(builder, "_write_ndx", _write)

    builder.build()
    assert calls == ["packing", "parameterize", "ndx"]


# ---------------------------------------------------------------- no-input


def test_build_requires_n_mol_or_weight_fraction(tmp_path, monkeypatch):
    cfg = BuildConfig(
        components=[
            ComponentSpec(smiles="C", name="methane"),  # no n_mol, no wf
        ],
        output_dir=str(tmp_path / "bad"),
    )

    def _prep(smiles="", sdf_path="", name="MOL"):
        return _FakeMol(name=name)

    monkeypatch.setattr(
        "abmptools.amorphous.builder.prepare_molecule", _prep)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.get_molecular_weight", lambda m: 16.0)
    monkeypatch.setattr(
        "abmptools.amorphous.builder.write_single_mol_pdb",
        lambda m, p: Path(p).write_text("fake\n") or str(p))

    builder = AmorphousBuilder(cfg)
    with pytest.raises(ValueError):
        builder.build()
