# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import copy
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock

import pytest

from geneva import DEFAULT_UPLOAD_DIR, connect
from geneva.cluster.builder import default_image
from geneva.manifest import GenevaManifest
from geneva.manifest.builder import (
    CondaManifestBuilder,
    PipManifestBuilder,
    SiteManifestBuilder,
)


@pytest.mark.slow
def test_manifest_upload_location_namespace(tmp_path: Path) -> None:
    """Test manifest zips uploaded to correct location for namespace connections."""
    import pyarrow.fs as pafs
    from lance_namespace import DescribeTableRequest

    # Create a directory namespace connection with a system namespace
    system_ns = "test_system"
    db_root = tmp_path / "db"
    db_root.mkdir()

    # Create zip output directory
    zip_dir = tmp_path / "zips"
    zip_dir.mkdir()

    geneva = connect(
        namespace_impl="dir",
        namespace_properties={"root": str(db_root)},
        system_namespace=[system_ns],
    )

    manifest_def = GenevaManifest(
        name="test-manifest",
        local_zip_output_dir=str(zip_dir),
        skip_site_packages=True,  # Skip to make test faster
        delete_local_zips=False,
        pip=["numpy"],
    )

    # Define the manifest (this should upload zips)
    geneva.define_manifest("test-manifest", manifest_def)

    # Verify the manifest was created and has zips
    manifests = geneva.list_manifests()
    assert len(manifests) == 1
    manifest = manifests[0]
    assert len(manifest.zips) > 0, "Expected manifest to have uploaded zips"

    # Get the manifest table location from namespace
    assert geneva.namespace_client is not None
    from geneva.manifest.mgr import MANIFEST_TABLE_NAME

    table_id = [system_ns, MANIFEST_TABLE_NAME]
    response = geneva.namespace_client.describe_table(DescribeTableRequest(id=table_id))
    manifest_table_location = response.location
    assert manifest_table_location is not None

    # Expected upload directory should be {manifest_table_location}/_geneva_uploads
    expected_upload_dir = f"{manifest_table_location.rstrip('/')}/{DEFAULT_UPLOAD_DIR}"

    # Verify that all uploaded zips are in the expected location
    for zip_path_list in manifest.zips:
        for zip_path in zip_path_list:
            # The zip path should start with the expected upload directory
            assert zip_path.startswith(expected_upload_dir), (
                f"Expected zip to be uploaded to {expected_upload_dir}, "
                f"but got {zip_path}"
            )

            # Verify the file actually exists at that location
            filesystem, path = pafs.FileSystem.from_uri(zip_path)
            file_info = filesystem.get_file_info(path)
            assert file_info.type == pafs.FileType.File, (
                f"Expected file to exist at {zip_path}"
            )


@pytest.mark.slow
def test_manifest_upload_location_local(tmp_path: Path) -> None:
    """Test manifest zips uploaded to correct location for local connections."""
    import pyarrow.fs as pafs

    # Create a local connection
    db_path = tmp_path / "db"

    # Create zip output directory
    zip_dir = tmp_path / "zips"
    zip_dir.mkdir()

    geneva = connect(db_path)

    manifest_def = GenevaManifest(
        name="test-manifest",
        local_zip_output_dir=str(tmp_path / "zips"),
        skip_site_packages=True,  # Skip to make test faster
        delete_local_zips=False,
        pip=["numpy"],
    )

    # Define the manifest (this should upload zips)
    geneva.define_manifest("test-manifest", manifest_def)

    # Verify the manifest was created and has zips
    manifests = geneva.list_manifests()
    assert len(manifests) == 1
    manifest = manifests[0]
    assert len(manifest.zips) > 0, "Expected manifest to have uploaded zips"

    # Verify upload directory is within the manifest table
    from geneva.manifest.mgr import MANIFEST_TABLE_NAME

    expected_upload_dir = (
        f"{str(db_path)}/__system/{MANIFEST_TABLE_NAME}.lance/{DEFAULT_UPLOAD_DIR}"
    )

    # Verify that all uploaded zips are in the expected location
    for zip_path_list in manifest.zips:
        for zip_path in zip_path_list:
            # The zip path should start with the expected upload directory
            assert zip_path.startswith(expected_upload_dir), (
                f"Expected zip to be uploaded to {expected_upload_dir}, "
                f"but got {zip_path}"
            )

            # Verify the file actually exists at that location
            filesystem, path = pafs.FileSystem.from_uri(zip_path)
            file_info = filesystem.get_file_info(path)
            assert file_info.type == pafs.FileType.File, (
                f"Expected file to exist at {zip_path}"
            )


