import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

try:
    import torch  # type: ignore
except Exception:  # pragma: no cover - fallback for test environment
    import types

    torch_module = types.ModuleType("torch")
    utils_module = types.ModuleType("torch.utils")
    data_module = types.ModuleType("torch.utils.data")
    setattr(data_module, "DataLoader", object)
    setattr(data_module, "TensorDataset", object)
    setattr(utils_module, "data", data_module)
    setattr(torch_module, "utils", utils_module)
    setattr(torch_module, "device", lambda *_args, **_kwargs: "cpu")
    sys.modules["torch"] = torch_module
    sys.modules["torch.utils"] = utils_module
    sys.modules["torch.utils.data"] = data_module
else:  # pragma: no cover
    import types
    try:
        import torch.utils.data as torch_utils_data  # type: ignore
    except Exception:
        torch_utils_data = types.SimpleNamespace()
        sys.modules["torch.utils.data"] = torch_utils_data  # type: ignore[assignment]

    if not hasattr(torch, "device"):
        torch.device = lambda *_args, **_kwargs: "cpu"  # type: ignore[attr-defined]
    if not hasattr(torch, "utils"):
        torch.utils = types.SimpleNamespace()  # type: ignore[attr-defined]
    if not hasattr(torch.utils, "data"):
        torch.utils.data = torch_utils_data  # type: ignore[attr-defined]
    if not hasattr(torch_utils_data, "DataLoader"):
        torch_utils_data.DataLoader = object  # type: ignore[attr-defined]
    if not hasattr(torch_utils_data, "TensorDataset"):
        torch_utils_data.TensorDataset = object  # type: ignore[attr-defined]


import importlib

device_mod = importlib.import_module("rebel_forge.device")
from rebel_forge.device import NebiusProvisioner, RemoteError


GPU_PLATFORM_RESPONSE = {
    "spec": {
        "presets": [
            {"name": "1gpu-16vcpu-200gb", "resources": {"gpu_count": 1}},
            {"name": "8gpu-128vcpu-1600gb", "resources": {"gpu_count": 8}},
        ]
    }
}


