from typing import List, Dict, Tuple, Set, Any, Optional, Union
#!/usr/bin/env python3
"""
collect.py — Palladium 远程数据采集推送脚本

在 Palladium 主机上运行，执行 test_server -short 并将结果推送到监控服务器。

用法:
  python3 collect.py                                          # 默认配置
  python3 collect.py --url http://monitor:5100/api/collect
  python3 collect.py --host scmp03 --env /edatools/vxe/VXE23.03.001/bin/env.csh
  python3 collect.py --dry-run                                # 只采集不推送
  python3 collect.py --log /var/log/palladium-collect.log     # 追加写日志

可配合 crontab 定时执行 (每 30 分钟):
  */30 * * * * /opt/palladium-monitor/collect.py --host scmp03 --url http://MONITOR:5100/api/collect --log /var/log/palladium-collect.log >> /var/log/palladium-collect.cron.log 2>&1
"""

import argparse
import os
import socket
import subprocess
import sys
import time
from datetime import datetime

try:
    import requests
except ImportError:
    print("ERROR: requests 未安装，请执行 pip install requests", file=sys.stderr)
    sys.exit(1)


# ═══════════════════════════════════════════════════════════════════════════
#  Defaults
# ═══════════════════════════════════════════════════════════════════════════
DEFAULT_URL = "http://127.0.0.1:5000/api/collect"
DEFAULT_CMD = "test_server -short"
DEFAULT_TIMEOUT = 300
DEFAULT_RETRIES = 3
DEFAULT_RETRY_DELAY = 5


