from __future__ import annotations

import zipfile
from pathlib import Path

import pytest

from cogames.cli.policy import PolicySpec
from cogames.cli.submit import create_submission_zip, validate_paths
from mettagrid.util.uri_resolvers.schemes import localize_uri


def test_validate_paths_accepts_absolute_within_cwd(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.chdir(tmp_path)
    (tmp_path / "weights.safetensors").write_text("ok")

    rel = validate_paths(["weights.safetensors"])
    assert rel == [Path("weights.safetensors")]

    abs_path = str((tmp_path / "weights.safetensors").resolve())
    rel2 = validate_paths([abs_path])
    assert rel2 == [Path("weights.safetensors")]


def test_validate_paths_rejects_outside_cwd(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.chdir(tmp_path)
    outside = tmp_path.parent / "outside.txt"

    with pytest.raises(ValueError):
        validate_paths([str(outside)])


def test_bundle_uri_directory_is_localized(tmp_path: Path) -> None:
    bundle_dir = tmp_path / "bundle"
    bundle_dir.mkdir()
    (bundle_dir / "policy_spec.json").write_text(
        '{"class_path": "foo.Bar", "data_path": "weights.safetensors", "init_kwargs": {}}'
    )
    (bundle_dir / "weights.safetensors").write_text("ok")

    local = localize_uri(bundle_dir.as_uri())
    assert local is not None
    assert local == bundle_dir
    assert local.is_dir()


def test_bundle_uri_zip_is_localized(tmp_path: Path) -> None:
    bundle_zip = tmp_path / "bundle.zip"
    with zipfile.ZipFile(bundle_zip, "w") as zf:
        zf.writestr("policy_spec.json", '{"class_path": "foo.Bar", "data_path": null, "init_kwargs": {}}')

    local = localize_uri(bundle_zip.as_uri())
    assert local == bundle_zip


def test_create_submission_zip_includes_ancestor_inits(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.chdir(tmp_path)

    # src/pkg/__init__.py
    # src/pkg/sub/__init__.py
    # src/pkg/sub/leaf/__init__.py
    # src/pkg/sub/leaf/policy.py
    leaf = tmp_path / "src" / "pkg" / "sub" / "leaf"
    leaf.mkdir(parents=True)
    (tmp_path / "src" / "pkg" / "__init__.py").write_text("")
    (tmp_path / "src" / "pkg" / "sub" / "__init__.py").write_text("")
    (leaf / "__init__.py").write_text("")
    (leaf / "policy.py").write_text("class P: pass")

    policy_spec = PolicySpec(class_path="pkg.sub.leaf.policy.P")
    include = [Path("src/pkg/sub/leaf")]

    zip_path = create_submission_zip(include, policy_spec)
    try:
        with zipfile.ZipFile(zip_path) as zf:
            names = set(zf.namelist())
        assert "src/pkg/__init__.py" in names
        assert "src/pkg/sub/__init__.py" in names
        assert "src/pkg/sub/leaf/__init__.py" in names
        assert "src/pkg/sub/leaf/policy.py" in names
    finally:
        zip_path.unlink()
