"""Self Updater: drone-task-manager 自我更新模块。

支持上传 deploy 生成的 tar.gz 压缩包，解压验证后原子性替换并重启服务。
失败时自动回滚到备份版本。
"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import tarfile
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

from fastapi import HTTPException


def _check_flying() -> bool:
    """检测飞行状态。"""
    return Path("/tmp/arm.lock").exists() or Path("/tmp/flight.lock").exists()


class SelfUpdater:
    """管理 drone-task-manager 的自我更新和回滚。"""

    REQUIRED_FILES = ("main.py", "api.py", "media_manager.py", "requirements.txt")
    STATUS_FILE = Path("/tmp/dtm_self_update_status.json")

    # rsync 排除模式（与 deploy.sh 保持一致）
    EXCLUDE_PATTERNS = [
        ".venv/", "node_modules/", "__pycache__/",
        ".git/", "*.pyc", ".pytest_cache/",
        "data/", "models/", "backup/", "updater/",
        "push_streams.json", "detection_config.json", "detection_alerts.jsonl",
    ]

    def __init__(
        self,
        install_dir: str | Path | None = None,
        updater_dir: str | Path | None = None,
    ):
        self._install_dir = (
            Path(install_dir).resolve()
            if install_dir
            else Path(__file__).parent.resolve()
        )
        self._updater_dir = (
            Path(updater_dir).expanduser().resolve()
            if updater_dir
            else self._install_dir / "updater"
        )
        self._backup_dir = self._install_dir / "backup"
        self._updater_dir.mkdir(parents=True, exist_ok=True)
        self._backup_dir.mkdir(parents=True, exist_ok=True)
        self._init_status()

    # ────────── 状态管理 ──────────

    def _init_status(self):
        # 仅在状态文件不存在或当前状态为终端状态时才重置为 idle
        # 避免覆盖后台脚本写入的 applying/restarting/rolling_back 状态
        try:
            current = self.get_status()
            if current.get("status") in ("applying", "restarting", "rolling_back"):
                return
        except Exception:
            pass
        self._write_status(
            {"status": "idle", "step": "", "error": "", "timestamp": time.time()}
        )

    def _write_status(self, data: dict):
        data["timestamp"] = time.time()
        try:
            with open(self.STATUS_FILE, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False)
        except Exception:
            pass

    def get_status(self) -> dict:
        try:
            with open(self.STATUS_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return {"status": "idle", "step": "", "error": ""}

    # ────────── 包管理 ──────────

    @property
    def updater_dir(self) -> Path:
        return self._updater_dir

    def list_packages(self) -> List[dict]:
        packages = []
        for f in sorted(self._updater_dir.iterdir()):
            if f.is_file() and f.suffix.lower() in (
                ".gz", ".tgz", ".tar", ".zip"
            ):
                stat = f.stat()
                packages.append(
                    {
                        "filename": f.name,
                        "size": stat.st_size,
                        "mtime": datetime.fromtimestamp(
                            stat.st_mtime
                        ).isoformat(),
                        "path": str(f),
                    }
                )
        return packages

    def upload(self, filename: str, content: bytes) -> Path:
        if ".." in filename or "/" in filename or "\\" in filename:
            raise ValueError("Invalid filename")
        dest = self._updater_dir / filename
        dest.write_bytes(content)
        os.sync()
        return dest

    def delete_package(self, filename: str) -> None:
        if ".." in filename or "/" in filename or "\\" in filename:
            raise ValueError("Invalid filename")
        target = self._updater_dir / filename
        if target.exists():
            target.unlink()

    def delete_backup(self, name: str) -> None:
        if ".." in name or "/" in name or "\\" in name:
            raise ValueError("Invalid name")
        target = self._backup_dir / name
        if target.exists() and target.is_dir():
            import shutil
            shutil.rmtree(target)

    def list_backups(self) -> List[dict]:
        backups = []
        for b in sorted(
            self._backup_dir.iterdir(), key=lambda p: p.name, reverse=True
        ):
            if b.is_dir():
                stat = b.stat()
                size = sum(
                    f.stat().st_size for f in b.rglob("*") if f.is_file()
                )
                backups.append(
                    {
                        "name": b.name,
                        "path": str(b),
                        "size": size,
                        "mtime": datetime.fromtimestamp(
                            stat.st_mtime
                        ).isoformat(),
                    }
                )
        return backups

    def get_version_info(self) -> dict:
        info: Dict[str, str] = {}
        build_info = self._install_dir / "BUILD_INFO"
        if build_info.exists():
            try:
                with open(build_info, encoding="utf-8") as f:
                    for line in f:
                        if ":" in line:
                            k, v = line.strip().split(":", 1)
                            info[
                                k.strip().lower().replace(" ", "_")
                            ] = v.strip()
            except Exception:
                pass
        return info

    # ────────── 更新核心 ──────────

    def apply(self, package_path: Path) -> dict:
        """应用更新：解压 → 验证 → 备份 → 启动后台替换脚本。"""
        if _check_flying():
            raise HTTPException(
                status_code=423, detail="Cannot update while flying"
            )

        status = self.get_status()
        if status.get("status") in ("applying", "restarting", "rolling_back"):
            raise HTTPException(
                status_code=409, detail="Update already in progress"
            )

        package_path = Path(package_path).resolve()
        if not package_path.exists() or not package_path.is_file():
            raise HTTPException(
                status_code=404, detail="Package not found"
            )

        extract_dir = Path(f"/tmp/dtm_update_{int(time.time())}")
        try:
            # 1. 解压
            self._extract_tar(package_path, extract_dir)

            # 2. 找到源码根目录
            source_dir = self._find_source_dir(extract_dir)

            # 3. 验证
            self._validate(source_dir)

            # 4. 创建备份
            backup_dir = self._create_backup()

            # 5. 记录状态
            self._write_status(
                {
                    "status": "applying",
                    "step": "replacing_files",
                    "package": str(package_path.name),
                    "backup_dir": str(backup_dir),
                }
            )

            # 6. 生成并启动后台更新脚本
            script = self._generate_apply_script(source_dir, backup_dir)
            subprocess.Popen(
                ["/bin/bash", str(script)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                start_new_session=True,
            )

            return {
                "ok": True,
                "status": "applying",
                "message": "Update started, service will restart shortly",
                "backup_dir": str(backup_dir),
            }

        except HTTPException:
            raise
        except Exception as e:
            self._write_status(
                {
                    "status": "failed",
                    "step": "validation",
                    "error": str(e),
                }
            )
            raise HTTPException(
                status_code=500, detail=f"Update failed: {e}"
            )

    def rollback(self) -> dict:
        """回滚到最新的备份版本。"""
        if _check_flying():
            raise HTTPException(
                status_code=423, detail="Cannot rollback while flying"
            )

        status = self.get_status()
        if status.get("status") in ("applying", "restarting", "rolling_back"):
            raise HTTPException(
                status_code=409, detail="Update/rollback already in progress"
            )

        backups = sorted(
            self._backup_dir.iterdir(), key=lambda p: p.name, reverse=True
        )
        backups = [b for b in backups if b.is_dir()]
        if not backups:
            raise HTTPException(
                status_code=404, detail="No backup available"
            )

        backup_dir = backups[0]
        self._write_status(
            {
                "status": "rolling_back",
                "step": "restoring_backup",
                "backup_dir": str(backup_dir),
            }
        )

        script = self._generate_rollback_script(backup_dir)
        subprocess.Popen(
            ["/bin/bash", str(script)],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            start_new_session=True,
        )

        return {
            "ok": True,
            "status": "rolling_back",
            "message": "Rollback started, service will restart shortly",
            "backup_dir": str(backup_dir),
        }

    # ────────── 内部方法 ──────────

    def _extract_tar(self, tar_path: Path, dest_dir: Path):
        """安全解压 tar.gz，防止路径遍历。"""
        dest_dir.mkdir(parents=True, exist_ok=True)
        with tarfile.open(tar_path, "r:gz") as tf:
            for member in tf.getmembers():
                member_path = (dest_dir / member.name).resolve()
                if not str(member_path).startswith(str(dest_dir.resolve())):
                    raise ValueError(
                        f"Path traversal detected: {member.name}"
                    )
                tf.extract(member, dest_dir)

    def _find_source_dir(self, extract_dir: Path) -> Path:
        """找到解压后的源码根目录。tar.gz 通常包含一个顶层目录。"""
        entries = [
            e for e in extract_dir.iterdir() if not e.name.startswith(".")
        ]
        dirs = [e for e in entries if e.is_dir()]
        if len(dirs) == 1 and len(entries) == 1:
            return dirs[0]
        return extract_dir

    def _validate(self, source_dir: Path):
        """验证解压后的包包含必要的文件。"""
        missing = [
            f for f in self.REQUIRED_FILES if not (source_dir / f).exists()
        ]
        if missing:
            raise ValueError(
                f"Package validation failed, missing files: {missing}"
            )

    def _create_backup(self) -> Path:
        """创建当前版本的备份。"""
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_dir = self._backup_dir / f"backup_{ts}"
        backup_dir.mkdir(parents=True, exist_ok=True)

        for item in self._install_dir.iterdir():
            if item.name in (
                ".venv",
                "__pycache__",
                ".git",
                "node_modules",
                "backup",
            ):
                continue
            dest = backup_dir / item.name
            if item.is_dir():
                shutil.copytree(
                    item,
                    dest,
                    ignore=shutil.ignore_patterns(
                        "__pycache__", "*.pyc", ".git"
                    ),
                )
            else:
                shutil.copy2(item, dest)

        return backup_dir

    def _generate_apply_script(self, source_dir: Path, backup_dir: Path) -> Path:
        """生成后台更新执行脚本。"""
        script_path = Path(f"/tmp/dtm_apply_{int(time.time())}.sh")
        exclude_args = " ".join(
            f"--exclude='{p}'" for p in self.EXCLUDE_PATTERNS
        )

        try:
            stat = self._install_dir.stat()
            owner = f"{stat.st_uid}:{stat.st_gid}"
        except Exception:
            owner = "ramy:ramy"

        template = """#!/bin/bash
