#!/usr/bin/env python3
"""
设备初始化自动化脚本
根据 init_device.md 文档自动执行设备初始化步骤
"""

import subprocess
import sys
import os
import argparse
import tempfile
import re
from pathlib import Path
from typing import Optional, List
from dataclasses import dataclass
from contextlib import contextmanager

@dataclass
class InitConfig:
    """初始化配置"""
    install_chromium: bool = False
    username: Optional[str] = None
    dry_run: bool = False

class Logger:
    """优雅的日志输出"""
    
    @staticmethod
    def section(title: str):
        """输出章节标题"""
        print(f"\n{'='*60}")
        print(f"  {title}")
        print(f"{'='*60}")
    
    @staticmethod
    def step(description: str):
        """输出步骤说明"""
        print(f"\n🔧 {description}")
    
    @staticmethod
    def success(message: str):
        """输出成功信息"""
        print(f"✅ {message}")
    
    @staticmethod
    def warning(message: str):
        """输出警告信息"""
        print(f"⚠️  {message}")
    
    @staticmethod
    def error(message: str):
        """输出错误信息"""
        print(f"❌ {message}")
    
    @staticmethod
    def info(message: str):
        """输出信息"""
        print(f"ℹ️  {message}")

class CommandRunner:
    """命令执行器"""
    
    def __init__(self, dry_run: bool = False):
        self.dry_run = dry_run
    
    def _execute_command(self, command: str) -> subprocess.CompletedProcess:
        """执行命令的核心逻辑"""
        with subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1
        ) as process:
            output_lines = []
            
            # 实时读取并显示输出
            for line in iter(process.stdout.readline, ''):
                print(line.rstrip())
                output_lines.append(line)
            
            # 等待进程完成
            process.wait()
            
            return subprocess.CompletedProcess(
                command, 
                process.returncode, 
                ''.join(output_lines), 
                ""
            )
    
    def run(self, command: str, description: str = "", check: bool = True) -> subprocess.CompletedProcess:
        """执行命令"""
        Logger.step(description or command)
        
        if self.dry_run:
            Logger.info(f"[预览模式] 命令: {command}")
            return subprocess.CompletedProcess(command, 0, "", "")
        
        try:
            result = self._execute_command(command)
            
            if check and result.returncode != 0:
                raise subprocess.CalledProcessError(result.returncode, command, result.stdout)
            
            Logger.success("执行成功")
            return result
            
        except subprocess.CalledProcessError as e:
            Logger.error(f"命令执行失败 (退出码: {e.returncode})")
            if e.stderr:
                print(f"错误详情: {e.stderr.strip()}")
            if check:
                sys.exit(1)
        except Exception as e:
            Logger.error(f"执行异常: {e}")
            if check:
                sys.exit(1)

@contextmanager
def temp_file_with_content(content: str, suffix: str = ""):
    """创建临时文件上下文管理器"""
    with tempfile.NamedTemporaryFile(mode='w', suffix=suffix, delete=False) as f:
        f.write(content)
        temp_path = f.name
    
    try:
        yield temp_path
    finally:
        Path(temp_path).unlink(missing_ok=True)

