import os
import pty
import select
import struct
import fcntl
import termios
import uuid
import threading
import signal
import logging
import json
import time
from pathlib import Path
from flask import Flask, render_template, request, send_from_directory, jsonify
from flask_socketio import SocketIO, emit
import gevent

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['SECRET_KEY'] = os.urandom(24).hex()
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB upload limit

socketio = SocketIO(
    app,
    cors_allowed_origins="*",
    async_mode="gevent",
    ping_timeout=60,
    ping_interval=25,
    max_http_buffer_size=10 * 1024 * 1024,
    logger=False,
    engineio_logger=False
)

WORKSPACE = '/workspace'
os.makedirs(WORKSPACE, exist_ok=True)

# session_id -> {fd, pid, thread, sid (socket id), cols, rows}
sessions = {}
# socket_id -> [session_ids]  (for cleanup on disconnect)
socket_sessions = {}

sessions_lock = threading.Lock()


def resize_pty(fd, rows, cols):
    try:
        winsize = struct.pack('HHHH', rows, cols, 0, 0)
        fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
    except Exception:
        pass


def read_loop(session_id, fd):
    """Read output from PTY and forward to client via SocketIO."""
    try:
        while True:
            try:
                ready, _, _ = select.select([fd], [], [], 0.04)
            except (ValueError, OSError):
                break
            if ready:
                try:
                    data = os.read(fd, 4096)
                    if not data:
                        break
                    socketio.emit(
                        'terminal_output',
                        {'id': session_id, 'data': data.decode('utf-8', errors='replace')},
                        room=session_id
                    )
                except OSError:
                    break
            else:
                gevent.sleep(0)
    except Exception as e:
        logger.warning(f"Read loop error {session_id}: {e}")
    finally:
        with sessions_lock:
            if session_id in sessions:
                try:
                    os.close(sessions[session_id]['fd'])
                except Exception:
                    pass
                del sessions[session_id]
        socketio.emit('terminal_closed', {'id': session_id}, room=session_id)
        logger.info(f"Session {session_id} ended")


def spawn_shell(session_id, socket_sid, rows=24, cols=80):
    """Fork a bash shell with a PTY."""
    master_fd, slave_fd = pty.openpty()
    resize_pty(master_fd, rows, cols)

    pid = os.fork()
    if pid == 0:
        # Child process
        try:
            os.setsid()
            import fcntl as _fcntl
            _fcntl.ioctl(slave_fd, termios.TIOCSCTTY, 0)
            os.close(master_fd)
            os.dup2(slave_fd, 0)
            os.dup2(slave_fd, 1)
            os.dup2(slave_fd, 2)
            if slave_fd > 2:
                os.close(slave_fd)
            env = os.environ.copy()
            env.update({
                'TERM': 'xterm-256color',
                'COLORTERM': 'truecolor',
                'HOME': WORKSPACE,
                'SHELL': '/bin/bash',
                'PATH': '/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin',
                'LANG': 'en_US.UTF-8',
                'LC_ALL': 'en_US.UTF-8',
            })
            os.chdir(WORKSPACE)
            os.execve('/bin/bash', ['bash', '--login'], env)
        except Exception as e:
            os._exit(1)
    else:
        # Parent process
        os.close(slave_fd)
        with sessions_lock:
            sessions[session_id] = {
                'fd': master_fd,
                'pid': pid,
                'sid': socket_sid,
                'rows': rows,
                'cols': cols,
            }
        t = threading.Thread(target=read_loop, args=(session_id, master_fd), daemon=True)
        t.start()
        sessions[session_id]['thread'] = t
        logger.info(f"Spawned session {session_id} (pid={pid})")


# ─── Socket Events ────────────────────────────────────────────────────────────

@socketio.on('connect')
def on_connect():
    sid = request.sid
    with sessions_lock:
        socket_sessions[sid] = []
    logger.info(f"Client connected: {sid}")