set -e

INSTALL_DIR="%(install_dir)s"
SOURCE_DIR="%(source_dir)s"
BACKUP_DIR="%(backup_dir)s"
STATUS_FILE="%(status_file)s"
SERVICE_NAME="%(service_name)s"
OWNER="%(owner)s"

write_status() {
    echo "$1" > "$STATUS_FILE"
}

# 忽略 systemd stop 时向整个 cgroup 发送的信号，确保脚本能执行完成
trap '' TERM INT

# 等待主服务退出，释放文件句柄
sleep 3

write_status '{"status":"applying","step":"replacing_files"}'

# 原子替换文件（排除配置和数据）
rsync -a --delete %(exclude_args)s "$SOURCE_DIR/" "$INSTALL_DIR/" || { RC=$?; [ "$RC" -eq 24 ] || exit "$RC"; }

# 修复权限
chown -R "$OWNER" "$INSTALL_DIR" 2>/dev/null || true

write_status '{"status":"restarting","step":"restarting_service"}'

# 获取当前主进程 PID（兼容旧版 systemd，不用 --value）
OLD_PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
if [ -n "$OLD_PID" ] && [ "$OLD_PID" != "0" ]; then
    kill -TERM "$OLD_PID" 2>/dev/null || true
else
    # fallback：如果取不到 PID，直接用 systemctl restart
    sudo -n systemctl restart "$SERVICE_NAME" 2>/dev/null || systemctl restart "$SERVICE_NAME" 2>/dev/null