# ═══════════════════════════════════════════════════════════════════════════
#  Logging
# ═══════════════════════════════════════════════════════════════════════════
class Logger:
    """同时输出到 stderr 和可选的日志文件。"""

    def __init__(self, log_file: Optional[str] = None):
        self.log_file = log_file

    def _emit(self, level: str, msg: str, ts: str) -> None:
        line = f"{ts} [{level}] {msg}"
        print(line, file=sys.stderr)
        if self.log_file:
            try:
                os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
                with open(self.log_file, "a", encoding="utf-8") as f:
                    f.write(line + "\n")
            except OSError as e:
                print(f"{ts} [WARN] 无法写入日志文件 {self.log_file}: {e}", file=sys.stderr)

    def info(self, msg: str) -> None:
        self._emit("INFO", msg, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    def warn(self, msg: str) -> None:
        self._emit("WARN", msg, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    def error(self, msg: str) -> None:
        self._emit("ERROR", msg, datetime.now().strftime("%Y-%m-%d %H:%M:%S"))


# ═══════════════════════════════════════════════════════════════════════════
#  Command Execution
# ═══════════════════════════════════════════════════════════════════════════
def build_shell_cmd(cmd: str, env_script: Optional[str]) -> str:
    """
    如果指定了 env_script，把 source env 之后再执行 cmd 拼成一个 shell 串。
    不指定则直接返回 cmd。
    """
    if not env_script:
        return cmd
    # env 文件可能是 csh 或 sh 后缀, 用 bash -c 强制走 sh 语义, 兼容大多数 vendor
    # 假设 vendor 提供的 env.* 是 POSIX 兼容 (setenv → export)
    return f"source {shlex_quote(env_script)} >/dev/null 2>&1; {cmd}"


def shlex_quote(s: str) -> str:
    """最小 shell 单引号转义 (用于 env_script 路径里有空格的情况)。"""
    return "'" + s.replace("'", "'\\''") + "'"


def run_test_server(cmd: str, timeout: int, env_script: Optional[str]) -> Tuple[str, int, str]:
    """
    执行 test_server -short, 返回 (stdout, returncode, stderr)。
    timeout 触发时 subprocess.run 不会抛 (>=3.12 subprocessing 时可能抛), 这里显式处理。
    """
    full_cmd = build_shell_cmd(cmd, env_script)
    try:
        result = subprocess.run(
            full_cmd,
            shell=True,
            executable="/bin/bash",
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            timeout=timeout,
        )
    except subprocess.TimeoutExpired as e:
        return (e.stdout.decode("utf-8", errors="replace") if e.stdout else "", -1,
                f"命令超时 (>={timeout}s)")
    except FileNotFoundError as e:
        return ("", 127, f"shell 不可用: {e}")

    stdout = result.stdout.decode("utf-8", errors="replace") if result.stdout else ""
    stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else ""
    return (stdout, result.returncode, stderr)


# ═══════════════════════════════════════════════════════════════════════════
#  Push to monitor
# ═══════════════════════════════════════════════════════════════════════════
def push_to_monitor(url: str, raw_output: str, source: str,
                    retries: int, retry_delay: int,
                    log: Logger) -> Tuple[bool, Optional[Dict], str]:
    """带重试的 HTTP 推送, 返回 (ok, json_body_or_None, error_msg)。"""
    payload = {"raw_output": raw_output, "source": source}
    last_err = ""
    for attempt in range(1, retries + 1):
        try:
            resp = requests.post(url, json=payload, timeout=30)
        except requests.exceptions.ConnectionError as e:
            last_err = f"无法连接 {url}: {e}"
        except requests.exceptions.Timeout as e:
            last_err = f"推送超时: {e}"
        except requests.exceptions.RequestException as e:
            last_err = f"请求异常: {e}"
        else:
            if resp.status_code == 200:
                try:
                    return (True, resp.json(), "")
                except ValueError:
                    return (True, None, "")
            last_err = f"HTTP {resp.status_code}: {resp.text[:200]}"

        if attempt < retries:
            log.warn(f"推送失败 (第 {attempt}/{retries} 次): {last_err} — {retry_delay}s 后重试")
            time.sleep(retry_delay)
    return (False, None, last_err)


# ═══════════════════════════════════════════════════════════════════════════
#  Main
# ═══════════════════════════════════════════════════════════════════════════
def main() -> int:
    parser = argparse.ArgumentParser(
        description="Palladium 数据采集推送 (cron 友好, 带重试和日志)",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--url", default=DEFAULT_URL,
                        help=f"监控服务器 API 地址 (默认: {DEFAULT_URL})")
    parser.add_argument("--cmd", default=DEFAULT_CMD,
                        help=f"采集命令 (默认: {DEFAULT_CMD})")
    parser.add_argument("--host", default=socket.gethostname(),
                        help="本机标签, 会写入 source 字段 (默认: hostname)")
    parser.add_argument("--env", default=None,
                        help="VXE env.csh/env.sh 路径, 自动 source 后再执行命令")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
                        help=f"命令超时秒数 (默认: {DEFAULT_TIMEOUT})")
    parser.add_argument("--retries", type=int, default=DEFAULT_RETRIES,
                        help=f"推送失败重试次数 (默认: {DEFAULT_RETRIES})")
    parser.add_argument("--retry-delay", type=int, default=DEFAULT_RETRY_DELAY,
                        help=f"重试间隔秒数 (默认: {DEFAULT_RETRY_DELAY})")
    parser.add_argument("--log", default=None,
                        help="追加写日志的文件路径 (例如 /var/log/palladium-collect.log)")
    parser.add_argument("--dry-run", action="store_true",
                        help="只采集不推送, 打印解析结果")
    args = parser.parse_args()

    log = Logger(args.log)
    source = f"collect.py@{args.host}"
    log.info(f"开始采集 — host={args.host} cmd={args.cmd!r} env={args.env!r}")

    # ─── 执行采集命令 ──────────────────────────────────────────
    stdout, rc, stderr = run_test_server(args.cmd, args.timeout, args.env)
    if rc != 0:
        log.warn(f"命令退出码 {rc} (stderr: {stderr.strip()[:300]})")
    if not stdout.strip():
        log.error("命令无输出, 终止")
        return 1

    log.info(f"采集成功 — {len(stdout)} 字节, rc={rc}")

    # ─── dry-run 模式 ──────────────────────────────────────────
    if args.dry_run:
        print("─" * 60)
        print(stdout)
        print("─" * 60)
        log.info("dry-run 模式, 不推送")
        return 0

    # ─── 推送到监控服务器 ──────────────────────────────────────
    log.info(f"推送到 {args.url} (source={source})")
    ok, body, err = push_to_monitor(
        args.url, stdout, source,
        args.retries, args.retry_delay, log,
    )

    if not ok:
        log.error(f"推送失败 (已重试 {args.retries} 次): {err}")
        return 2

    # 成功 — 打印 stats 摘要
    snapshot_id = body.get("snapshot_id") if body else "?"
    stats = body.get("stats") if body else None
    if stats:
        log.info(
            f"推送成功 — snapshot #{snapshot_id} | "
            f"板 {stats['total_boards']} (在线 {stats['online_boards']}/离线 {stats['offline_boards']}) | "
            f"利用率 {stats['utilization']}% (使用 {stats['used_boards']}/{stats['total_boards']}) | "
            f"作业 {stats['active_jobs']}"
        )
    else:
        log.info(f"推送成功 — snapshot #{snapshot_id} (无 stats 字段)")

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\n[!] 用户中断", file=sys.stderr)
        sys.exit(130)