@pytest.mark.slow
def test_manifest_crud(tmp_path: Path) -> None:
    mock_uploader = MagicMock()
    mock_uploader.upload_dir = "/mock/upload/dir"
    mock_uploader._file_exists.return_value = False
    mock_uploader.upload.side_effect = lambda path: f"mock://{path.name}"

    geneva = connect(tmp_path)

    manifest_def = GenevaManifest(
        name="test-manifest-1",
        local_zip_output_dir=str(tmp_path),
        skip_site_packages=False,
        delete_local_zips=False,
        pip=["numpy", "pandas"],
        py_modules=["pyarrow"],
    )

    # upload and create
    geneva.define_manifest("test-manifest-1", manifest_def, uploader=mock_uploader)
    m = geneva.list_manifests()[0]
    _assert_manifest_eq(m.as_dict(), manifest_def.as_dict())

    upload_count = mock_uploader.upload.call_count
    assert upload_count >= 1, "files were not uploaded"

    # update - should update metadata and upload new artifacts
    manifest_def.skip_site_packages = True
    geneva.define_manifest("test-manifest-1", manifest_def, uploader=mock_uploader)
    manifests = geneva.list_manifests()
    assert len(manifests) == 1, "expected single manifest"
    m1 = manifests[0].as_dict()
    m2 = manifest_def.as_dict()
    assert m1["checksum"] != m2["checksum"], "checksum should change"
    _assert_manifest_eq(m1, m2)
    assert mock_uploader.upload.call_count >= upload_count, "files were not uploaded"

    # delete
    geneva.delete_manifest("test-manifest-1")
    assert geneva.list_manifests() == []


def _assert_manifest_eq(m1: dict, m2: dict) -> bool:
    m1 = copy.deepcopy(m1)
    m2 = copy.deepcopy(m2)
    # exclude transient fields from comparison
    for f in {"checksum", "zips"}:
        if f in m1:
            del m1[f]
        if f in m2:
            del m2[f]
    assert m1 == m2, "manifests should match"


def test_manifest_conda_passed_to_ray(monkeypatch: Any) -> None:
    """Test that conda dependencies from manifest are passed to Ray runtime_env."""
    import geneva.runners.ray._mgr as ray_mgr_mod

    ray_init_called: dict[str, Any] = {}

    def mock_ray_init(*args: Any, **kwargs: Any) -> None:
        ray_init_called["args"] = args
        ray_init_called["kwargs"] = kwargs

    monkeypatch.setattr("ray.init", mock_ray_init)
    monkeypatch.setattr("ray.shutdown", lambda: None)
    monkeypatch.setattr("ray.is_initialized", lambda: False)

    # Use a realistic conda structure with nested dependencies
    conda_deps = {
        "dependencies": [
            "python=3.8",
            "pip",
            {"pip": ["requests", "chess"]},
        ],
        "channels": ["conda-forge"],
    }

    manifest = CondaManifestBuilder.create("conda-manifest").conda(conda_deps).build()

    with ray_mgr_mod.ray_cluster(local=True, manifest=manifest):
        pass

    assert "kwargs" in ray_init_called
    runtime_env = ray_init_called["kwargs"]["runtime_env"]
    assert "conda" in runtime_env
    assert runtime_env["conda"] == conda_deps
    assert "pip" not in runtime_env


