import socketio
import threading
import os
import select
import uuid
import time
import sys

# 检查并导入 pty 模块（Unix/Linux/MacOS 专用）
try:
    import pty
    import termios
    import tty
    PTY_AVAILABLE = True
except ImportError:
    print("警告: pty 模块不可用，这可能是 Windows 系统")
    PTY_AVAILABLE = False

class TerminalClient:
    def __init__(self, server_url, client_name=None):
        self.sio = socketio.Client()
        self.server_url = server_url
        self.client_id = str(uuid.uuid4())
        self.client_name = client_name or f"Client-{self.get_hostname()}-{self.client_id[:8]}"
        self.master_fd = None
        self.running = False
        self.connected = False
        self.child_pid = None

        self.setup_handlers()

    def get_hostname(self):
        """获取主机名"""
        try:
            import socket
            return socket.gethostname()
        except:
            return "unknown-host"

    def find_available_shell(self):
        """查找可用的 shell"""
        shells = [
            '/bin/bash',
            '/usr/bin/bash',
            '/bin/zsh',
            '/usr/bin/zsh',
            '/bin/sh',
            '/usr/bin/sh',
            '/bin/ksh',
            '/usr/bin/ksh',
            '/bin/tcsh',
            '/usr/bin/tcsh'
        ]

        for shell in shells:
            if os.path.exists(shell):
                print(f"找到可用 shell: {shell}")
                return shell

        # 如果找不到任何标准 shell，尝试使用环境变量中的 SHELL
        fallback_shell = os.environ.get('SHELL')
        if fallback_shell and os.path.exists(fallback_shell):
            print(f"使用环境变量中的 shell: {fallback_shell}")
            return fallback_shell

        print("警告: 未找到可用的 shell，将尝试使用 /bin/sh")
        return '/bin/sh'  # 最后尝试

    def setup_handlers(self):
        @self.sio.event
        def connect():
            print(f"✓ 连接到服务器: {self.server_url}")
            self.connected = True
            # 立即注册客户端
            self.sio.emit('client_register', {
                'client_id': self.client_id,
                'name': self.client_name
            })
            print(f"✓ 客户端注册信息已发送: {self.client_name} (ID: {self.client_id})")

        @self.sio.event
        def disconnect():
            print("✗ 与服务器断开连接")
            self.connected = False
            self.stop_terminal()

        @self.sio.on('registration_success')
        def handle_registration_success(data):
            print("✓ 客户端注册成功")

        @self.sio.on('terminal_command')
        def handle_terminal_command(data):
            """处理来自服务器的终端命令"""
            target_client_id = data.get('client_id')
            command_data = data.get('command', '')

            # 确保命令是发送给这个客户端的
            if target_client_id != self.client_id:
                return

            print(f"收到终端输入: {repr(command_data)}")

            if self.master_fd and self.running:
                try:
                    # 将数据写入伪终端
                    if isinstance(command_data, str):
                        command_data = command_data.encode('utf-8')
                    written = os.write(self.master_fd, command_data)
                    print(f"成功写入 {written} 字节到伪终端")
                except Exception as e:
                    print(f"写入伪终端失败: {e}")
            else:
                print("伪终端未就绪，无法写入数据")

    def start_terminal(self):
        """启动伪终端 - 修复 shell 不存在的问题"""
        try:
            print("正在启动终端...")

            if not PTY_AVAILABLE:
                print("✗ pty 模块不可用，无法创建伪终端")
                print("请确保在 Unix/Linux/MacOS 系统上运行此客户端")
                return False

            # 查找可用的 shell
            shell_path = self.find_available_shell()
            shell_name = os.path.basename(shell_path)

            print(f"使用 shell: {shell_path}")

            # 创建伪终端
            self.master_fd, slave_fd = pty.openpty()

            # 设置终端属性
            try:
                # 保存原始属性
                old_settings = termios.tcgetattr(slave_fd)
                # 设置原始模式
                tty.setraw(slave_fd)
            except Exception as e:
                print(f"设置终端模式时警告: {e}")

            # 启动shell进程
            self.child_pid = os.fork()

            if self.child_pid == 0:  # 子进程
                try:
                    # 关闭主终端文件描述符
                    os.close(self.master_fd)

                    # 创建新的会话
                    os.setsid()

                    # 将伪终端设置为控制终端
                    os.dup2(slave_fd, 0)   # stdin
                    os.dup2(slave_fd, 1)   # stdout
                    os.dup2(slave_fd, 2)   # stderr

                    # 关闭从终端
                    os.close(slave_fd)

                    # 设置环境变量
                    env = os.environ.copy()
                    env['TERM'] = 'xterm-256color'
                    env['COLORTERM'] = 'truecolor'

                    # 为不同 shell 设置不同的提示符
                    if 'bash' in shell_name:
                        env['PS1'] = '\\[\\033[1;32m\\]\\u@\\h:\\w\\$\\[\\033[0m\\] '
                    elif 'zsh' in shell_name:
                        env['PROMPT'] = '%F{green}%n@%m:%~%#%f '

                    # 设置工作目录
                    os.chdir(os.path.expanduser("~"))

                    print(f"正在启动 {shell_path}...")

                    # 启动交互式shell - 关键修复：处理启动失败
                    if shell_name in ['bash', 'zsh', 'ksh']:
                        os.execve(shell_path, [shell_name, '-i'], env)
                    else:
                        os.execve(shell_path, [shell_name], env)

                except Exception as e:
                    # 如果 execve 失败，打印错误并退出子进程
                    print(f"启动 shell 失败: {e}")
                    sys.stderr.write(f"无法启动 shell: {e}\n")
                    sys.exit(1)

            else:  # 父进程
                # 关闭从终端
                os.close(slave_fd)

                # 等待子进程状态，检查是否成功启动
                time.sleep(0.5)  # 给子进程一点时间启动

                # 检查子进程是否还在运行
                try:
                    pid, status = os.waitpid(self.child_pid, os.WNOHANG)
                    if pid != 0:
                        # 子进程已经退出，说明启动失败
                        print(f"✗ 终端启动失败，子进程退出状态: {status}")
                        return False
                except ChildProcessError:
                    # 子进程不存在
                    print("✗ 子进程不存在，终端启动失败")
                    return False

                self.running = True

                # 启动线程读取终端输出
                output_thread = threading.Thread(target=self.read_terminal_output)
                output_thread.daemon = True
                output_thread.start()

                print("✓ 终端启动成功")
                return True

        except Exception as e:
            print(f"✗ 启动终端失败: {e}")
            import traceback
            traceback.print_exc()
            return False

    def start_terminal_simple(self):
        """简单的终端启动方案，不依赖 fork"""
        try:
            print("使用简单终端方案...")

            if not PTY_AVAILABLE:
                return self.start_terminal_windows()

            # 创建伪终端
            self.master_fd, slave_fd = pty.openpty()

            # 使用 subprocess 启动 shell，避免 fork 问题
            import subprocess

            shell_path = self.find_available_shell()

            # 启动进程
            self.process = subprocess.Popen(
                [shell_path, '-i'],
                stdin=slave_fd,
                stdout=slave_fd,
                stderr=slave_fd,
                preexec_fn=os.setsid,
                start_new_session=True
            )

            # 关闭从终端（subprocess 会管理它）
            os.close(slave_fd)

            self.running = True

            # 启动线程读取终端输出
            output_thread = threading.Thread(target=self.read_terminal_output)
            output_thread.daemon = True
            output_thread.start()

            print("✓ 简单终端启动成功")
            return True

        except Exception as e:
            print(f"✗ 简单终端启动失败: {e}")
            return False

    def start_terminal_windows(self):
        """Windows 系统的备用终端实现"""
        try:
            print("Windows 系统使用备用终端实现...")

            import subprocess

            # 启动一个子进程
            self.process = subprocess.Popen(
                ['cmd.exe'] if os.name == 'nt' else ['/bin/sh', '-i'],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=False,  # 使用二进制模式
                bufsize=0
            )

            self.running = True

            # 启动线程读取输出
            output_thread = threading.Thread(target=self.read_process_output)
            output_thread.daemon = True
            output_thread.start()

            # 启动线程处理输入
            input_thread = threading.Thread(target=self.handle_process_input)
            input_thread.daemon = True
            input_thread.start()

            print("✓ 备用终端启动成功")
            return True

        except Exception as e:
            print(f"✗ 启动备用终端失败: {e}")
            return False

    def handle_process_input(self):
        """处理进程输入 (Windows 备用方案)"""
        while self.running and hasattr(self, 'process') and self.process.poll() is None:
            time.sleep(0.1)  # 这个线程主要是为了保持进程运行

    def read_process_output(self):
        """读取进程输出 (Windows 备用方案)"""
        print("开始读取进程输出...")

        while self.running and hasattr(self, 'process') and self.process.poll() is None:
            try:
                output = self.process.stdout.read1(1024)  # 读取可用数据
                if output and self.connected:
                    # 解码输出
                    try:
                        output_str = output.decode('utf-8', errors='replace')
                    except:
                        output_str = output.decode('latin-1', errors='replace')

                    self.sio.emit('terminal_data', {
                        'client_id': self.client_id,
                        'data': output_str
                    })
                    print(f"发送输出: {repr(output_str[:100])}")
            except Exception as e:
                if self.running:  # 只在运行状态下打印错误
                    print(f"读取进程输出时出错: {e}")
                break

        print("进程输出读取线程结束")

    def read_terminal_output(self):
        """读取终端输出并发送到服务器"""
        print("开始读取终端输出...")

        while self.running and self.master_fd:
            try:
                # 使用select检查是否有数据可读
                r, w, e = select.select([self.master_fd], [], [], 0.1)

                if self.master_fd in r:
                    # 读取数据
                    data = os.read(self.master_fd, 1024)
                    if data:
                        output = data.decode('utf-8', errors='replace')

                        if self.connected:
                            # 发送终端输出到服务器
                            self.sio.emit('terminal_data', {
                                'client_id': self.client_id,
                                'data': output
                            })
                            print(f"发送终端输出: {repr(output[:100])}...")

                # 短暂休眠避免CPU占用过高
                time.sleep(0.01)

            except (OSError, ValueError, Exception) as e:
                if self.running:  # 只在运行状态下打印错误
                    print(f"读取终端输出时出错: {e}")
                break

        print("终端输出读取线程结束")

    def stop_terminal(self):
        """停止终端"""
        print("正在停止终端...")
        self.running = False

        # 停止 subprocess (如果使用)
        if hasattr(self, 'process') and self.process:
            try:
                self.process.terminate()
                self.process.wait(timeout=5)
            except:
                try:
                    self.process.kill()
                except:
                    pass
            self.process = None

        # 关闭伪终端
        if self.master_fd:
            try:
                # 发送退出命令
                os.write(self.master_fd, b'exit\n')
                time.sleep(0.2)
                os.close(self.master_fd)
            except Exception as e:
                print(f"关闭伪终端时出错: {e}")
            finally:
                self.master_fd = None

        # 终止子进程 (如果使用 fork)
        if self.child_pid:
            try:
                os.kill(self.child_pid, 9)
            except:
                pass
            self.child_pid = None

        print("终端已停止")

    def connect_to_server(self):
        """连接到服务器"""
        try:
            print(f"正在连接到服务器: {self.server_url}")
            self.sio.connect(self.server_url)
            return True
        except Exception as e:
            print(f"连接服务器失败: {e}")
            return False

    def run(self):
        """运行客户端"""
        if self.connect_to_server():
            # 根据平台选择合适的终端启动方式
            terminal_started = False

            if PTY_AVAILABLE:
                # 首先尝试标准方法
                terminal_started = self.start_terminal()

                # 如果标准方法失败，尝试简单方法
                if not terminal_started:
                    print("标准终端启动失败，尝试简单方案...")
                    terminal_started = self.start_terminal_simple()
            else:
                terminal_started = self.start_terminal_windows()

            if terminal_started:
                try:
                    # 保持连接
                    print("客户端运行中，按 Ctrl+C 退出...")
                    self.sio.wait()
                except KeyboardInterrupt:
                    print("\n收到中断信号，正在关闭客户端...")
                finally:
                    self.stop_terminal()
                    self.sio.disconnect()
            else:
                print("无法启动终端，客户端退出")
                self.sio.disconnect()
        else:
            print("无法连接到服务器，客户端退出")

def start():
    import argparse

    parser = argparse.ArgumentParser(description='终端客户端')
    parser.add_argument('--name', type=str, help='客户端名称', default=os.getenv("NAME", '客户端名称'))
    parser.add_argument('--server', type=str, help='服务器地址',
                        default=os.getenv("SERVER_URL", 'http://localhost:5000'))

    args = parser.parse_args()

    # 检查系统兼容性
    if not PTY_AVAILABLE and os.name != 'nt':
        print("警告: 当前系统可能不支持 pty，终端功能可能受限")

    client = TerminalClient(
        server_url=args.server,
        client_name=args.name
    )

    try:
        client.run()
    except Exception as e:
        print(f"客户端运行错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == '__main__':
    start()