fi

# 先等待主进程真正退出（systemd 随后会自动重启）
for i in {1..30}; do
    PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
    if [ "$PID" = "0" ] || [ -z "$PID" ]; then
        break
    fi
    sleep 1
done

# 再等待新进程启动（最多 60 秒）
for i in {1..60}; do
    NEW_PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
    if [ -n "$NEW_PID" ] && [ "$NEW_PID" != "0" ] && [ "$NEW_PID" != "$OLD_PID" ]; then
        write_status '{"status":"completed","step":"done"}'
        exit 0
    fi
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        write_status '{"status":"completed","step":"done"}'
        exit 0
    fi
    sleep 1
done

# 服务未运行，尝试自动回滚
write_status '{"status":"failed","step":"service_not_running","error":"Service failed to start, attempting rollback"}'

rsync -a --delete %(exclude_args)s "$BACKUP_DIR/" "$INSTALL_DIR/" || { RC=$?; [ "$RC" -eq 24 ] || exit "$RC"; }
chown -R "$OWNER" "$INSTALL_DIR" 2>/dev/null || true

# 回滚后再次触发主进程退出，让 systemd 自动重启
OLD_PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
if [ -n "$OLD_PID" ] && [ "$OLD_PID" != "0" ]; then
    kill -TERM "$OLD_PID" 2>/dev/null || true