@socketio.on('disconnect')
def on_disconnect():
    sid = request.sid
    with sessions_lock:
        owned = socket_sessions.pop(sid, [])
    for session_id in owned:
        _kill_session(session_id)
    logger.info(f"Client disconnected: {sid}, cleaned {len(owned)} sessions")


def _kill_session(session_id):
    with sessions_lock:
        sess = sessions.pop(session_id, None)
    if sess:
        try:
            os.kill(sess['pid'], signal.SIGHUP)
        except ProcessLookupError:
            pass
        try:
            os.close(sess['fd'])
        except OSError:
            pass


@socketio.on('create_terminal')
def handle_create_terminal(data=None):
    data = data or {}
    sid = request.sid
    session_id = str(uuid.uuid4())
    rows = max(10, int(data.get('rows', 24)))
    cols = max(20, int(data.get('cols', 80)))

    # Join a room for this terminal
    from flask_socketio import join_room
    join_room(session_id)

    with sessions_lock:
        if sid not in socket_sessions:
            socket_sessions[sid] = []
        socket_sessions[sid].append(session_id)

    spawn_shell(session_id, sid, rows, cols)
    emit('terminal_created', {'id': session_id})
    logger.info(f"Terminal {session_id} created for {sid}")


@socketio.on('terminal_input')
def handle_input(data):
    session_id = data.get('id')
    raw = data.get('data', '')
    if not session_id or not raw:
        return
    with sessions_lock:
        sess = sessions.get(session_id)
    if sess:
        try:
            os.write(sess['fd'], raw.encode('utf-8'))
        except OSError:
            pass


@socketio.on('terminal_resize')
def handle_resize(data):
    session_id = data.get('id')
    rows = max(10, int(data.get('rows', 24)))
    cols = max(20, int(data.get('cols', 80)))
    with sessions_lock:
        sess = sessions.get(session_id)
    if sess:
        resize_pty(sess['fd'], rows, cols)
        sess['rows'] = rows
        sess['cols'] = cols


@socketio.on('disconnect_terminal')
def handle_disconnect_terminal(data):
    session_id = data.get('id')
    if session_id:
        sid = request.sid
        with sessions_lock:
            if sid in socket_sessions and session_id in socket_sessions[sid]:
                socket_sessions[sid].remove(session_id)
        _kill_session(session_id)
        logger.info(f"Terminal {session_id} closed by client")


# ─── HTTP Routes ──────────────────────────────────────────────────────────────

@app.route('/')
def index():
    return render_template('index.html')


@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    files = request.files.getlist('file')
    uploaded = []
    for f in files:
        if f.filename:
            safe = Path(f.filename).name
            dest = os.path.join(WORKSPACE, safe)
            f.save(dest)
            uploaded.append(safe)
    return jsonify({'uploaded': uploaded, 'count': len(uploaded)}), 200


@app.route('/download/<path:filename>')
def download_file(filename):
    return send_from_directory(WORKSPACE, filename, as_attachment=True)


@app.route('/list-files')
def list_files():
    try:
        entries = []
        for name in sorted(os.listdir(WORKSPACE)):
            full = os.path.join(WORKSPACE, name)
            stat = os.stat(full)
            entries.append({
                'name': name,
                'size': stat.st_size,
                'is_dir': os.path.isdir(full),
                'modified': int(stat.st_mtime),
            })
        return jsonify({'files': entries}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/delete/<path:filename>', methods=['DELETE'])
def delete_file(filename):
    try:
        full = os.path.join(WORKSPACE, Path(filename).name)
        if os.path.isfile(full):
            os.remove(full)
            return jsonify({'deleted': filename}), 200
        return jsonify({'error': 'Not a file'}), 400
    except Exception as e:
        return jsonify({'error': str(e)}), 500


if __name__ == '__main__':
    socketio.run(app, host='0.0.0.0', port=7860, debug=False, log_output=False)