import sys
import os

# Ensure the parent directory (libs/python) is in sys.path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from flash_sandbox import SandboxClient

def test_lifecycle():
    print("Testing Sandbox Lifecycle via gRPC...")

    try:
        client = SandboxClient(host="localhost", port=50051)
    except Exception as e:
        print(f"Failed to initialize SandboxClient: {e}")
        return

    sandbox_ids = []

    try:
        # 1. Start Docker Sandbox
        print("1. Starting Docker sandbox...")
        docker_id = client.start_sandbox(
            type="docker",
            image="python:3.11-alpine",
            command=["tail", "-f", "/dev/null"],
            memory_mb=128,
            cpu_cores=0.5
        )
        print(f"   Created Docker sandbox with ID: {docker_id}")
        sandbox_ids.append(docker_id)

        # 2. Skip gVisor Sandbox since it hangs on this host
        # 3. Get Status
        print("3. Getting status...")
        for sid in sandbox_ids:
            status = client.get_status(sid)
            print(f"   Sandbox {sid} Status: {status}")
            if status != "running":
                 print(f"   Warning: Status is {status}, expected running")

        # 4. Exec Command before Snapshot
        print("4. Executing command before snapshot 'echo hello'...")
        for sid in sandbox_ids:
            exec_res = client.exec_command(sid, ["echo", f"hello from {sid}"])
            print(f"   Sandbox {sid} Output: {exec_res.stdout.strip()}")

            if f"hello from {sid}" not in exec_res.stdout:
                print(f"   Error: Output did not contain expected text")

        # 5. Get Platform Info and Run Python
        print("5. Running python and getting platform info...")
        plat_info = client.get_platform_info(docker_id)
        print(f"   Platform Info: {plat_info}")
        if not isinstance(plat_info, dict):
            print(f"   Error: Platform info should be a dict, got {type(plat_info).__name__}")
        if not plat_info:
            print(f"   Error: Platform info empty")

        py_res = client.run_python(docker_id, "print('sdk test')")
        print(f"   Python execution: {py_res.strip()}")
        if "sdk test" not in py_res:
            print(f"   Error: Python snippet output missing")

        # 6. Snapshot and Resume (Docker)
        print("6. Snapshotting and Resuming Docker sandbox...")
        try:
             snap_res = client.snapshot_sandbox(docker_id)
             print(f"   Snapshot created: {snap_res}")
             client.resume_sandbox(docker_id)
             print(f"   Resumed successfully!")

             # Run exec again to verify
             exec_res_post = client.exec_command(docker_id, ["echo", "post-resume docker"])
             print(f"   Post-Resume Output: {exec_res_post.stdout.strip()}")
        except Exception as e:
             print(f"   Failed snapshot/resume on Docker: {e}")

    except Exception as e:
        print(f"Test failed with error: {e}")
    finally:
        # 7. Cleanup
        print("7. Cleaning up sandboxes...")
        for sid in sandbox_ids:
             try:
                 client.stop_sandbox(sid)
                 print(f"   Stopped {sid}")
             except:
                 pass
        client.close()

if __name__ == "__main__":
    test_lifecycle()