class DeviceInitializer:
    """设备初始化器"""
    
    # 常量定义
    SECTION_OPTIMIZE_BOOT = "优化系统启动时间"
    SECTION_UPDATE_SYSTEM = "更新系统包列表"
    SECTION_INSTALL_BASIC = "安装基础依赖"
    SECTION_ADD_REPOS = "添加第三方软件源"
    SECTION_INSTALL_PACKAGES = "安装主要软件包"
    SECTION_CONFIGURE_SYSTEM = "配置系统设置"
    SECTION_AUTO_STARTX = "配置图形界面自动启动"
    SECTION_AUTO_LOGIN = "配置自动登录"
    SECTION_CONFIGURE_I3 = "配置 i3 窗口管理器"
    SECTION_MANUAL_STEPS = "手动配置步骤"
    SECTION_COMPLETE = "初始化完成"
    
    def __init__(self, config: InitConfig):
        self.config = config
        # 获取实际用户名（sudo 下运行时从 SUDO_USER 获取）
        self.username = config.username or os.getenv('SUDO_USER') or os.getenv('USER')
        self.user_home = Path(f"/home/{self.username}")
        self.runner = CommandRunner(config.dry_run)
        
        # 软件包配置
        self.base_packages = [
            "fonts-noto-cjk",
            "fonts-noto-color-emoji", 
            "unclutter",
            "xorg",
            "i3-wm",
            "libvlc-dev",
            "vlc",
            "libasound2-dev",
            "alsa-utils",
            "trzsz",
            "wireguard",
            "wireguard-tools"
        ]
        
        self.repositories = [
            "ppa:xtradeb/apps",
            "ppa:trzsz/ppa"
        ]
    
    def _edit_systemd_service(self, service_name: str, override_content: str, description: str):
        """兼容地编辑 systemd 服务配置"""
        # 直接创建 override 目录和文件，兼容所有 systemctl 版本
        override_dir = f"/etc/systemd/system/{service_name}.d"
        override_file = f"{override_dir}/override.conf"
        
        # 创建目录
        self.runner.run(f"sudo mkdir -p {override_dir}", f"创建服务 override 目录: {override_dir}")
        
        # 写入配置文件
        with temp_file_with_content(override_content) as temp_file:
            self.runner.run(f"sudo cp {temp_file} {override_file}", f"创建服务配置文件: {override_file}")
        
        # 重新加载 systemd 配置
        self.runner.run("sudo systemctl daemon-reload", "重新加载 systemd 配置")
        
        Logger.success(f"{description} - 已完成")
    
    def optimize_boot_time(self):
        """优化 Ubuntu 24 开机时间"""
        Logger.section(self.SECTION_OPTIMIZE_BOOT)
        
        override_content = "[Service]\nTimeoutStartSec=2sec\n"
        self._edit_systemd_service(
            "systemd-networkd-wait-online.service",
            override_content,
            "配置网络等待服务超时时间为2秒"
        )
    
    def update_system(self):
        """更新系统包"""
        Logger.section(self.SECTION_UPDATE_SYSTEM)
        self.runner.run("sudo apt-get update", "刷新软件包索引")
    
    def install_basic_packages(self):
        """安装基础包"""
        Logger.section(self.SECTION_INSTALL_BASIC)
        basic_packages = ["curl", "gpg"]
        packages_str = " ".join(basic_packages)
        self.runner.run(f"sudo apt-get install -y {packages_str}", f"安装基础包: {', '.join(basic_packages)}")
    
    def add_repositories(self):
        """添加软件源"""
        Logger.section(self.SECTION_ADD_REPOS)
        
        for repo in self.repositories:
            self.runner.run(f"sudo add-apt-repository -y {repo}", f"添加软件源: {repo}")
        
        self.runner.run("sudo apt-get update", "更新软件包索引")
    
    def install_packages(self):
        """安装主要软件包"""
        Logger.section(self.SECTION_INSTALL_PACKAGES)
        
        packages = self.base_packages.copy()
        if self.config.install_chromium:
            packages.insert(0, "ungoogled-chromium")
            Logger.info("已包含 ungoogled-chromium")
        
        packages_str = " ".join(packages)
        package_list = ", ".join(packages)
        # 安装软件包，实时显示进度
        self.runner.run(f"sudo apt-get install -y {packages_str}", f"安装软件包: {package_list}")
    
    def configure_system(self):
        """配置系统设置"""
        Logger.section(self.SECTION_CONFIGURE_SYSTEM)
        
        # 设置时区
        self.runner.run("sudo timedatectl set-timezone Asia/Shanghai", "设置时区为上海")
        
        # 添加用户到相关组
        groups = "audio,video,dialout"
        self.runner.run(
            f"sudo usermod -aG {groups} {self.username}", 
            f"将用户 {self.username} 添加到 {groups} 组"
        )
    
    def setup_auto_startx(self):
        """设置自动启动 Xorg"""
        Logger.section(self.SECTION_AUTO_STARTX)
        
        bashrc_path = self.user_home / ".bashrc"
        startx_config = """
# 自动启动 Xorg 和窗口管理器
if [ -z "$DISPLAY" ] && [ "$(tty)" = "/dev/tty1" ]; then
    startx
fi
"""
        
        if self.config.dry_run:
            Logger.step("配置自动启动 Xorg")
            Logger.info(f"[预览模式] 将添加配置到 {bashrc_path}")
            return
        
        # 检查是否已存在配置
        if bashrc_path.exists():
            try:
                content = bashrc_path.read_text(encoding='utf-8')
                if 'startx' in content and 'tty1' in content:
                    Logger.warning("自动启动 Xorg 配置已存在，跳过")
                    return
            except Exception as e:
                Logger.warning(f"读取 .bashrc 文件失败: {e}")
        
        # 添加配置
        try:
            with open(bashrc_path, 'a', encoding='utf-8') as f:
                f.write(startx_config)
            
            # 设置文件所有权为实际用户
            self.runner.run(f"sudo chown {self.username}:{self.username} {bashrc_path}", 
                          f"设置 .bashrc 所有权为 {self.username}")
            
            Logger.success(f"已配置自动启动 Xorg: {bashrc_path}")
        except Exception as e:
            Logger.error(f"配置自动启动 Xorg 失败: {e}")
    
    def setup_auto_login(self):
        """设置自动登录"""
        Logger.section(self.SECTION_AUTO_LOGIN)
        
        override_content = f"""[Service]
ExecStart=
ExecStart=-/sbin/agetty --autologin {self.username} --noclear %I $TERM
"""
        
        self._edit_systemd_service(
            "getty@tty1.service",
            override_content,
            f"配置用户 {self.username} 自动登录"
        )
    
    def configure_i3(self):
        """配置 i3 窗口管理器"""
        Logger.section(self.SECTION_CONFIGURE_I3)
        
        i3_config_dir = self.user_home / ".config" / "i3"
        i3_config_path = i3_config_dir / "config"
        
        if self.config.dry_run:
            Logger.step("配置 i3 窗口管理器")
            Logger.info(f"[预览模式] 将配置 i3 配置文件: {i3_config_path}")
            return
        
        # 创建 i3 配置目录
        i3_config_dir.mkdir(parents=True, exist_ok=True)
        
        # 检查是否已有 i3 配置文件
        if not i3_config_path.exists():
            Logger.info("i3 配置文件不存在，从系统默认配置复制")
            self._copy_system_i3_config(i3_config_path)
        else:
            Logger.info("i3 配置文件已存在，直接更新配置")
        
        # 无论是新生成还是已存在，都更新配置
        self._update_i3_config(i3_config_path)
        
        # 设置配置文件所有权为实际用户
        self.runner.run(f"sudo chown -R {self.username}:{self.username} {i3_config_dir}", f"设置 i3 配置目录所有权为 {self.username}")
        
        Logger.success(f"已配置 i3 窗口管理器: {i3_config_path}")
    
    def _copy_system_i3_config(self, config_path: Path):
        """从系统默认配置复制 i3 配置文件"""
        Logger.step("从系统默认配置复制 i3 配置文件")
        
        system_config = "/etc/i3/config"
        self.runner.run(f"cp {system_config} {config_path}", f"复制系统默认 i3 配置到 {config_path}")
        
        # 移除 i3-config-wizard 指令
        try:
            content = config_path.read_text(encoding='utf-8')
            if 'i3-config-wizard' in content:
                # 移除包含 i3-config-wizard 的行
                lines = content.splitlines()
                filtered_lines = [line for line in lines if 'i3-config-wizard' not in line]
                
                with open(config_path, 'w', encoding='utf-8') as f:
                    f.write('\n'.join(filtered_lines) + '\n')
                
                Logger.info("已移除系统配置中的 i3-config-wizard 指令")
            
        except Exception as e:
            Logger.warning(f"移除 i3-config-wizard 指令失败: {e}")
        
        Logger.success("已从系统默认配置复制 i3 配置文件")
    
    def _update_i3_config(self, config_path: Path):
        """更新现有 i3 配置文件"""
        Logger.step("更新现有 i3 配置文件")
        
        try:
            content = config_path.read_text(encoding='utf-8')
        except Exception as e:
            Logger.error(f"读取 i3 配置文件失败: {e}")
            return
        
        # 添加启动配置项
        startup_configs = [
            "exec --no-startup-id unclutter -root # 隐藏鼠标",
            "exec --no-startup-id xset dpms 0 0 0 # 关闭屏幕自动关闭", 
            "exec --no-startup-id xset s off # 关闭屏幕保护"
        ]
        
        config_modified = False
        for config_line in startup_configs:
            if config_line not in content:
                content += f"\n{config_line}\n"
                Logger.info(f"添加配置: {config_line}")
                config_modified = True
        
        # 移除 i3bar 配置块
        if "bar {" in content:
            # 使用正则表达式直接移除整个 bar 配置块
            bar_pattern = r'bar\s*\{[^}]*\}'
            new_content = re.sub(bar_pattern, '', content, flags=re.DOTALL)
            if new_content != content:
                content = new_content
                Logger.info("已移除 i3bar 状态栏配置")
                config_modified = True
        
        # 只在有修改时才写回文件
        if config_modified:
            try:
                with open(config_path, 'w', encoding='utf-8') as f:
                    f.write(content)
                Logger.success("i3 配置文件更新完成")
            except Exception as e:
                Logger.error(f"写入 i3 配置文件失败: {e}")
        else:
            Logger.info("i3 配置无需修改")
    
    def show_manual_steps(self):
        """显示需要手动完成的步骤"""
        Logger.section(self.SECTION_MANUAL_STEPS)
        
        Logger.info("请手动完成以下配置步骤:")
        
        print("\n📋 WireGuard 配置:")
        print("   1. 从服务器获取 WireGuard 配置文件")
        print("   2. 保存到 /etc/wireguard/wg0.conf")
        print("   3. 启用并启动 WireGuard:")
        print("      sudo systemctl enable wg-quick@wg0")
        print("      sudo systemctl start wg-quick@wg0")
        
        print("\n📋 系统音量配置:")
        print("   1. 运行 alsamixer 进入音量控制界面")
        print("   2. 调节 master 音量到合适级别")
        print("   3. 按 ESC 退出")
        print("   4. 运行 sudo alsactl store 保存音量设置")
        
        print("\n📋 i3 窗口管理器:")
        print("   1. 重启系统后会自动进入 i3")
        print("   2. 脚本已自动配置必要的设置(隐藏鼠标、禁用屏保等)")
        print("   3. 可以根据需要进一步自定义 ~/.config/i3/config")
    
    def run_initialization(self):
        """运行完整的初始化流程"""
        Logger.section("设备初始化开始")
        
        Logger.info(f"目标用户: {self.username}")
        Logger.info(f"用户主目录: {self.user_home}")
        Logger.info(f"安装 Chromium: {'是' if self.config.install_chromium else '否'}")
        Logger.info(f"预览模式: {'是' if self.config.dry_run else '否'}")
        
        # 初始化步骤配置
        initialization_steps = [
            ("优化系统启动", self.optimize_boot_time),
            ("更新系统包", self.update_system),
            ("安装基础依赖", self.install_basic_packages),
            ("添加软件源", self.add_repositories),
            ("安装主要软件", self.install_packages),
            ("配置系统设置", self.configure_system),
            ("配置自动启动X", self.setup_auto_startx),
            ("配置自动登录", self.setup_auto_login),
            ("配置i3窗口管理器", self.configure_i3),
        ]
        
        try:
            for step_name, step_func in initialization_steps:
                Logger.info(f"正在执行: {step_name}")
                step_func()
            
            self.show_manual_steps()
            
            Logger.section(self.SECTION_COMPLETE)
            Logger.success("🎉 所有自动化步骤已完成")
            Logger.warning("⚠️  请重启系统以使所有更改生效")
            Logger.info("💡 重启后请完成 WireGuard 和音量的手动配置")
            
        except KeyboardInterrupt:
            Logger.error("❌ 用户中断了初始化过程")
            sys.exit(1)
        except Exception as e:
            Logger.error(f"❌ 初始化过程中发生错误: {e}")
            sys.exit(1)