else
    sudo -n systemctl restart "$SERVICE_NAME" 2>/dev/null || systemctl restart "$SERVICE_NAME" 2>/dev/null
fi

# 等待主进程退出
for i in {1..30}; do
    PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
    if [ "$PID" = "0" ] || [ -z "$PID" ]; then
        break
    fi
    sleep 1
done

# 等待新进程启动
for i in {1..60}; do
    NEW_PID=$(systemctl show --property=MainPID "$SERVICE_NAME" 2>/dev/null | cut -d= -f2)
    if [ -n "$NEW_PID" ] && [ "$NEW_PID" != "0" ] && [ "$NEW_PID" != "$OLD_PID" ]; then
        write_status '{"status":"rolled_back","step":"rollback_success"}'
        exit 0
    fi
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        write_status '{"status":"rolled_back","step":"rollback_success"}'
        exit 0
    fi
    sleep 1
done

write_status '{"status":"rollback_failed","step":"rollback_restart_failed","error":"Rollback failed, manual intervention required"}'
exit 1
"""

        script = template % {
            "install_dir": self._install_dir,
            "source_dir": source_dir,
            "backup_dir": backup_dir,
            "status_file": self.STATUS_FILE,
            "service_name": "drone-task-manager",
            "owner": owner,
            "exclude_args": exclude_args,
        }
        script_path.write_text(script, encoding="utf-8")
        os.sync()
        script_path.chmod(0o755)
        return script_path

    def _generate_rollback_script(self, backup_dir: Path) -> Path:
        """生成后台回滚执行脚本。"""
        script_path = Path(f"/tmp/dtm_rollback_{int(time.time())}.sh")
        exclude_args = " ".join(
            f"--exclude='{p}'" for p in self.EXCLUDE_PATTERNS
        )

        try:
            stat = self._install_dir.stat()
            owner = f"{stat.st_uid}:{stat.st_gid}"
        except Exception:
            owner = "ramy:ramy"

        template = """#!/bin/bash
set -e

INSTALL_DIR="%(install_dir)s"
BACKUP_DIR="%(backup_dir)s"
STATUS_FILE="%(status_file)s"
SERVICE_NAME="%(service_name)s"
OWNER="%(owner)s"

write_status() {
    echo "$1" > "$STATUS_FILE"
}

# 忽略 systemd stop 时向整个 cgroup 发送的信号，确保脚本能执行完成
trap '' TERM INT

# 等待主服务退出
sleep 3

write_status '{"status":"rolling_back","step":"restoring_backup"}'

# 恢复备份
rsync -a --delete %(exclude_args)s "$BACKUP_DIR/" "$INSTALL_DIR/" || { RC=$?; [ "$RC" -eq 24 ] || exit "$RC"; }
chown -R "$OWNER" "$INSTALL_DIR" 2>/dev/null || true

write_status '{"status":"restarting","step":"restarting_service"}'

# 获取当前主进程 PID，发送 SIGTERM 使其退出，systemd 会自动重启
OLD_PID=$(systemctl show --property=MainPID --value "$SERVICE_NAME" 2>/dev/null)
if [ -n "$OLD_PID" ] && [ "$OLD_PID" != "0" ]; then
    kill -TERM "$OLD_PID" 2>/dev/null || true
else
    sudo -n systemctl restart "$SERVICE_NAME" 2>/dev/null || systemctl restart "$SERVICE_NAME" 2>/dev/null
fi

sleep 8
for i in {1..20}; do
    NEW_PID=$(systemctl show --property=MainPID --value "$SERVICE_NAME" 2>/dev/null)
    if [ -n "$NEW_PID" ] && [ "$NEW_PID" != "0" ] && [ "$NEW_PID" != "$OLD_PID" ]; then
        write_status '{"status":"rolled_back","step":"rollback_success"}'
        exit 0
    fi
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        write_status '{"status":"rolled_back","step":"rollback_success"}'
        exit 0
    fi
    sleep 1
done

write_status '{"status":"rollback_failed","step":"rollback_restart_failed","error":"Rollback failed, manual intervention required"}'
exit 1
"""

        script = template % {
            "install_dir": self._install_dir,
            "backup_dir": backup_dir,
            "status_file": self.STATUS_FILE,
            "service_name": "drone-task-manager",
            "owner": owner,
            "exclude_args": exclude_args,
        }
        script_path.write_text(script, encoding="utf-8")
        os.sync()
        script_path.chmod(0o755)
        return script_path
