#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Docker远程控制脚本
功能：
1. 远程控制Docker容器的启动、停止、重启
2. 修改配置文件
3. 查看日志文件
"""

import json
import os
import sys
import argparse
from pathlib import Path
from typing import Optional, Dict, Any, List,Tuple
import paramiko
from io import StringIO


class DockerRemoteController:
    """Docker远程控制器"""
    
    def __init__(self, host: str, port: int = 22, username: str = "root", 
                 password: Optional[str] = None, key_file: Optional[str] = None,
                 project_path: str = "/root/troov-asyncio"):
        """
        初始化远程控制器
        
        Args:
            host: 远程服务器地址
            port: SSH端口
            username: SSH用户名
            password: SSH密码
            key_file: SSH私钥文件路径
            project_path: 项目在远程服务器上的路径
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_file = key_file
        self.project_path = project_path
        self.ssh_client = None
        self.sftp_client = None
        
    def connect(self):
        """建立SSH连接"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            if self.key_file:
                # 使用密钥文件连接
                self.ssh_client.connect(
                    hostname=self.host,
                    port=self.port,
                    username=self.username,
                    key_filename=self.key_file
                )
            else:
                # 使用密码连接
                self.ssh_client.connect(
                    hostname=self.host,
                    port=self.port,
                    username=self.username,
                    password=self.password
                )
            
            # 创建SFTP客户端用于文件传输
            self.sftp_client = self.ssh_client.open_sftp()
            print(f"✓ 成功连接到 {self.host}:{self.port}")
            return True
        except Exception as e:
            print(f"✗ 连接失败: {e}")
            return False
    
    def disconnect(self):
        """关闭连接"""
        if self.sftp_client:
            self.sftp_client.close()
        if self.ssh_client:
            self.ssh_client.close()
    
    def execute_command(self, command: str) -> Tuple[str, str, int]:
        """
        执行远程命令
        
        Returns:
            (stdout, stderr, exit_code)
        """
        if not self.ssh_client:
            raise Exception("未建立SSH连接，请先调用connect()")
        
        stdin, stdout, stderr = self.ssh_client.exec_command(command)
        exit_code = stdout.channel.recv_exit_status()
        stdout_text = stdout.read().decode('utf-8')
        stderr_text = stderr.read().decode('utf-8')
        return stdout_text, stderr_text, exit_code
    
    # ========== Docker控制方法 ==========
    
    def docker_compose_up(self, services: Optional[List[str]] = None) -> bool:
        """启动Docker Compose服务"""
        cmd = f"cd {self.project_path} && docker-compose up -d"
        if services:
            cmd += " " + " ".join(services)
        
        print(f"正在启动Docker服务...")
        stdout, stderr, exit_code = self.execute_command(cmd)
        
        if exit_code == 0:
            print("✓ Docker服务启动成功")
            if stdout:
                print(stdout)
            return True
        else:
            print(f"✗ Docker服务启动失败: {stderr}")
            return False
    
    def docker_compose_down(self, services: Optional[List[str]] = None) -> bool:
        """停止Docker Compose服务"""
        cmd = f"cd {self.project_path} && docker-compose down"
        if services:
            # docker-compose down不支持指定服务，需要单独停止
            for service in services:
                self.docker_stop(service)
            return True
        
        print(f"正在停止Docker服务...")
        stdout, stderr, exit_code = self.execute_command(cmd)
        
        if exit_code == 0:
            print("✓ Docker服务停止成功")
            if stdout:
                print(stdout)
            return True
        else:
            print(f"✗ Docker服务停止失败: {stderr}")
            return False
    
    def docker_restart(self, container_name: str) -> bool:
        """重启指定容器"""
        print(f"正在重启容器: {container_name}")
        stdout, stderr, exit_code = self.execute_command(
            f"docker restart {container_name}"
        )
        
        if exit_code == 0:
            print(f"✓ 容器 {container_name} 重启成功")
            if stdout:
                print(stdout)
            return True
        else:
            print(f"✗ 容器 {container_name} 重启失败: {stderr}")
            return False
    
    def docker_start(self, container_name: str) -> bool:
        """启动指定容器"""
        print(f"正在启动容器: {container_name}")
        stdout, stderr, exit_code = self.execute_command(
            f"docker start {container_name}"
        )
        
        if exit_code == 0:
            print(f"✓ 容器 {container_name} 启动成功")
            if stdout:
                print(stdout)
            return True
        else:
            print(f"✗ 容器 {container_name} 启动失败: {stderr}")
            return False
    
    def docker_stop(self, container_name: str) -> bool:
        """停止指定容器"""
        print(f"正在停止容器: {container_name}")
        stdout, stderr, exit_code = self.execute_command(
            f"docker stop {container_name}"
        )
        
        if exit_code == 0:
            print(f"✓ 容器 {container_name} 停止成功")
            if stdout:
                print(stdout)
            return True
        else:
            print(f"✗ 容器 {container_name} 停止失败: {stderr}")
            return False
    
    def docker_status(self) -> Dict[str, Any]:
        """查看所有容器状态"""
        stdout, stderr, exit_code = self.execute_command(
            "docker ps -a --format '{{.Names}}\t{{.Status}}\t{{.Image}}'"
        )
        
        if exit_code == 0:
            containers = {}
            for line in stdout.strip().split('\n'):
                if line.strip():
                    parts = line.split('\t')
                    if len(parts) >= 2:
                        name = parts[0]
                        status = parts[1]
                        image = parts[2] if len(parts) > 2 else ""
                        containers[name] = {
                            'status': status,
                            'image': image
                        }
            return containers
        else:
            print(f"✗ 获取容器状态失败: {stderr}")
            return {}
    
    def docker_logs(self, container_name: str, lines: int = 100, follow: bool = False) -> str:
        """查看容器日志"""
        cmd = f"docker logs --tail {lines} {container_name}"
        if follow:
            cmd += " -f"
        
        stdout, stderr, exit_code = self.execute_command(cmd)
        
        if exit_code == 0:
            return stdout
        else:
            return f"获取日志失败: {stderr}"
    
    # ========== 配置文件操作方法 ==========
    
    def read_config(self, config_file: str) -> Optional[Dict[str, Any]]:
        """读取远程配置文件"""
        remote_path = f"{self.project_path}/{config_file}"
        
        try:
            # 读取文件内容
            with self.sftp_client.open(remote_path, 'r') as f:
                content = f.read().decode('utf-8')
                return json.loads(content)
        except FileNotFoundError:
            print(f"✗ 配置文件不存在: {remote_path}")
            return None
        except json.JSONDecodeError as e:
            print(f"✗ JSON解析失败: {e}")
            return None
        except Exception as e:
            print(f"✗ 读取配置文件失败: {e}")
            return None
    
    def write_config(self, config_file: str, config_data: Dict[str, Any]) -> bool:
        """写入远程配置文件"""
        remote_path = f"{self.project_path}/{config_file}"
        
        try:
            # 将配置转换为JSON字符串
            json_str = json.dumps(config_data, indent=2, ensure_ascii=False)
            
            # 写入文件
            with self.sftp_client.open(remote_path, 'w') as f:
                f.write(json_str.encode('utf-8'))
            
            print(f"✓ 配置文件已更新: {config_file}")
            return True
        except Exception as e:
            print(f"✗ 写入配置文件失败: {e}")
            return False
    
    def update_config(self, config_file: str, key_path: str, value: Any) -> bool:
        """
        更新配置文件中的指定键值
        
        Args:
            config_file: 配置文件路径（相对于项目根目录）
            key_path: 键路径，支持嵌套，如 "lockV1.sessionLimit" 或 "releaseTimePoint.hour"
            value: 新值
        """
        config = self.read_config(config_file)
        if config is None:
            return False
        
        # 解析键路径
        keys = key_path.split('.')
        current = config
        
        # 导航到目标位置
        for key in keys[:-1]:
            if key not in current:
                current[key] = {}
            current = current[key]
        
        # 设置值
        old_value = current.get(keys[-1])
        current[keys[-1]] = value
        
        print(f"配置更新: {key_path}")
        print(f"  旧值: {old_value}")
        print(f"  新值: {value}")
        
        return self.write_config(config_file, config)
    
    # ========== 日志文件查看方法 ==========
    
    def view_log(self, log_file: str, lines: int = 100, follow: bool = False) -> str:
        """
        查看日志文件
        
        Args:
            log_file: 日志文件路径（相对于项目根目录或绝对路径）
            lines: 显示的行数
            follow: 是否实时跟踪（注意：follow模式在SSH中可能不稳定）
        """
        # 如果是相对路径，添加项目路径前缀
        if not log_file.startswith('/'):
            remote_path = f"{self.project_path}/{log_file}"
        else:
            remote_path = log_file
        
        if follow:
            # follow模式需要特殊处理
            cmd = f"tail -f {remote_path}"
        else:
            cmd = f"tail -n {lines} {remote_path}"
        
        stdout, stderr, exit_code = self.execute_command(cmd)
        
        if exit_code == 0:
            return stdout
        else:
            return f"查看日志失败: {stderr}"
    
    def read_log(
        self,
        log_file: str,
        lines: int = 100,
        from_head: bool = False,
        full: bool = False,
        output_path: Optional[str] = None,
    ) -> str:
        """
        远程读取 logs 目录下指定的日志文件。
        若 log_file 仅为文件名（如 slot_monitor.log），自动视为 logs/ 下文件。
        
        Args:
            log_file: 日志文件名或路径。纯文件名则自动加 logs/ 前缀。
            lines: 读取行数（默认 100）。full 时忽略。
            from_head: True 则读开头 N 行，否则读末尾 N 行。
            full: True 则读取整个文件。
            output_path: 若指定，将读取内容保存到本地该路径。
        
        Returns:
            读取到的日志内容；失败时返回错误信息。
        """
        # 纯文件名（无路径）视为 logs 下
        if not log_file.startswith('/') and '/' not in log_file:
            remote_path = f"{self.project_path}/logs/{log_file}"
        elif not log_file.startswith('/'):
            remote_path = f"{self.project_path}/{log_file}"
        else:
            remote_path = log_file
        
        if full:
            cmd = f"cat {remote_path}"
        elif from_head:
            cmd = f"head -n {lines} {remote_path}"
        else:
            cmd = f"tail -n {lines} {remote_path}"
        
        stdout, stderr, exit_code = self.execute_command(cmd)
        
        if exit_code != 0:
            return f"读取日志失败: {stderr}"
        
        if output_path:
            try:
                path = Path(output_path)
                path.parent.mkdir(parents=True, exist_ok=True)
                path.write_text(stdout, encoding='utf-8')
                print(f"✓ 已保存到本地: {output_path}")
            except Exception as e:
                return f"保存到本地失败: {e}\n\n{stdout}"
        
        return stdout
    
    def list_logs(self) -> List[str]:
        """列出所有日志文件"""
        logs_dir = f"{self.project_path}/logs"
        stdout, stderr, exit_code = self.execute_command(
            f"ls -1 {logs_dir}/*.log 2>/dev/null || echo ''"
        )
        
        if exit_code == 0 and stdout.strip():
            log_files = [os.path.basename(f) for f in stdout.strip().split('\n') if f.strip()]
            return log_files
        return []


def main():
    parser = argparse.ArgumentParser(
        description='Docker远程控制脚本',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 连接并查看容器状态
  python docker_remote_control.py --host 192.168.1.100 --password yourpass status

  # 重启容器
  python docker_remote_control.py --host 192.168.1.100 --password yourpass restart troov_monitor

  # 启动所有服务
  python docker_remote_control.py --host 192.168.1.100 --password yourpass up

  # 查看日志
  python docker_remote_control.py --host 192.168.1.100 --password yourpass logs troov_monitor --lines 50

  # 读取配置文件
  python docker_remote_control.py --host 192.168.1.100 --password yourpass read-config config/troov_config.json

  # 更新配置文件
  python docker_remote_control.py --host 192.168.1.100 --password yourpass update-config config/troov_config.json lockV1.sessionLimit 10

  # 查看日志文件
  python docker_remote_control.py --host 192.168.1.100 --password yourpass view-log logs/slot_monitor.log --lines 100

  # 远程读取 logs 下指定日志（仅文件名自动加 logs/ 前缀）
  python docker_remote_control.py --host 192.168.1.100 --password yourpass read-log slot_monitor.log --lines 200
  python docker_remote_control.py --host 192.168.1.100 --password yourpass read-log slot_monitor.log --full --output ./downloaded.log
  python docker_remote_control.py --host 192.168.1.100 --password yourpass read-log session_creator.log --head --lines 50
        """
    )
    
    # 连接参数
    parser.add_argument('--host', required=True, help='远程服务器地址')
    parser.add_argument('--port', type=int, default=22, help='SSH端口 (默认: 22)')
    parser.add_argument('--username', default='root', help='SSH用户名 (默认: root)')
    parser.add_argument('--password', help='SSH密码')
    parser.add_argument('--key-file', help='SSH私钥文件路径')
    parser.add_argument('--project-path', default='/root/troov-asyncio', 
                       help='项目在远程服务器上的路径 (默认: /root/troov-asyncio)')
    
    # 子命令
    subparsers = parser.add_subparsers(dest='command', help='可用命令')
    
    # status命令
    subparsers.add_parser('status', help='查看所有容器状态')
    
    # docker控制命令
    up_parser = subparsers.add_parser('up', help='启动Docker Compose服务')
    up_parser.add_argument('services', nargs='*', help='要启动的服务名称（可选）')
    
    down_parser = subparsers.add_parser('down', help='停止Docker Compose服务')
    down_parser.add_argument('services', nargs='*', help='要停止的服务名称（可选）')
    
    restart_parser = subparsers.add_parser('restart', help='重启容器')
    restart_parser.add_argument('container', help='容器名称')
    
    start_parser = subparsers.add_parser('start', help='启动容器')
    start_parser.add_argument('container', help='容器名称')
    
    stop_parser = subparsers.add_parser('stop', help='停止容器')
    stop_parser.add_argument('container', help='容器名称')
    
    logs_parser = subparsers.add_parser('logs', help='查看容器日志')
    logs_parser.add_argument('container', help='容器名称')
    logs_parser.add_argument('--lines', type=int, default=100, help='显示行数 (默认: 100)')
    logs_parser.add_argument('--follow', action='store_true', help='实时跟踪日志')
    
    # 配置文件命令
    read_config_parser = subparsers.add_parser('read-config', help='读取配置文件')
    read_config_parser.add_argument('config_file', help='配置文件路径')
    
    update_config_parser = subparsers.add_parser('update-config', help='更新配置文件')
    update_config_parser.add_argument('config_file', help='配置文件路径')
    update_config_parser.add_argument('key_path', help='键路径，如 lockV1.sessionLimit')
    update_config_parser.add_argument('value', help='新值')
    
    # 日志文件命令
    view_log_parser = subparsers.add_parser('view-log', help='查看日志文件')
    view_log_parser.add_argument('log_file', help='日志文件路径')
    view_log_parser.add_argument('--lines', type=int, default=100, help='显示行数 (默认: 100)')
    view_log_parser.add_argument('--follow', action='store_true', help='实时跟踪日志')
    
    read_log_parser = subparsers.add_parser(
        'read-log',
        help='远程读取 logs 下指定日志文件；仅文件名时自动加 logs/ 前缀',
    )
    read_log_parser.add_argument(
        'log_file',
        help='日志文件名（如 slot_monitor.log）或路径（如 logs/slot_monitor.log）',
    )
    read_log_parser.add_argument('--lines', type=int, default=100, help='行数 (默认: 100)')
    read_log_parser.add_argument(
        '--head',
        action='store_true',
        help='读开头 N 行；默认读末尾 N 行',
    )
    read_log_parser.add_argument('--full', action='store_true', help='读取整个文件')
    read_log_parser.add_argument(
        '--output',
        dest='output_path',
        metavar='PATH',
        help='将读取内容保存到本地文件',
    )
    
    list_logs_parser = subparsers.add_parser('list-logs', help='列出所有日志文件')
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        sys.exit(1)
    
    # 检查认证方式
    if not args.password and not args.key_file:
        print("错误: 必须提供 --password 或 --key-file")
        sys.exit(1)
    
    # 创建控制器并连接
    controller = DockerRemoteController(
        host=args.host,
        port=args.port,
        username=args.username,
        password=args.password,
        key_file=args.key_file,
        project_path=args.project_path
    )
    
    if not controller.connect():
        sys.exit(1)
    
    try:
        # 执行命令
        if args.command == 'status':
            containers = controller.docker_status()
            print("\n容器状态:")
            print("-" * 80)
            for name, info in containers.items():
                print(f"{name:30} {info['status']:30} {info['image']}")
        
        elif args.command == 'up':
            controller.docker_compose_up(args.services if args.services else None)
        
        elif args.command == 'down':
            controller.docker_compose_down(args.services if args.services else None)
        
        elif args.command == 'restart':
            controller.docker_restart(args.container)
        
        elif args.command == 'start':
            controller.docker_start(args.container)
        
        elif args.command == 'stop':
            controller.docker_stop(args.container)
        
        elif args.command == 'logs':
            output = controller.docker_logs(args.container, args.lines, args.follow)
            print(output)
        
        elif args.command == 'read-config':
            config = controller.read_config(args.config_file)
            if config:
                print(json.dumps(config, indent=2, ensure_ascii=False))
        
        elif args.command == 'update-config':
            # 尝试将值转换为合适的类型
            value = args.value
            try:
                # 尝试转换为数字
                if value.isdigit() or (value.startswith('-') and value[1:].isdigit()):
                    value = int(value)
                elif '.' in value and value.replace('.', '').replace('-', '').isdigit():
                    value = float(value)
                # 尝试转换为布尔值
                elif value.lower() in ('true', 'false'):
                    value = value.lower() == 'true'
                # 尝试转换为JSON
                elif value.startswith('{') or value.startswith('['):
                    value = json.loads(value)
            except:
                pass  # 保持为字符串
            
            controller.update_config(args.config_file, args.key_path, value)
        
        elif args.command == 'view-log':
            output = controller.view_log(args.log_file, args.lines, args.follow)
            print(output)
        
        elif args.command == 'read-log':
            output = controller.read_log(
                args.log_file,
                lines=args.lines,
                from_head=args.head,
                full=args.full,
                output_path=args.output_path,
            )
            print(output)
        
        elif args.command == 'list-logs':
            log_files = controller.list_logs()
            print("日志文件列表:")
            for log_file in log_files:
                print(f"  - {log_file}")
    
    finally:
        controller.disconnect()


if __name__ == '__main__':
    main()