#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
服务器日志下载脚本

从生产服务器下载后端日志到本地 logs/server/ 目录，支持多种查询模式。

使用方法:
  python tools/download-server-log.py                       # 默认: 下载完整日志
  python tools/download-server-log.py latest                # 下载完整日志
  python tools/download-server-log.py tail [N]              # 只看最近 N 行 (默认 200)
  python tools/download-server-log.py errors [N]            # 只看最近 N 条 ERROR/Exception
  python tools/download-server-log.py grep <keyword> [N]    # 按关键词搜索 (默认 200 条)

公共参数:
  --lines N        tail/errors/grep 返回的最大行数，默认 200
  --context N      grep 时显示的上下文行数，默认 3
  -q, --quiet      静默模式，只写文件不打印到 stdout

Author: Peanut
Created: 2026-06-03
Purpose: 从服务器下载 emotion-single.log 到本地，方便排查问题
"""

import argparse
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

# ============================================================================
# 配置
# ============================================================================
SERVER_IP = "101.200.208.45"
USERNAME = "root"
REMOTE_LOG = "/data/logs/emotion-museum/emotion-single.log"

PROJECT_DIR = Path(__file__).parent.parent.absolute()
LOCAL_LOG_DIR = PROJECT_DIR / "logs" / "server"

SSH_COMMON = ["-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=no"]


# ============================================================================
# 日志输出
# ============================================================================
def log_info(msg):
    print(f"\033[32m[INFO]\033[0m {msg}", flush=True)


def log_warn(msg):
    print(f"\033[33m[WARN]\033[0m {msg}", flush=True)


def log_error(msg):
    print(f"\033[31m[ERROR]\033[0m {msg}", flush=True)


# ============================================================================
# 基础工具
# ============================================================================
def timestamp_suffix():
    return datetime.now().strftime("%Y%m%d-%H%M%S")


def ensure_local_dir():
    LOCAL_LOG_DIR.mkdir(parents=True, exist_ok=True)


def check_ssh():
    try:
        result = subprocess.run(
            ["ssh", *SSH_COMMON, f"{USERNAME}@{SERVER_IP}", "echo ok"],
            capture_output=True, text=True, timeout=15
        )
        if result.returncode == 0:
            log_info("SSH 连接正常")
            return True
        log_error(f"SSH 连接失败: {result.stderr.strip()}")
        log_error(f"请确认已配置免密登录: ssh {USERNAME}@{SERVER_IP}")
        return False
    except subprocess.TimeoutExpired:
        log_error("SSH 连接超时")
        return False
    except FileNotFoundError:
        log_error("未找到 ssh 命令，请确认已安装 OpenSSH client")
        return False


def get_remote_file_size():
    try:
        result = subprocess.run(
            ["ssh", *SSH_COMMON, f"{USERNAME}@{SERVER_IP}",
             f"stat -c %s {REMOTE_LOG} 2>/dev/null || echo -1"],
            capture_output=True, text=True, timeout=15
        )
        size = int(result.stdout.strip())
        return size if size > 0 else -1
    except (subprocess.TimeoutExpired, ValueError, OSError) as e:
        log_error(f"获取远程文件大小失败: {e}")
        return -1


def format_size(size_bytes):
    if size_bytes < 0:
        return "未知"
    for unit in ["B", "KB", "MB", "GB"]:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} TB"


def backup_local_log(target: Path):
    bak = target.with_suffix(target.suffix + ".bak")
    if target.exists():
        shutil.copy2(target, bak)
        log_info(f"已备份旧文件: {bak}")


def write_local_file(path: Path, content: str, quiet: bool):
    path.write_text(content, encoding="utf-8")
    log_info(f"已写入本地文件: {path} ({format_size(len(content.encode('utf-8')))})")
    if not quiet:
        sys.stdout.write(content)
        if not content.endswith("\n"):
            sys.stdout.write("\n")
        sys.stdout.flush()


def ssh_capture(remote_cmd, timeout=120):
    """在远端执行命令并捕获 stdout。失败返回 None。"""
    try:
        result = subprocess.run(
            ["ssh", *SSH_COMMON, f"{USERNAME}@{SERVER_IP}", remote_cmd],
            capture_output=True, text=True, encoding="utf-8", errors="replace",
            timeout=timeout
        )
        if result.returncode != 0:
            log_error(f"远端命令失败: {remote_cmd}")
            if result.stderr.strip():
                log_error(result.stderr.strip())
            return None
        return result.stdout
    except subprocess.TimeoutExpired:
        log_error(f"远端命令超时 ({timeout}s): {remote_cmd}")
        return None
    except Exception as e:
        log_error(f"远端命令异常: {e}")
        return None


def shlex_quote(s):
    """简易 shell 转义：用单引号包裹，内部单引号用 '\\'' 处理。"""
    return "'" + s.replace("'", "'\\''") + "'"


