import eventlet
eventlet.monkey_patch()
import pty
import subprocess
import select 
import socket
import os
from flask import Flask
from flask_socketio import SocketIO
from .engine import JanuEngine

class JanuTerminalBridge:
    """The Multi-Tab Web Bridge."""
    def __init__(self, workspace_path="./workspace", exposed_ports=None):
        self.engine = JanuEngine(workspace_path, exposed_ports)
        self.app = Flask(__name__)
        self.socketio = SocketIO(self.app, cors_allowed_origins="*", async_mode='eventlet')
        self.port = self._get_free_port()
        
        # Track the raw bash sockets for each tab
        self.exec_sockets = {}
        
        self._setup_routes()

    def _get_free_port(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('', 0))
            return s.getsockname()[1]

    def _setup_routes(self):
        @self.app.route('/')
        def index():
            html_path = os.path.join(os.path.dirname(__file__), 'static', 'terminal.html')
            with open(html_path, 'r') as f:
                return f.read()

        @self.socketio.on('connect')
        def on_connect():
            print(f"[Janu PTY] Browser connected. Awaiting tab spawns...")

        # --- NEW: Spawn a specific tab ---
        # --- NEW: Spawn a specific tab ---
        @self.socketio.on('spawn-terminal')
        def on_spawn_terminal(data):
            tab_id = data['tabId']
            language = data['language']
            
            container = self.engine.boot(tab_id, language)
            
            if self.engine.mode == "docker":
                image_name = self.engine.images.get(language.lower(), "ubuntu:20.04")
                shell_cmd = "/bin/sh" if "alpine" in image_name else "/bin/bash"

                exec_id = self.engine.client.api.exec_create(
                    container.id, 
                    cmd=shell_cmd, 
                    stdin=True, tty=True, stdout=True, stderr=True
                )
                
                # FIX: Request the raw, two-way socket!
                exec_sock = self.engine.client.api.exec_start(exec_id, socket=True, tty=True)
                
                # Extract the raw OS socket object so we can read and write to it
                raw_sock = getattr(exec_sock, '_sock', exec_sock)
                self.exec_sockets[tab_id] = raw_sock
                
                # Force the bash prompt to wake up and appear!
                raw_sock.send(b"\n")
                
                self.socketio.start_background_task(target=self._read_from_docker, tab_id=tab_id)

            elif self.engine.mode == "native":
                # 1. Create a true Linux Pseudo-Terminal (PTY)
                master_fd, slave_fd = pty.openpty()
                
                # 2. Spawn a local bash session attached to the PTY slave
                proc = subprocess.Popen(
                    ["/bin/bash"],
                    stdin=slave_fd,
                    stdout=slave_fd,
                    stderr=slave_fd,
                    cwd=self.engine.actual_mount_path, 
                    preexec_fn=os.setsid 
                )
                
                # 3. Store the master file descriptor and process
                self.exec_sockets[tab_id] = {
                    "fd": master_fd,
                    "proc": proc
                }
                
                # ---> ADD THIS NEW LINE <---
                # Force the terminal to instantly cd into the Drive cache and clear the screen
                os.write(master_fd, f"cd {self.engine.actual_mount_path} && clear\n".encode('utf-8'))
                
                # 4. Start the background streaming task
                self.socketio.start_background_task(target=self._read_from_native, tab_id=tab_id)
        # --- UPDATED: Route input to the correct tab ---
        @self.socketio.on('pty-input')
        def on_pty_input(data):
            tab_id = data['tabId']
            if self.engine.mode == "docker" and tab_id in self.exec_sockets:
                # Use .send() for Docker sockets
                self.exec_sockets[tab_id].send(data['input'].encode('utf-8'))
            elif self.engine.mode == "native" and tab_id in self.exec_sockets:
                # Use os.write() for Native file descriptors
                fd = self.exec_sockets[tab_id]["fd"]
                os.write(fd, data['input'].encode('utf-8'))

        # --- NEW: Cleanup when a tab closes ---
        @self.socketio.on('close-terminal')
        def on_close_terminal(data):
            tab_id = data['tabId']
            if self.engine.mode == "docker" and tab_id in self.exec_sockets:
                self.exec_sockets[tab_id].close()
                del self.exec_sockets[tab_id]
            elif self.engine.mode == "native" and tab_id in self.exec_sockets:
                # Terminate the local process and close the PTY
                proc = self.exec_sockets[tab_id]["proc"]
                fd = self.exec_sockets[tab_id]["fd"]
                proc.terminate()
                os.close(fd)
                del self.exec_sockets[tab_id]
                
            self.engine.shutdown(tab_id)


    def _read_from_docker(self, tab_id):
        """Continuously streams output for a Docker PTY tab."""
        raw_sock = self.exec_sockets.get(tab_id)
        if not raw_sock:
            return
        try:
            while True:
                # Let eventlet breathe
                eventlet.sleep(0.01)
                
                # Wait until Docker actually has text to show us
                r, _, _ = select.select([raw_sock], [], [], 0)
                if r:
                    out = raw_sock.recv(1024)
                    if not out:
                        break
                    # Beam it to the browser!
                    self.socketio.emit('pty-output', {'tabId': tab_id, 'output': out.decode('utf-8', errors='replace')})
        except Exception as e:
            print(f"[Janu PTY] ⚠️ Docker stream closed for {tab_id}. Reason: {repr(e)}")

    def _read_from_native(self, tab_id):
        """Continuously streams output for a Native PTY tab."""
        if tab_id not in self.exec_sockets:
            return
            
        fd = self.exec_sockets[tab_id]["fd"]
        try:
            while True:
                # Let eventlet breathe so the server doesn't freeze
                eventlet.sleep(0.01) 
                
                # Check if there is data ready to be read from the PTY
                r, _, _ = select.select([fd], [], [], 0)
                if r:
                    out = os.read(fd, 1024)
                    if not out:
                        break
                    # Beam it back to the frontend!
                    self.socketio.emit('pty-output', {'tabId': tab_id, 'output': out.decode('utf-8', errors='replace')})
        except Exception as e:
            print(f"[Janu PTY] ⚠️ Native stream closed for {tab_id}. Reason: {repr(e)}")

    def start(self):
        print("="*60)
        print(f" 🚀 JANU PTY MULTIPLEXER LIVE! http://localhost:{self.port}")
        print("="*60)
        self.socketio.run(self.app, host='0.0.0.0', port=self.port, debug=False, use_reloader=False)

    def shutdown(self):
        for sock in self.exec_sockets.values():
            sock.close()
        self.engine.shutdown_all()

# --- NEW: Real-time Cloud Sync Trigger ---
        @self.socketio.on('sync-cloud')
        def on_sync_cloud():
            print("[Janu PTY] ⚡ Manual Cloud Sync Requested by UI!")
            if hasattr(self.engine, 'drive') and self.engine.drive:
                try:
                    # Run the upload in a background task so it doesn't freeze the terminal
                    self.socketio.start_background_task(
                        target=self.engine.drive.upload_all, 
                        local_dir=self.engine.actual_mount_path
                    )
                    # Tell the frontend it worked!
                    self.socketio.emit('sync-status', {'status': 'success', 'message': 'Cloud Sync Complete!'})
                except Exception as e:
                    print(f"[Janu PTY] ❌ Manual Cloud Sync Failed: {e}")
                    self.socketio.emit('sync-status', {'status': 'error', 'message': str(e)})
            else:
                 self.socketio.emit('sync-status', {'status': 'error', 'message': 'Drive not initialized.'})