def test_manifest_requirements_path_passed_to_ray(
    monkeypatch: Any, tmp_path: Path
) -> None:
    """Test that requirements_path from manifest is passed to Ray runtime_env."""
    import geneva.runners.ray._mgr as ray_mgr_mod

    ray_init_called: dict[str, Any] = {}

    def mock_ray_init(*args: Any, **kwargs: Any) -> None:
        ray_init_called["args"] = args
        ray_init_called["kwargs"] = kwargs

    monkeypatch.setattr("ray.init", mock_ray_init)
    monkeypatch.setattr("ray.shutdown", lambda: None)
    monkeypatch.setattr("ray.is_initialized", lambda: False)

    req_file = tmp_path / "requirements.txt"
    req_file.write_text("numpy>=1.0\npandas\n")
    req_path = str(req_file)

    manifest = (
        PipManifestBuilder.create("req-manifest").requirements_path(req_path).build()
    )

    with ray_mgr_mod.ray_cluster(local=True, manifest=manifest):
        pass

    assert "kwargs" in ray_init_called
    runtime_env = ray_init_called["kwargs"]["runtime_env"]
    assert "pip" in runtime_env
    assert runtime_env["pip"] == req_path
    assert "conda" not in runtime_env


def test_manifest_conda_environment_path_passed_to_ray(
    monkeypatch: Any, tmp_path: Path
) -> None:
    """Test that conda_environment_path from manifest is passed to Ray runtime_env."""
    import geneva.runners.ray._mgr as ray_mgr_mod

    ray_init_called: dict[str, Any] = {}

    def mock_ray_init(*args: Any, **kwargs: Any) -> None:
        ray_init_called["args"] = args
        ray_init_called["kwargs"] = kwargs

    monkeypatch.setattr("ray.init", mock_ray_init)
    monkeypatch.setattr("ray.shutdown", lambda: None)
    monkeypatch.setattr("ray.is_initialized", lambda: False)

    env_file = tmp_path / "environment.yml"
    env_file.write_text(
        "name: test-env\n"
        "channels:\n"
        "  - conda-forge\n"
        "dependencies:\n"
        "  - python=3.10\n"
        "  - numpy\n"
    )
    env_path = str(env_file)

    manifest = (
        CondaManifestBuilder.create("conda-env-manifest")
        .conda_environment_path(env_path)
        .build()
    )

    with ray_mgr_mod.ray_cluster(local=True, manifest=manifest):
        pass

    assert "kwargs" in ray_init_called
    runtime_env = ray_init_called["kwargs"]["runtime_env"]
    assert "conda" in runtime_env
    assert runtime_env["conda"] == env_path
    assert "pip" not in runtime_env


def test_make_manifest_conda_none() -> None:
    """Test that _make_manifest handles conda=None from stored data."""
    from datetime import datetime, timezone

    from geneva.manifest.mgr import _make_manifest

    args = {
        "name": "test",
        "version": None,
        "pip": [],
        "py_modules": [],
        "head_image": None,
        "worker_image": None,
        "skip_site_packages": True,
        "delete_local_zips": False,
        "local_zip_output_dir": None,
        "zips": [[]],
        "checksum": None,
        "created_at": datetime.now(timezone.utc),
        "created_by": "test",
        "requirements_path": None,
        "conda": None,
        "conda_environment_path": None,
    }
    manifest = _make_manifest(args)
    assert manifest.conda == {}


# =============================================================================
# Type-Safe Builder Tests
# =============================================================================


