"""
Unit and regression test for the openff_nagl_models package.
"""

import importlib.resources
import os
from importlib.metadata import entry_points

import pytest

from openff.nagl_models import list_available_nagl_models, validate_nagl_model_path
from openff.nagl_models.openff_nagl_models import get_models_by_type


def find_model_files():
    base = importlib.resources.files("openff.nagl_models")
    filenames = sorted(base.glob("models/*/*.pt"))
    assert len(filenames) > 0
    return filenames


@pytest.mark.parametrize("model_name", find_model_files())
def test_validate_nagl_model_path(model_name):
    """Test that we can find files."""
    model_path = validate_nagl_model_path(model_name)
    assert os.path.exists(model_path)


def test_validate_nagl_model_path_failed():
    """Test that we cannot find false files."""
    with pytest.raises(FileNotFoundError):
        validate_nagl_model_path("does-not-exist.pt")


def test_local_validation(tmpdir):
    """Test that we can find local files."""
    with tmpdir.as_cwd():
        with pytest.raises(FileNotFoundError):
            validate_nagl_model_path("test.pt")

        with open("test.pt", "w") as f:
            f.write("test")
        model_path = validate_nagl_model_path("test.pt")
        assert os.path.exists(model_path)


def test_list_models():
    """Test that we can list models."""
    model_names = find_model_files()

    # also returns cached files, if they exist
    listed_model_names = list_available_nagl_models()

    # just make sure all files from entry points are in the list
    for model_name in model_names:
        assert model_name in listed_model_names


def test_entry_points():
    for entry_point in entry_points().select(
        group="openforcefield.nagl_model_directory",
    ):
        paths = entry_point.load()()
        for path in paths:
            assert os.path.exists(path)


def test_get_models_by_type():
    all_model_stems = [path.stem for path in get_models_by_type(model_type="am1bcc")]
    expected_stems = [
        "openff-gnn-am1bcc-0.0.1-alpha.1",
        "openff-gnn-am1bcc-0.1.0-rc.1",
        "openff-gnn-am1bcc-0.1.0-rc.2",
        "openff-gnn-am1bcc-0.1.0-rc.3",
        "openff-gnn-am1bcc-1.0.0",
    ]

    assert all_model_stems == expected_stems


def test_get_models_by_type_does_not_exist():
    err = "Model type does-not-exist not found in openff-nagl-models."
    with pytest.raises(ValueError, match=err):
        get_models_by_type(model_type="does-not-exist")


def test_test_get_models_by_type_production():
    latest_models = get_models_by_type(model_type="am1bcc", production_only=True)
    assert len(latest_models) == 1
    assert latest_models[0].stem == "openff-gnn-am1bcc-1.0.0"
