import json
import unittest
from unittest.mock import patch, MagicMock
import base64
from instavm import InstaVM

LIVE_API_KEY = ""

class TestInstaVMClient(unittest.TestCase):
    @unittest.skipIf(not LIVE_API_KEY, "No live API key set in LIVE_API_KEY")
    def test_live_api_execute(self):
        client = InstaVM(api_key=LIVE_API_KEY, base_url="https://api.instavm.io")
        # 2. Execute a simple command
        exec_result = client.execute("print('Hello from live test')")
        print(exec_result)
        self.assertTrue(isinstance(exec_result, dict))
        print("Execution result:", exec_result)

        # 3. Get usage info for this session
        usage_info = client.get_usage()
        self.assertTrue(isinstance(usage_info, dict))
        print("Usage info:", usage_info)




    @patch('instavm.sandbox_client.InstaVM._make_request')
    def setUp(self, mock_make_request):
        # Mock _make_request for session initialization
        mock_response = MagicMock()
        mock_response.json.return_value = {"session_id": "init_session_id"}
        mock_make_request.return_value = mock_response

        self.client = InstaVM(api_key="test_api_key", base_url="http://api.instavm.io")
        self.mock_session_id = "test_session_id"


    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_start_session(self, mock_make_request):

        mock_response = MagicMock()
        mock_response.json.return_value = {"session_id": self.mock_session_id}
        mock_make_request.return_value = mock_response

        session_id = self.client.start_session()
        self.assertEqual(session_id, self.mock_session_id)
        mock_make_request.assert_called_with(
            "POST",
            f"{self.client.base_url}/v1/sessions/session",
            json={
                "api_key": self.client.api_key,
                "vm_lifetime_seconds": self.client.timeout
            }
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_start_session_with_bootstrap_options(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.json.return_value = {"session_id": "session_with_resources"}
        mock_make_request.return_value = mock_response

        client = InstaVM(
            api_key="test_api_key",
            base_url="http://api.instavm.io",
            cpu_count=4,
            memory_mb=2048,
            env={"APP_ENV": "test"},
            metadata={"project": "sdk"}
        )

        self.assertEqual(client.session_id, "session_with_resources")
        mock_make_request.assert_called_with(
            "POST",
            f"{client.base_url}/v1/sessions/session",
            json={
                "api_key": client.api_key,
                "vm_lifetime_seconds": client.timeout,
                "memory_mb": 2048,
                "vcpu_count": 4,
                "metadata": {"project": "sdk"},
                "env": {"APP_ENV": "test"}
            }
        )


    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_execute(self, mock_make_request):

        self.client.session_id = self.mock_session_id
        command = "test_command"

        mock_response = MagicMock()
        mock_response.json.return_value = {"result": "success"}
        mock_make_request.return_value = mock_response

        response = self.client.execute(command)
        self.assertEqual(response, {"result": "success"})
        mock_make_request.assert_called_once_with(
            "POST",
            f"{self.client.base_url}/execute",
            headers={"X-API-Key": self.client.api_key},
            json={
                "command": command,
                "session_id": self.client.session_id
            },
            timeout=self.client.timeout
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_execute_async(self, mock_make_request):
        self.client.session_id = self.mock_session_id
        command = "test_async_command"

        mock_response = MagicMock()
        mock_response.json.return_value = {"result": "async_success"}
        mock_make_request.return_value = mock_response

        response = self.client.execute_async(command, language="python", timeout=60)
        self.assertEqual(response, {"result": "async_success"})
        mock_make_request.assert_called_once_with(
            "POST",
            f"{self.client.base_url}/execute_async",
            headers={"X-API-Key": self.client.api_key},
            json={
                "command": command,
                "session_id": self.client.session_id,
                "language": "python",
                "timeout": 60
            }
        )


    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_get_usage(self, mock_make_request):

        self.client.session_id = self.mock_session_id

        mock_response = MagicMock()
        mock_response.json.return_value = {"usage": "test_usage"}
        mock_make_request.return_value = mock_response

        usage = self.client.get_usage()
        self.assertEqual(usage, {"usage": "test_usage"})
        mock_make_request.assert_called_once_with(
            "GET",
            f"{self.client.base_url}/v1/sessions/usage/{self.client.session_id}"
        )


    @patch('instavm.sandbox_client.InstaVM._make_request')
    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data="data")
    def test_upload_file(self, mock_open, mock_make_request):
        self.client.session_id = self.mock_session_id
        mock_response = MagicMock()
        mock_response.json.return_value = {"status": "file_uploaded"}
        mock_make_request.return_value = mock_response

        response = self.client.upload_file('test.txt', '/remote/test.txt')
        self.assertEqual(response, {"status": "file_uploaded"})
        mock_make_request.assert_called_once()

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_snapshots_create_with_build_args(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"id": "snap_123", "status": "building"}
        mock_make_request.return_value = mock_response

        response = self.client.snapshots.create(
            oci_image="docker.io/library/python:3.11-slim",
            name="python-base",
            build_args={
                "git_clone_url": "https://github.com/example/repo.git",
                "git_clone_branch": "main",
                "envs": {"APP_ENV": "test"}
            }
        )

        self.assertEqual(response["id"], "snap_123")
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "POST")
        self.assertEqual(args[1], f"{self.client.base_url}/v1/snapshots")
        self.assertIn("build_args", kwargs["json"])
        self.assertEqual(
            kwargs["json"]["build_args"]["git_clone_url"],
            "https://github.com/example/repo.git"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_vm_clone(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"vm_id": "vm_clone_1"}
        mock_make_request.return_value = mock_response

        response = self.client.vms.clone(
            vm_id="vm_src_1",
            wait=True,
            snapshot_name="clone-seed"
        )

        self.assertEqual(response["vm_id"], "vm_clone_1")
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "POST")
        self.assertEqual(args[1], f"{self.client.base_url}/v1/vms/vm_src_1/clone")
        self.assertEqual(kwargs["params"]["wait"], "true")
        self.assertEqual(kwargs["json"]["snapshot_name"], "clone-seed")

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_share_create_uses_current_session(self, mock_make_request):
        self.client.session_id = "session_current_1"
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"share_id": "share_123", "port": 3000}
        mock_make_request.return_value = mock_response

        response = self.client.shares.create(port=3000, is_public=True)

        self.assertEqual(response["share_id"], "share_123")
        _, kwargs = mock_make_request.call_args
        self.assertEqual(kwargs["json"]["session_id"], "session_current_1")
        self.assertEqual(kwargs["json"]["port"], 3000)
        self.assertTrue(kwargs["json"]["is_public"])

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_computer_use_proxy_get(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"ok": True}
        mock_make_request.return_value = mock_response

        response = self.client.computer_use.get("session_1", "state")

        self.assertEqual(response, {"ok": True})
        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/computeruse/session_1/state"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_get_session_info_escapes_path_segment(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"ok": True}
        mock_make_request.return_value = mock_response

        self.client.get_session_info("session/../../etc")

        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/sessions/session/session%2F..%2F..%2Fetc"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_vms_get_escapes_path_segment(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"vm_id": "vm_1"}
        mock_make_request.return_value = mock_response

        self.client.vms.get("vm/../../snapshots")

        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/vms/vm%2F..%2F..%2Fsnapshots"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_volumes_get_escapes_path_segment(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"id": "vol_1"}
        mock_make_request.return_value = mock_response

        self.client.volumes.get("vol/../../etc")

        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/volumes/vol%2F..%2F..%2Fetc"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    @patch('builtins.open', new_callable=unittest.mock.mock_open, read_data=b"hello")
    def test_volumes_upload_file_uses_multipart(self, mock_open_file, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"path": "docs/readme.txt", "size": 5}
        mock_make_request.return_value = mock_response

        result = self.client.volumes.upload_file(
            volume_id="vol_1",
            file_path="/tmp/readme.txt",
            path="docs/readme.txt",
            overwrite=True
        )

        self.assertEqual(result["path"], "docs/readme.txt")
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "POST")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/volumes/vol_1/files/upload"
        )
        self.assertEqual(kwargs["data"]["path"], "docs/readme.txt")
        self.assertEqual(kwargs["data"]["overwrite"], "true")
        self.assertIn("file", kwargs["files"])

    @patch('instavm.sandbox_client.InstaVM._make_request')
    @patch('builtins.open', new_callable=unittest.mock.mock_open)
    def test_volumes_download_file_to_local_path(self, mock_open_file, mock_make_request):
        encoded = base64.b64encode(b"hello").decode("utf-8")
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {
            "path": "docs/readme.txt",
            "filename": "readme.txt",
            "size": 5,
            "content": encoded
        }
        mock_make_request.return_value = mock_response

        result = self.client.volumes.download_file(
            volume_id="vol_1",
            path="docs/readme.txt",
            local_path="/tmp/readme.txt"
        )

        self.assertEqual(result["local_path"], "/tmp/readme.txt")
        self.assertEqual(result["saved_bytes"], 5)
        mock_open_file.assert_called_with("/tmp/readme.txt", "wb")
        mock_open_file().write.assert_called_with(b"hello")

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_vms_mount_and_unmount_volume(self, mock_make_request):
        mount_response = MagicMock()
        mount_response.content = b'{}'
        mount_response.json.return_value = {
            "volume_id": "vol_1",
            "mount_path": "/data",
            "mode": "rw",
            "status": "mounted"
        }
        mock_make_request.return_value = mount_response

        mounted = self.client.vms.mount_volume(
            vm_id="vm/1",
            volume_id="vol_1",
            mount_path="/data",
            mode="rw",
            checkpoint_id="cp_1",
            wait=False
        )
        self.assertEqual(mounted["status"], "mounted")
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "POST")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/vms/vm%2F1/volumes"
        )
        self.assertEqual(kwargs["params"]["wait"], "false")
        self.assertEqual(kwargs["json"]["checkpoint_id"], "cp_1")

        list_response = MagicMock()
        list_response.content = b'[]'
        list_response.json.return_value = [mounted]
        mock_make_request.return_value = list_response

        volumes = self.client.vms.list_volumes("vm/1")
        self.assertEqual(len(volumes), 1)
        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/vms/vm%2F1/volumes"
        )

        unmount_response = MagicMock()
        unmount_response.content = b''
        mock_make_request.return_value = unmount_response

        self.client.vms.unmount_volume(
            vm_id="vm/1",
            volume_id="vol/1",
            mount_path="/data",
            wait=True
        )
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "DELETE")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/vms/vm%2F1/volumes/vol%2F1"
        )
        self.assertEqual(kwargs["params"]["mount_path"], "/data")
        self.assertEqual(kwargs["params"]["wait"], "true")

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_path_segment_rejects_dot_and_dotdot(self, mock_make_request):
        calls_before = mock_make_request.call_count

        with self.assertRaises(ValueError):
            self.client.volumes.get("..")
        with self.assertRaises(ValueError):
            self.client.vms.unmount_volume("vm_1", ".", mount_path="/data")

        self.assertEqual(mock_make_request.call_count, calls_before)

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_computer_use_proxy_rejects_path_traversal(self, mock_make_request):
        calls_before = mock_make_request.call_count
        with self.assertRaises(ValueError):
            self.client.computer_use.get("session_1", "../state")

        # Traversal is blocked before an API request is made.
        self.assertEqual(mock_make_request.call_count, calls_before)

    # --- Drift-fix tests (v0.13.0) ---

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_get_session_app_url_with_port(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"app_url": "https://example.com:8080"}
        mock_make_request.return_value = mock_response

        result = self.client.get_session_app_url(port=8080)

        self.assertEqual(result["app_url"], "https://example.com:8080")
        args, kwargs = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(kwargs["params"], {"port": 8080})

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_get_session_app_url_without_port(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"app_url": "https://example.com"}
        mock_make_request.return_value = mock_response

        self.client.get_session_app_url()

        _, kwargs = mock_make_request.call_args
        self.assertIsNone(kwargs.get("params"))

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_list_sandboxes_with_metadata_and_limit(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'[]'
        mock_response.json.return_value = [{"session_id": "s1"}]
        mock_make_request.return_value = mock_response

        result = self.client.list_sandboxes(
            metadata={"env": "dev"},
            limit=50
        )

        self.assertEqual(len(result), 1)
        _, kwargs = mock_make_request.call_args
        self.assertEqual(kwargs["params"]["metadata"], json.dumps({"env": "dev"}))
        self.assertEqual(kwargs["params"]["limit"], 50)

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_list_sandboxes_without_params(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'[]'
        mock_response.json.return_value = []
        mock_make_request.return_value = mock_response

        self.client.list_sandboxes()

        _, kwargs = mock_make_request.call_args
        self.assertIsNone(kwargs.get("params"))

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_computer_use_head(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"ok": True}
        mock_make_request.return_value = mock_response

        self.client.computer_use.head("session_1", "status")

        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "HEAD")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/computeruse/session_1/status"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_computer_use_vnc_websockify(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"ws_url": "wss://example.com/vnc"}
        mock_make_request.return_value = mock_response

        result = self.client.computer_use.vnc_websockify("session_1")

        self.assertEqual(result["ws_url"], "wss://example.com/vnc")
        args, _ = mock_make_request.call_args
        self.assertEqual(args[0], "GET")
        self.assertEqual(
            args[1],
            f"{self.client.base_url}/v1/computeruse/session_1/vnc/websockify"
        )

    @patch('instavm.sandbox_client.InstaVM._make_request')
    def test_vms_create_with_volumes(self, mock_make_request):
        mock_response = MagicMock()
        mock_response.content = b'{}'
        mock_response.json.return_value = {"vm_id": "vm_456"}
        mock_make_request.return_value = mock_response

        result = self.client.vms.create(
            vm_lifetime_seconds=3600,
            volumes=[
                {"volume_id": "vol_1", "mount_path": "/data", "mode": "rw"},
                {"volume_id": "vol_2", "mount_path": "/cache", "mode": "ro"},
            ]
        )

        self.assertEqual(result["vm_id"], "vm_456")
        _, kwargs = mock_make_request.call_args
        self.assertEqual(len(kwargs["json"]["volumes"]), 2)
        self.assertEqual(kwargs["json"]["volumes"][0]["volume_id"], "vol_1")
        self.assertEqual(kwargs["json"]["vm_lifetime_seconds"], 3600)

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