class AutoPresetTest(unittest.TestCase):
    def test_auto_preset_selects_gpu_count(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_home, patch.dict(
            os.environ, {"HOME": tmp_home}, clear=False
        ):
            provisioner = NebiusProvisioner(
                project_id="project-test",
                service_account_id="sa-test",
                public_key_id="pub-test",
                private_key_pem="-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----",
                profile="test-profile",
                endpoint="api.nebius.cloud:443",
            )
            with patch.object(
                provisioner,
                "_run_cli",
                return_value=json.dumps(GPU_PLATFORM_RESPONSE),
            ):
                preset = provisioner.auto_preset(platform="gpu-h200-sxm", gpu_count=8)
        self.assertEqual(preset, "8gpu-128vcpu-1600gb")

    def test_auto_preset_rejects_unknown_count(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_home, patch.dict(
            os.environ, {"HOME": tmp_home}, clear=False
        ):
            provisioner = NebiusProvisioner(
                project_id="project-test",
                service_account_id="sa-test",
                public_key_id="pub-test",
                private_key_pem="-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----",
                profile="test-profile",
                endpoint="api.nebius.cloud:443",
            )
            with patch.object(
                provisioner,
                "_run_cli",
                return_value=json.dumps(GPU_PLATFORM_RESPONSE),
            ):
                with self.assertRaises(RemoteError):
                    provisioner.auto_preset(platform="gpu-h200-sxm", gpu_count=4)

    def test_auto_preset_cpu_defaults_without_cli(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_home, patch.dict(
            os.environ, {"HOME": tmp_home}, clear=False
        ):
            provisioner = NebiusProvisioner(
                project_id="project-test",
                service_account_id="sa-test",
                public_key_id="pub-test",
                private_key_pem="-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----",
                profile="test-profile",
                endpoint="api.nebius.cloud:443",
            )
            with patch.object(provisioner, "_run_cli", MagicMock()) as mock_cli:
                preset = provisioner.auto_preset(platform="cpu-d3", gpu_count=None)
        self.assertEqual(preset, "8vcpu-32gb")
        mock_cli.assert_not_called()

    def test_resolve_credentials_falls_back_on_remote_error(self) -> None:
        device_mod._CACHED_BUNDLE = None
        env_overrides = {
            "NEBIUS_PROJECT_ID": "project-test",
            "NEBIUS_SERVICE_ACCOUNT_ID": "sa-test",
            "NEBIUS_AUTHORIZED_KEY_ID": "key-test",
            "NEBIUS_AUTHORIZED_PRIVATE_KEY": "-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----",
            "ssh_key_public": "ssh-rsa dummy rebel@test",
            "ssh_key_private": "-----BEGIN OPENSSH PRIVATE KEY-----\n...\n-----END OPENSSH PRIVATE KEY-----",
        }
        original_load_env = device_mod.load_local_env
        original_session = device_mod._load_session_payload
        original_fetch = device_mod._fetch_remote_bundle
        mock_load_env = MagicMock()
        mock_session = MagicMock(return_value={"token": "tok", "userId": "user"})
        mock_fetch = MagicMock(side_effect=RemoteError("remote error"))
        device_mod.load_local_env = mock_load_env  # type: ignore[assignment]
        device_mod._load_session_payload = mock_session  # type: ignore[assignment]
        device_mod._fetch_remote_bundle = mock_fetch  # type: ignore[assignment]
        env_with_flag = dict(env_overrides)
        env_with_flag["REBEL_FORGE_ALLOW_LEGACY_CREDENTIALS"] = "1"
        try:
            with patch.dict(os.environ, env_with_flag, clear=False):
                bundle = device_mod._resolve_credentials()
            self.assertEqual(bundle["projectId"], "project-test")
            self.assertEqual(bundle["serviceAccountId"], "sa-test")
            mock_load_env.assert_called_once()
            mock_session.assert_called_once()
            mock_fetch.assert_called_once()
        finally:
            device_mod.load_local_env = original_load_env  # type: ignore[assignment]
            device_mod._load_session_payload = original_session  # type: ignore[assignment]
            device_mod._fetch_remote_bundle = original_fetch  # type: ignore[assignment]
            device_mod._CACHED_BUNDLE = None

    def test_normalize_portal_base_adds_scheme(self) -> None:
        self.assertEqual(device_mod._normalize_portal_base("portal.example.com"), "https://portal.example.com")
        self.assertEqual(
            device_mod._normalize_portal_base("https://portal.example.com/"), "https://portal.example.com"
        )
        self.assertIsNone(device_mod._normalize_portal_base(""))

    def test_resolve_portal_base_updates_env(self) -> None:
        session = {"token": "tok", "userId": "user", "login_url": "portal.example.com"}
        env_snapshot = dict(os.environ)
        try:
            with patch.dict(os.environ, {}, clear=True):
                base, updated = device_mod._resolve_portal_base(session)
                self.assertEqual(base, "https://portal.example.com")
                self.assertEqual(updated["login_url"], "https://portal.example.com")
                self.assertEqual(os.environ.get("REBEL_FORGE_PORTAL_URL"), "https://portal.example.com")
        finally:
            os.environ.clear()
            os.environ.update(env_snapshot)

    def test_project_fallback_via_handshake_on_server_error(self) -> None:
        session = {"token": "tok", "userId": "user", "login_url": "https://portal.example.com"}
        env_snapshot = dict(os.environ)
        try:
            with patch.dict(os.environ, {}, clear=True):
                with patch.object(device_mod, "_fetch_default_project_metadata", return_value=None), patch.object(
                    device_mod, "_fetch_handshake_bundle", return_value={
                        "projectId": "project-handshake",
                        "endpoint": "api.nebius.cloud:443",
                        "imageId": "image-abc",
                        "subnetId": "subnet-xyz",
                    }
                ):
                    project_id, updated = device_mod._resolve_project_id(session)
                    self.assertEqual(project_id, "project-handshake")
                    self.assertEqual(updated["projectId"], "project-handshake")
                    self.assertEqual(os.environ.get("NEBIUS_PROJECT_ID"), "project-handshake")
        finally:
            os.environ.clear()
            os.environ.update(env_snapshot)

    def test_server_bootstrap_on_404_then_retry(self) -> None:
        # Prepare a minimal session and environment
        session = {"token": "tok", "userId": "user", "login_url": "https://portal.example.com"}
        with patch.object(device_mod, "_load_session_payload", return_value=session), \
            patch.object(device_mod, "_resolve_project_id", return_value=("project-abc", session)), \
            patch.object(device_mod, "_fetch_handshake_bundle", return_value={
                "serviceAccountId": "sa-1",
                "authorizedKeyId": "key-1",
                "authorizedPrivateKey": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----",
                "endpoint": "api.nebius.cloud:443",
            }), \
            patch.object(device_mod, "_bootstrap_server_credentials", return_value=True) as mock_bootstrap:

            # Simulate first provision attempt 404 then successful JSON
            def _fake_urlopen(req, timeout=0):
                # Return different responses based on a simple counter
                if not hasattr(_fake_urlopen, "calls"):
                    _fake_urlopen.calls = 0  # type: ignore[attr-defined]
                _fake_urlopen.calls += 1  # type: ignore[attr-defined]
                if _fake_urlopen.calls == 1:  # type: ignore[attr-defined]
                    from urllib.error import HTTPError
                    raise HTTPError(req.full_url, 404, "Not Found", hdrs=None, fp=None)
                class _Resp:
                    status = 200
                    def read(self):
                        return json.dumps({
                            "instance": {
                                "instanceId": "i-1",
                                "host": "203.0.113.11",
                                "username": "ubuntu",
                                "bootDiskId": "d-1",
                                "platform": "gpu-h100-sxm",
                                "preset": "1gpu-16vcpu-200gb",
                                "gpuCount": 1,
                            },
                            "sshPrivateKey": "-----BEGIN OPENSSH PRIVATE KEY-----\n...\n-----END OPENSSH PRIVATE KEY-----",
                            "sshPublicKey": "ssh-rsa abc",
                            "sshFingerprint": "fp:xx",
                        }).encode("utf-8")
                    def __enter__(self):
                        return self
                    def __exit__(self, *args):
                        return False
                return _Resp()

            with patch.object(device_mod.urlrequest, "urlopen", side_effect=_fake_urlopen):
                result = device_mod._server_provision(
                    gpu_type="H100",
                    gpu_count=1,
                    storage_gib=512,
                    preset=None,
                    image_id=None,
                    subnet_id=None,
                    ttl_seconds=None,
                )
                self.assertIsInstance(result, dict)
                self.assertIn("instance", result)
                self.assertIn("ssh_private_key", result)
            mock_bootstrap.assert_called_once()

    def test_device_accepts_string_args_and_fallback(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_home, patch.dict(
            os.environ, {"HOME": tmp_home, "REBEL_FORGE_ALLOW_LEGACY_CREDENTIALS": "1"}, clear=False
        ):
            cli_dir = Path(tmp_home) / ".nebius" / "bin"
            cli_dir.mkdir(parents=True, exist_ok=True)
            (cli_dir / "nebius").write_text("#!/bin/sh\n", encoding="utf-8")

            bundle = {
                "projectId": "project-test",
                "serviceAccountId": "sa-test",
                "authorizedKeyId": "key-test",
                "authorizedPrivateKey": "-----BEGIN PRIVATE KEY-----\nabc\n-----END PRIVATE KEY-----",
                "sshPublicKey": "ssh-rsa dummy rebel@test",
                "sshPrivateKey": "-----BEGIN OPENSSH PRIVATE KEY-----\nxyz\n-----END OPENSSH PRIVATE KEY-----",
            }
            instance = device_mod.NebiusInstance(
                instance_id="instance-123",
                name="forge-instance",
                host="203.0.113.10",
                username="ubuntu",
                disk_id="disk-456",
                platform="gpu-h200-sxm",
                preset="preset-h200",
            )
            provisioner = MagicMock()
            provisioner.auto_preset.side_effect = [
                RemoteError("not found"),
                "preset-h200",
            ]
            provisioner.provision_instance.return_value = instance

            with patch.object(device_mod, "_server_provision", side_effect=RemoteError("offline")), patch.object(
                device_mod, "_resolve_credentials", return_value=bundle
            ), patch.object(
                device_mod, "_wait_for_ssh"
            ), patch.object(device_mod, "ensure_remote", side_effect=SystemExit), patch.object(
                device_mod, "NebiusProvisioner", return_value=provisioner
            ):
                with self.assertRaises(SystemExit):
                    device_mod.device("H100", "2 GPUs", "512 gib")

        self.assertEqual(provisioner.auto_preset.call_count, 2)
        first_call = provisioner.auto_preset.call_args_list[0]
        self.assertEqual(first_call.kwargs["platform"], "gpu-h100-sxm")
        self.assertEqual(first_call.kwargs["gpu_count"], 2)
        second_call = provisioner.auto_preset.call_args_list[1]
        self.assertEqual(second_call.kwargs["platform"], "gpu-h200-sxm")
        self.assertEqual(second_call.kwargs["gpu_count"], 2)
        provisioner.provision_instance.assert_called_once()
        kwargs = provisioner.provision_instance.call_args.kwargs
        self.assertEqual(kwargs["platform"], "gpu-h200-sxm")
        self.assertEqual(kwargs["boot_disk_gib"], 512)

    def test_device_raises_when_remote_unavailable_without_legacy(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_home, patch.dict(
            os.environ, {"HOME": tmp_home}, clear=False
        ):
            with patch.object(
                device_mod, "_server_provision", side_effect=RemoteError("service down")
            ), patch.object(device_mod, "_resolve_credentials") as mock_resolve:
                with self.assertRaises(RemoteError):
                    device_mod.device("H100", "512")
            mock_resolve.assert_not_called()

    def test_resolve_project_id_fetches_metadata(self) -> None:
        session_payload = {
            "token": "token-123",
            "userId": "user-456",
            "login_url": "https://portal.example.com",
        }
        metadata = {
            "projectId": "project-xyz",
            "endpoint": "api.nebius.cloud:443",
            "defaultImageId": "image-abc",
            "defaultSubnetId": "subnet-789",
        }
        original_env = {
            "project_id": os.environ.get("project_id"),
            "NEBIUS_PROJECT_ID": os.environ.get("NEBIUS_PROJECT_ID"),
            "NEBIUS_ENDPOINT": os.environ.get("NEBIUS_ENDPOINT"),
            "NEBIUS_IMAGE_ID": os.environ.get("NEBIUS_IMAGE_ID"),
            "NEBIUS_SUBNET_ID": os.environ.get("NEBIUS_SUBNET_ID"),
        }
        try:
            with patch.object(device_mod, "_fetch_default_project_metadata", return_value=metadata):
                project_id, updated = device_mod._resolve_project_id(session_payload)
            self.assertEqual(project_id, "project-xyz")
            self.assertEqual(updated["projectId"], "project-xyz")
            self.assertEqual(os.environ.get("project_id"), "project-xyz")
            self.assertEqual(os.environ.get("NEBIUS_PROJECT_ID"), "project-xyz")
            self.assertEqual(os.environ.get("NEBIUS_ENDPOINT"), "api.nebius.cloud:443")
            self.assertEqual(os.environ.get("NEBIUS_IMAGE_ID"), "image-abc")
            self.assertEqual(os.environ.get("NEBIUS_SUBNET_ID"), "subnet-789")
        finally:
            for key, value in original_env.items():
                if value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = value


if __name__ == "__main__":
    unittest.main()
