#!/usr/bin/env python3
"""
ProxMenux Terminal WebSocket Routes
Provides a WebSocket endpoint for interactive terminal sessions
"""

from flask import Blueprint, jsonify, request
from flask_sock import Sock
import subprocess
import os
import pty
import select
import struct
import fcntl
import termios
import threading
import time
import requests

terminal_bp = Blueprint('terminal', __name__)
sock = Sock()

# Active terminal sessions - now stores session by session_id string
sessions = {}

def set_winsize(fd, rows, cols):
    """Set terminal window size"""
    try:
        winsize = struct.pack('HHHH', rows, cols, 0, 0)
        fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
    except Exception as e:
        print(f"Error setting window size: {e}")

def read_and_forward_output(master_fd, ws):
    """Read from PTY and send to WebSocket"""
    while True:
        try:
            # Use select with timeout to check if data is available
            r, _, _ = select.select([master_fd], [], [], 0.01)
            if master_fd in r:
                try:
                    data = os.read(master_fd, 4096)
                    if data:
                        ws.send(data.decode('utf-8', errors='ignore'))
                    else:
                        break
                except OSError:
                    break
        except Exception as e:
            print(f"Error reading from PTY: {e}")
            break

@terminal_bp.route('/api/terminal/health', methods=['GET'])
def terminal_health():
    """Health check for terminal service"""
    return {'success': True, 'active_sessions': len(sessions)}

@terminal_bp.route('/api/terminal/<session_id>/resize', methods=['POST'])
def resize_terminal(session_id):
    """Resize the PTY for a given terminal session."""
    if session_id not in sessions:
        return jsonify({'error': 'Session not found'}), 404
    
    try:
        data = request.get_json()
        cols = int(data.get('cols', 120))
        rows = int(data.get('rows', 30))
        
        # Resize the PTY to match the frontend terminal dimensions
        master_fd = sessions[session_id]['master_fd']
        set_winsize(master_fd, rows, cols)
        
        print(f"[v0] Terminal {session_id} resized to {cols}x{rows}")
        
        return jsonify({'status': 'success', 'cols': cols, 'rows': rows})
    except Exception as e:
        print(f"Error resizing terminal {session_id}: {e}")
        return jsonify({'error': str(e)}), 500

@sock.route('/ws/terminal')
def terminal_websocket(ws):
    """WebSocket endpoint for terminal sessions"""
    
    # Create pseudo-terminal
    master_fd, slave_fd = pty.openpty()
    
    # Start bash process
    shell_process = subprocess.Popen(
        ['/bin/bash', '-i'],
        stdin=slave_fd,
        stdout=slave_fd,
        stderr=slave_fd,
        preexec_fn=os.setsid,
        cwd='/',
        env=dict(os.environ, TERM='xterm-256color', PS1='\\u@\\h:\\w\\$ ')
    )
    
    session_id = str(int(time.time() * 1000))
    sessions[session_id] = {
        'process': shell_process,
        'master_fd': master_fd
    }
    
    # Set non-blocking mode for master_fd
    flags = fcntl.fcntl(master_fd, fcntl.F_GETFL)
    fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
    
    set_winsize(master_fd, 30, 80)
    
    # Send session_id to frontend
    ws.send(f"\x1b]0;SESSION_ID:{session_id}\x07")
    
    # Start thread to read PTY output and forward to WebSocket
    output_thread = threading.Thread(
        target=read_and_forward_output,
        args=(master_fd, ws),
        daemon=True
    )
    output_thread.start()
    
    try:
        while True:
            # Receive data from WebSocket (blocking)
            data = ws.receive(timeout=None)
            
            if data is None:
                # Client closed connection
                break
            
            if data.startswith('\x1b[8;'):
                try:
                    # Parse: \x1b[8;{rows};{cols}t
                    parts = data[4:-1].split(';')
                    if len(parts) >= 2:
                        rows, cols = int(parts[0]), int(parts[1])
                        set_winsize(master_fd, rows, cols)
                        print(f"[v0] Terminal resized via WebSocket to {rows}x{cols}")
                    continue
                except Exception as e:
                    print(f"Error parsing resize: {e}")
                    pass

            # Send input to bash
            try:
                os.write(master_fd, data.encode('utf-8'))
            except OSError as e:
                print(f"Error writing to PTY: {e}")
                break
            
            # Check if process is still alive
            if shell_process.poll() is not None:
                break
                
    except Exception as e:
        print(f"Terminal session error: {e}")
    finally:
        # Cleanup
        try:
            shell_process.terminate()
            shell_process.wait(timeout=1)
        except:
            try:
                shell_process.kill()
            except:
                pass
        
        try:
            os.close(master_fd)
        except:
            pass
        
        try:
            os.close(slave_fd)
        except:
            pass
        
        if session_id in sessions:
            del sessions[session_id]

def init_terminal_routes(app):
    """Initialize terminal routes with Flask app"""
    sock.init_app(app)
    app.register_blueprint(terminal_bp)