class TestPipManifestBuilder:
    """Tests for PipManifestBuilder."""

    def test_no_conda_methods(self) -> None:
        """PipManifestBuilder should not have conda methods."""
        builder = PipManifestBuilder()
        assert not hasattr(builder, "conda")
        assert not hasattr(builder, "conda_environment_path")

    def test_has_pip_methods(self) -> None:
        """PipManifestBuilder should have pip methods."""
        builder = PipManifestBuilder()
        assert hasattr(builder, "pip")
        assert hasattr(builder, "add_pip")
        assert hasattr(builder, "requirements_path")

    def test_pip_packages(self) -> None:
        """Test setting pip packages."""
        manifest = PipManifestBuilder.create("test").pip(["numpy", "pandas"]).build()
        assert manifest.pip == ["numpy", "pandas"]
        assert manifest.conda == {}

    def test_add_pip(self) -> None:
        """Test adding pip packages incrementally."""
        manifest = (
            PipManifestBuilder.create("test").add_pip("numpy").add_pip("pandas").build()
        )
        assert manifest.pip == ["numpy", "pandas"]

    def test_requirements_path(self, tmp_path: Path) -> None:
        """Test using requirements.txt path."""
        req_file = tmp_path / "requirements.txt"
        req_file.write_text("numpy\npandas\n")

        manifest = (
            PipManifestBuilder.create("test").requirements_path(str(req_file)).build()
        )
        assert manifest.requirements_path == str(req_file)
        assert manifest.pip == []

    def test_pip_and_requirements_mutually_exclusive(self) -> None:
        """Test that pip and requirements_path cannot both be set."""
        with pytest.raises(ValueError, match="Cannot set both pip and requirements"):
            PipManifestBuilder.create("test").pip(["numpy"]).requirements_path(
                "/path"
            ).build()

    def test_common_methods(self) -> None:
        """Test that common methods work."""
        manifest = (
            PipManifestBuilder.create("test")
            .version("1.0")
            .py_modules(["mymodule"])
            .head_image("custom:latest")
            .worker_image("worker:latest")
            .upload_site_packages(True)
            .pip(["numpy"])
            .build()
        )
        assert manifest.version == "1.0"
        assert manifest.py_modules == ["mymodule"]
        assert manifest.head_image == "custom:latest"
        assert manifest.worker_image == "worker:latest"
        assert manifest.skip_site_packages is False

    def test_requires_name(self) -> None:
        """Test that builder requires a name."""
        with pytest.raises(ValueError, match="Manifest name is required"):
            PipManifestBuilder().build()


class TestCondaManifestBuilder:
    """Tests for CondaManifestBuilder."""

    def test_no_pip_methods(self) -> None:
        """CondaManifestBuilder should not have pip methods."""
        builder = CondaManifestBuilder()
        assert not hasattr(builder, "pip")
        assert not hasattr(builder, "add_pip")
        assert not hasattr(builder, "requirements_path")

    def test_has_conda_methods(self) -> None:
        """CondaManifestBuilder should have conda methods."""
        builder = CondaManifestBuilder()
        assert hasattr(builder, "conda")
        assert hasattr(builder, "conda_environment_path")

    def test_conda_dependencies(self) -> None:
        """Test setting conda dependencies."""
        conda_deps = {"dependencies": ["python=3.10", "numpy"]}
        manifest = CondaManifestBuilder.create("test").conda(conda_deps).build()
        assert manifest.conda == conda_deps
        assert manifest.pip == []

    def test_conda_environment_path(self, tmp_path: Path) -> None:
        """Test using environment.yml path."""
        env_file = tmp_path / "environment.yml"
        env_file.write_text("name: test\ndependencies:\n  - numpy\n")

        manifest = (
            CondaManifestBuilder.create("test")
            .conda_environment_path(str(env_file))
            .build()
        )
        assert manifest.conda_environment_path == str(env_file)
        assert manifest.conda == {}

    def test_conda_and_path_mutually_exclusive(self) -> None:
        """Test that conda and conda_environment_path cannot both be set."""
        with pytest.raises(ValueError, match="Cannot set both conda and conda_env"):
            CondaManifestBuilder.create("test").conda(
                {"dependencies": ["numpy"]}
            ).conda_environment_path("/path").build()

    def test_common_methods(self) -> None:
        """Test that common methods work."""
        manifest = (
            CondaManifestBuilder.create("test")
            .version("2.0")
            .add_py_module("mymodule")
            .default_head_image()
            .upload_site_packages()
            .conda({"dependencies": ["numpy"]})
            .build()
        )
        assert manifest.version == "2.0"
        assert manifest.py_modules == ["mymodule"]
        assert manifest.head_image == default_image()
        assert manifest.skip_site_packages is False

    def test_requires_name(self) -> None:
        """Test that builder requires a name."""
        with pytest.raises(ValueError, match="Manifest name is required"):
            CondaManifestBuilder().build()