class SystemValidator:
    """系统环境验证器"""
    
    @staticmethod
    def validate_environment():
        """验证运行环境"""
        Logger.section("验证系统环境")
        
        # 检查是否通过 sudo 运行
        if os.geteuid() != 0:
            Logger.error("此脚本需要使用 sudo 运行")
            Logger.info("请使用: sudo python3 init-device.py")
            sys.exit(1)
        
        # 检查是否有原始用户信息
        if not os.getenv('SUDO_USER'):
            Logger.error("请不要直接以 root 用户运行此脚本")
            Logger.info("请使用普通用户通过 sudo 运行: sudo python3 init-device.py")
            sys.exit(1)
        
        Logger.info(f"检测到通过 sudo 运行脚本，原用户: {os.getenv('SUDO_USER')}")
        Logger.success("系统环境验证通过")

def create_argument_parser() -> argparse.ArgumentParser:
    """创建命令行参数解析器"""
    parser = argparse.ArgumentParser(
        description="设备初始化自动化脚本",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  %(prog)s                    # 基础安装
  %(prog)s --chromium         # 包含 Chromium 的完整安装
  %(prog)s --dry-run          # 预览模式，不实际执行
  %(prog)s --username user1   # 指定用户名
        """
    )
    
    parser.add_argument(
        "--chromium", 
        action="store_true", 
        help="安装 ungoogled-chromium 浏览器"
    )
    parser.add_argument(
        "--username", 
        type=str, 
        help="指定用户名（默认使用当前用户）"
    )
    parser.add_argument(
        "--dry-run", 
        action="store_true",
        help="预览模式：仅显示将要执行的操作，不实际执行"
    )
    parser.add_argument(
        "--version", 
        action="version", 
        version="%(prog)s 1.0.0"
    )
    
    return parser

def show_preview(config: InitConfig):
    """显示预览信息"""
    Logger.section("预览模式")
    
    steps = [
        DeviceInitializer.SECTION_OPTIMIZE_BOOT,
        DeviceInitializer.SECTION_UPDATE_SYSTEM, 
        f"{DeviceInitializer.SECTION_INSTALL_BASIC} (curl, gpg)",
        DeviceInitializer.SECTION_ADD_REPOS,
        f"{DeviceInitializer.SECTION_INSTALL_PACKAGES}" + (" (包含 ungoogled-chromium)" if config.install_chromium else ""),
        f"{DeviceInitializer.SECTION_CONFIGURE_SYSTEM} (时区、用户组)",
        DeviceInitializer.SECTION_AUTO_STARTX, 
        DeviceInitializer.SECTION_AUTO_LOGIN,
        DeviceInitializer.SECTION_CONFIGURE_I3,
        "显示手动配置步骤"
    ]
    
    Logger.info("将要执行的操作步骤:")
    for i, step in enumerate(steps, 1):
        print(f"  {i:2d}. {step}")
    
    Logger.warning("这是预览模式，不会实际执行任何操作")

def main():
    """主函数"""
    parser = create_argument_parser()
    args = parser.parse_args()
    
    # 创建配置
    config = InitConfig(
        install_chromium=args.chromium,
        username=args.username,
        dry_run=args.dry_run
    )
    
    # 预览模式
    if config.dry_run:
        show_preview(config)
        return
    
    # 验证环境
    SystemValidator.validate_environment()
    
    # 运行初始化
    initializer = DeviceInitializer(config)
    initializer.run_initialization()

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        Logger.error("用户中断操作")
        sys.exit(1)
    except Exception as e:
        Logger.error(f"程序异常退出: {e}")
        sys.exit(1)