from __future__ import annotations

import shutil
from importlib.resources import as_file, files
from pathlib import Path

import numpy as np
import onnxruntime as ort
import pytest

from mouse_control import load_model, predict


def _bundled_path() -> Path:
    """Resolve the bundled mouse.onnx to a stable filesystem path for tests."""
    with as_file(files("mouse_control.assets").joinpath("mouse.onnx")) as p:
        return Path(p)


def test_load_model_returns_inference_session(bundled_session) -> None:
    assert isinstance(bundled_session, ort.InferenceSession)


def test_bundled_input_output_names(bundled_session) -> None:
    assert bundled_session.get_inputs()[0].name == "input"
    assert bundled_session.get_outputs()[0].name == "output"


def test_load_model_explicit_path() -> None:
    session = load_model(_bundled_path())
    assert isinstance(session, ort.InferenceSession)


def test_load_model_missing_path_raises() -> None:
    with pytest.raises(Exception):
        load_model("/no/such/path.onnx")


def test_load_model_cwd_override(tmp_path, monkeypatch) -> None:
    """When ./mouse.onnx exists, it should take precedence over the bundle."""
    monkeypatch.chdir(tmp_path)
    shutil.copy(_bundled_path(), tmp_path / "mouse.onnx")
    session = load_model()  # implicit
    assert isinstance(session, ort.InferenceSession)
    # Must produce the same output as the bundle since the file is identical.
    pkg_session = load_model(_bundled_path())
    a = predict(session, 75, -40)
    b = predict(pkg_session, 75, -40)
    np.testing.assert_array_equal(a, b)


def test_load_model_falls_back_to_bundled(tmp_path, monkeypatch) -> None:
    """With no ./mouse.onnx, load_model() must succeed via the bundle."""
    monkeypatch.chdir(tmp_path)
    assert not (tmp_path / "mouse.onnx").exists()
    session = load_model()
    assert isinstance(session, ort.InferenceSession)


def test_predict_shape_and_dtype(bundled_session) -> None:
    out = predict(bundled_session, 100.0, 200.0)
    assert out.shape == (10, 2)
    assert out.dtype == np.float32


def test_predict_is_deterministic(bundled_session) -> None:
    a = predict(bundled_session, 100, 100)
    b = predict(bundled_session, 100, 100)
    np.testing.assert_array_equal(a, b)


def test_predict_int_and_float_args_equivalent(bundled_session) -> None:
    a = predict(bundled_session, 50, 50)
    b = predict(bundled_session, 50.0, 50.0)
    np.testing.assert_array_equal(a, b)


@pytest.mark.parametrize(
    "dx,dy",
    [
        (150, 0), (-150, 0), (0, 150), (0, -150),
        (100, 100), (-100, 100), (100, -100), (-100, -100),
    ],
)
def test_predict_endpoint_near_target(bundled_session, dx: int, dy: int) -> None:
    """Trained model's endpoint should fall within ~50px of the target."""
    end = predict(bundled_session, dx, dy)[-1]
    err = float(np.linalg.norm(end - np.array([dx, dy])))
    assert err < 50, f"target ({dx},{dy}) endpoint err={err:.1f}"


def test_predict_zero_displacement_stays_near_origin(bundled_session) -> None:
    end = predict(bundled_session, 0, 0)[-1]
    assert abs(end[0]) < 20
    assert abs(end[1]) < 20