# ============================================================================
# 子命令实现
# ============================================================================
def cmd_latest(args):
    ensure_local_dir()
    remote_size = get_remote_file_size()
    if remote_size < 0:
        log_error(f"远程日志文件不存在或无法访问: {REMOTE_LOG}")
        return False

    log_info(f"远程日志大小: {format_size(remote_size)}")
    local_path = LOCAL_LOG_DIR / f"emotion-single-{timestamp_suffix()}.log"
    stable_link = LOCAL_LOG_DIR / "emotion-single-latest.log"

    log_info(f"正在下载: {USERNAME}@{SERVER_IP}:{REMOTE_LOG}")
    log_info(f"        -> {local_path}")

    try:
        result = subprocess.run(
            ["scp", *SSH_COMMON,
             f"{USERNAME}@{SERVER_IP}:{REMOTE_LOG}", str(local_path)],
            capture_output=True, text=True, timeout=300
        )
        if result.returncode != 0:
            log_error(f"下载失败: {result.stderr.strip()}")
            return False
    except subprocess.TimeoutExpired:
        log_error("下载超时 (300s)")
        return False
    except Exception as e:
        log_error(f"下载异常: {e}")
        return False

    # 维护一个固定名的拷贝，方便快速打开
    try:
        if stable_link.exists() or stable_link.is_symlink():
            stable_link.unlink()
        shutil.copy2(local_path, stable_link)
    except OSError:
        pass

    log_info(f"=== 下载完成: {local_path} ({format_size(local_path.stat().st_size)}) ===")
    return True


def cmd_tail(args):
    ensure_local_dir()
    n = args.lines
    log_info(f"拉取远程日志最后 {n} 行...")
    output = ssh_capture(f"tail -n {n} {REMOTE_LOG}")
    if output is None:
        return False
    local_path = LOCAL_LOG_DIR / f"tail-{timestamp_suffix()}.log"
    write_local_file(local_path, output, args.quiet)
    return True


def cmd_errors(args):
    ensure_local_dir()
    n = args.lines
    log_info(f"拉取远程日志中最近 {n} 条 ERROR/Exception ...")
    cmd = f"grep -nE 'ERROR|Exception' {REMOTE_LOG} | tail -n {n}"
    output = ssh_capture(cmd)
    if output is None:
        return False
    local_path = LOCAL_LOG_DIR / f"errors-{timestamp_suffix()}.log"
    write_local_file(local_path, output, args.quiet)
    return True


def cmd_grep(args):
    ensure_local_dir()
    keyword = args.keyword
    if not keyword:
        log_error("grep 子命令必须提供 keyword")
        return False
    # 防御：禁止 keyword 中包含 shell 元字符
    if any(ch in keyword for ch in "`$()|;&<>'\"\\"):
        log_error(f"keyword 含有非法字符: {keyword}")
        return False

    n = args.lines
    c = args.context
    log_info(f"按关键词搜索: {keyword!r} (上下文 {c} 行, 最多 {n} 条)")
    cmd = f"grep -nF -C {c} -- {shlex_quote(keyword)} {REMOTE_LOG} | tail -n {n}"
    output = ssh_capture(cmd)
    if output is None:
        return False
    local_path = LOCAL_LOG_DIR / f"grep-{timestamp_suffix()}.log"
    write_local_file(local_path, output, args.quiet)
    return True


# ============================================================================
# 主入口
# ============================================================================
def build_parser():
    parser = argparse.ArgumentParser(
        description="从服务器下载 emotion-single.log，支持多种查询模式"
    )
    parser.add_argument("--lines", type=int, default=200,
                        help="tail/errors/grep 返回的最大行数 (默认 200)")
    parser.add_argument("--context", type=int, default=3,
                        help="grep 时显示的上下文行数 (默认 3)")
    parser.add_argument("-q", "--quiet", action="store_true",
                        help="静默模式：只写文件，不打印到 stdout")

    sub = parser.add_subparsers(dest="command")

    sub.add_parser("latest", help="下载完整日志 (默认)")

    p_tail = sub.add_parser("tail", help="只看最后 N 行")
    p_tail.add_argument("n", nargs="?", type=int, help="行数 (默认用 --lines)")

    p_err = sub.add_parser("errors", help="只看最近 N 条 ERROR/Exception")
    p_err.add_argument("n", nargs="?", type=int, help="条数 (默认用 --lines)")

    p_grep = sub.add_parser("grep", help="按关键词搜索")
    p_grep.add_argument("keyword", help="搜索关键词")
    p_grep.add_argument("n", nargs="?", type=int, help="最多返回条数 (默认用 --lines)")

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    log_info("=== 情绪博物馆 - 日志下载工具 ===")
    log_info(f"服务器: {USERNAME}@{SERVER_IP}")
    log_info(f"远程日志: {REMOTE_LOG}")

    if not check_ssh():
        sys.exit(1)

    # 兼容老版本无参数调用 → 当作 latest
    if args.command is None:
        args.command = "latest"

    # 把位置参数 n 覆盖到 args.lines (如果用户显式传了)
    if hasattr(args, "n") and args.n is not None:
        args.lines = args.n

    dispatch = {
        "latest": cmd_latest,
        "tail": cmd_tail,
        "errors": cmd_errors,
        "grep": cmd_grep,
    }

    ok = dispatch[args.command](args)
    if ok:
        log_info("=== 完成 ===")
    else:
        log_error("操作失败")
        sys.exit(1)


if __name__ == "__main__":
    main()