class TestSiteManifestBuilder:
    """Tests for SiteManifestBuilder."""

    def test_no_pip_methods(self) -> None:
        """SiteManifestBuilder should not have pip methods."""
        builder = SiteManifestBuilder()
        assert not hasattr(builder, "pip")
        assert not hasattr(builder, "add_pip")
        assert not hasattr(builder, "requirements_path")

    def test_no_conda_methods(self) -> None:
        """SiteManifestBuilder should not have conda methods."""
        builder = SiteManifestBuilder()
        assert not hasattr(builder, "conda")
        assert not hasattr(builder, "conda_environment_path")

    def test_upload_site_packages_default_true(self) -> None:
        """SiteManifestBuilder should default to uploading site packages."""
        manifest = SiteManifestBuilder.create("test").build()
        assert manifest.skip_site_packages is False  # upload=True means skip=False

    def test_can_disable_site_packages(self) -> None:
        """Test that site packages can be disabled."""
        manifest = (
            SiteManifestBuilder.create("test").upload_site_packages(False).build()
        )
        assert manifest.skip_site_packages is True

    def test_no_external_deps(self) -> None:
        """Test that no external dependencies are set."""
        manifest = SiteManifestBuilder.create("test").build()
        assert manifest.pip == []
        assert manifest.conda == {}
        assert manifest.requirements_path is None
        assert manifest.conda_environment_path is None

    def test_common_methods(self) -> None:
        """Test that common methods work."""
        manifest = (
            SiteManifestBuilder.create("test")
            .version("3.0")
            .py_modules(["local_module"])
            .head_image("site:latest")
            .delete_local_zips(True)
            .build()
        )
        assert manifest.version == "3.0"
        assert manifest.py_modules == ["local_module"]
        assert manifest.head_image == "site:latest"
        assert manifest.delete_local_zips is True

    def test_requires_name(self) -> None:
        """Test that builder requires a name."""
        with pytest.raises(ValueError, match="Manifest name is required"):
            SiteManifestBuilder().build()


class TestGenevaManifestFactories:
    """Tests for GenevaManifest static factory methods."""

    def test_create_pip_factory(self) -> None:
        """GenevaManifest.create_pip() should return PipManifestBuilder."""
        builder = GenevaManifest.create_pip("test")
        assert isinstance(builder, PipManifestBuilder)
        manifest = builder.pip(["numpy"]).build()
        assert manifest.name == "test"
        assert manifest.pip == ["numpy"]

    def test_create_conda_factory(self) -> None:
        """GenevaManifest.create_conda() should return CondaManifestBuilder."""
        builder = GenevaManifest.create_conda("test")
        assert isinstance(builder, CondaManifestBuilder)
        manifest = builder.conda({"dependencies": ["numpy"]}).build()
        assert manifest.name == "test"
        assert manifest.conda == {"dependencies": ["numpy"]}

    def test_create_site_factory(self) -> None:
        """GenevaManifest.create_site() should return SiteManifestBuilder."""
        builder = GenevaManifest.create_site("test")
        assert isinstance(builder, SiteManifestBuilder)
        manifest = builder.build()
        assert manifest.name == "test"
        assert manifest.skip_site_packages is False  # upload enabled by default

    def test_full_workflow_pip(self) -> None:
        """Test complete workflow using pip factory."""
        manifest = (
            GenevaManifest.create_pip("my-pip-manifest")
            .version("1.0.0")
            .pip(["numpy", "pandas"])
            .add_py_module("mymodule")
            .head_image("custom:latest")
            .build()
        )

        assert manifest.name == "my-pip-manifest"
        assert manifest.version == "1.0.0"
        assert manifest.pip == ["numpy", "pandas"]
        assert manifest.py_modules == ["mymodule"]
        assert manifest.head_image == "custom:latest"
        assert manifest.conda == {}

    def test_full_workflow_conda(self) -> None:
        """Test complete workflow using conda factory."""
        conda_deps = {
            "channels": ["conda-forge"],
            "dependencies": ["python=3.10", "numpy", {"pip": ["requests"]}],
        }
        manifest = (
            GenevaManifest.create_conda("my-conda-manifest")
            .conda(conda_deps)
            .upload_site_packages(True)
            .build()
        )

        assert manifest.name == "my-conda-manifest"
        assert manifest.conda == conda_deps
        assert manifest.pip == []
        assert manifest.skip_site_packages is False
