import docker
import os

# 1. Attempt to load your Drive library so the Engine can use it
try:
    from janu_vfs.drive import JanuDrive
except ImportError:
    print("[Janu PTY] ⚠️ Warning: janu_vfs library not found on the host system!")


class JanuEngine:
    """
    The JanuOS Multi-Session Sandbox Spawner.
    Handles parallel containers, Offline Vaults, and Google Drive VFS Sync.
    """
    def __init__(self, workspace_path="./workspace", exposed_ports=None):
        self.workspace_path = workspace_path  
        self.exposed_ports = exposed_ports or []
        self.containers = {}
        self.drive = None # Initialize drive variable
        
        # 2. Smart Environment Detector & Infinite Timeout
        self.mode = "docker"
        try:
            self.client = docker.from_env(timeout=None)
            self.client.ping()
            print("[Janu PTY] 🐳 Docker Engine detected! Running in pure Sandbox Mode.")
        except Exception:
            print("[Janu PTY] ⚠️ Docker not found! Switching to Native Host Mode...")
            self.mode = "native"
            self.client = None

        self.images = {
            "python": "python:3.9-slim",
            "c": "gcc:12",               
            "c++": "gcc:12",              
            "java": "eclipse-temurin:17-jdk-alpine", 
            "ruby": "ruby:3.2-slim",       
            "node": "node:18-slim"
        }

        # 3. Storage Routing: Local vs Cloud
        if self.workspace_path.startswith("drive://") or "Google Drive" in self.workspace_path:
            folder_id = self.workspace_path.replace("drive://", "")
            print(f"[Janu PTY] ☁️ Cloud Drive detected! Target ID: {folder_id}")
            
            # Create the hidden cache folder
            self.actual_mount_path = os.path.abspath("./.janu_vfs_cache")
            os.makedirs(self.actual_mount_path, exist_ok=True)
            
            print(f"[Janu PTY] 🔄 Connecting to Google Drive API from Host...")
            try:
                # Initialize your Drive library to download the files BEFORE booting the terminal!
                self.drive = JanuDrive(master_id=folder_id)
                
                # IMPORTANT: If your download function is named differently in janu_vfs, change it here!
                self.drive.download_all(self.actual_mount_path) 
                
                print(f"[Janu PTY] ✅ Drive files synced to {self.actual_mount_path}")
            except Exception as e:
                print(f"[Janu PTY] ❌ Drive Sync failed: {e}")
                
        else:
            # Local Storage Mode
            self.actual_mount_path = os.path.abspath(self.workspace_path)
            if not os.path.exists(self.actual_mount_path):
                os.makedirs(self.actual_mount_path, exist_ok=True)


    def boot(self, tab_id, language):
        """Boots a specific language container for a specific UI tab."""
        image_name = self.images.get(language.lower(), "ubuntu:20.04")
        
        if self.mode == "native":
            print(f"[Janu PTY] Native Mode initialized for {tab_id}.")
            self.containers[tab_id] = "NATIVE_MOCK"
            return "NATIVE_MOCK"

        print(f"[Janu PTY] Preparing Sandbox for {tab_id}: {image_name}...")
        
        # 4. The Offline-First Detective
        try:
            self.client.images.get(image_name)
        except docker.errors.ImageNotFound:
            safe_name = image_name.replace(':', '_').replace('/', '_')
            offline_file = f"{safe_name}_offline.tar"
            tar_path = os.path.join(os.getcwd(), offline_file)
            
            if os.path.exists(tar_path):
                print(f"[Janu PTY] 🕵️ Offline vault found: {tar_path}!")
                with open(tar_path, 'rb') as f:
                    self.client.images.load(f)
            else:
                print(f"[Janu PTY] 🌐 Downloading {image_name} from internet...")
                self.client.images.pull(image_name)

        # 5. Boot the container with Collision-Proof EXPOSED PORTS
        import socket
        port_bindings = {}
        
        for port in self.exposed_ports:
            target_port = port + len(self.containers)
            # Search for an open port if this one is already taken!
            while True:
                try:
                    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                        s.bind(('', target_port))
                    break  # If we get here without an error, the port is free!
                except OSError:
                    target_port += 1  # The port was taken, try the next one!
                    
            port_bindings[f"{port}/tcp"] = target_port

        container = self.client.containers.run(
            image_name,
            command="tail -f /dev/null",
            detach=True,
            tty=True,
            stdin_open=True,
            working_dir="/workspace",
            volumes={self.actual_mount_path: {'bind': '/workspace', 'mode': 'rw'}},
            ports=port_bindings,  # <--- Uses the collision-proof ports!
            remove=True
        )

        self.containers[tab_id] = container
        print(f"[Janu PTY] Sandbox {tab_id} Live! ID: {container.short_id}")
        return container

    def shutdown(self, tab_id):
        """Kills a specific tab's container to free up RAM."""
        if self.mode == "docker" and tab_id in self.containers:
            container = self.containers[tab_id]
            print(f"[Janu PTY] Destroying Sandbox for {tab_id}...")
            container.stop()
            del self.containers[tab_id]

    def shutdown_all(self):
        """Emergency cleanup on exit and Cloud Sync."""
        # 1. Kill all active terminal sessions
        for tab_id in list(self.containers.keys()):
            self.shutdown(tab_id)
            
        # 2. Push all local changes back to Google Drive!
        if self.drive:
            print(f"[Janu PTY] ☁️ Uploading changes back to Google Drive...")
            try:
                # IMPORTANT: If your upload function in janu_vfs is named differently 
                # (like sync_to_cloud or upload_folder), change it here!
                self.drive.upload_all(self.actual_mount_path) 
                print(f"[Janu PTY] ✅ Cloud sync complete!")
            except Exception as e:
                print(f"[Janu PTY] ❌ Cloud sync failed: {e}")
