# -*- coding: utf-8 -*-
"""
runner_service.py — 巡检执行编排

职责：
  - 选择执行器（local/SSH）
  - 渲染 config、上传资产、逐模块执行
  - 生成 SSE 事件流（progress / module_done / finished / error）
  - 支持取消运行

运行模型：每次 run_id 独立工作目录与线程，结果累积后落盘。
"""

from __future__ import annotations

import logging
import threading
import time
import traceback
from datetime import datetime
from typing import Dict, Optional

from ..utils.check_modules import get_suite, render_config, CheckModule
from ..utils.parser import parse_module_output, summarize
from ..utils.executor import BaseExecutor
from ..services import target_service, report_service

logger = logging.getLogger("service_monitor.runner")

# 运行中任务注册表（支持取消）
_runs: Dict[str, dict] = {}
_runs_lock = threading.Lock()


def _now() -> str:
    return datetime.now().strftime("%Y-%m-%dT%H:%M:%S")


def _register_run(run_id: str, target_id: str, suite: str, total: int) -> None:
    with _runs_lock:
        _runs[run_id] = {
            "target_id": target_id,
            "suite": suite,
            "total": total,
            "done": 0,
            "current": "",
            "cancel": threading.Event(),
            "finished": False,
            "error": None,
        }


def _update_run(run_id: str, **kwargs) -> None:
    with _runs_lock:
        if run_id in _runs:
            _runs[run_id].update(kwargs)


def _pop_run(run_id: str) -> Optional[dict]:
    with _runs_lock:
        return _runs.get(run_id)


def cancel_run(run_id: str) -> bool:
    """请求取消运行。"""
    with _runs_lock:
        run = _runs.get(run_id)
    if not run:
        return False
    run["cancel"].set()
    logger.info("请求取消运行: %s", run_id)
    return True


def get_run_status(run_id: str) -> Optional[dict]:
    """获取运行状态（供轮询/SSE 心跳）。"""
    with _runs_lock:
        run = _runs.get(run_id)
    if not run:
        return None
    return {
        "done": run["done"],
        "total": run["total"],
        "current": run["current"],
        "finished": run["finished"],
        "error": run["error"],
    }


def run_inspection(target_id: str, suite: str):
    """巡检生成器：yield SSE 事件 dict。

    事件格式：
      {"event": "progress", "run_id", "done", "total", "current"}
      {"event": "module_done", "run_id", "module", "items_count", "summary"}
      {"event": "finished", "run_id", "report_id", "summary"}
      {"event": "error", "run_id", "message"}
      {"event": "cancelled", "run_id"}
    """
    from datetime import datetime as _dt
    run_id = _dt.now().strftime("%Y%m%d_%H%M%S_") + str(threading.get_ident())[-4:]

    target = target_service.get_target(target_id)
    if not target:
        yield {"event": "error", "run_id": run_id, "message": "监测目标不存在"}
        return

    modules = get_suite(suite)
    if not modules:
        yield {"event": "error", "run_id": run_id, "message": f"套件 {suite} 无可用模块"}
        return

    _register_run(run_id, target_id, suite, len(modules))
    started_at = _now()

    # 准备 config（含解密后的凭据）
    target_for_render = dict(target)
    target_for_render["credential_overrides"] = target_service.resolve_credentials(target)
    config_text = render_config(target_for_render)

    executor: Optional[BaseExecutor] = None
    modules_result = []
    try:
        executor = target_service.make_executor(target, run_id)

        # 连通性检查（远程）
        if not executor.is_local:
            yield {"event": "progress", "run_id": run_id, "done": 0, "total": len(modules),
                   "current": "测试SSH连接..."}
            if not executor.test_connection():
                yield {"event": "error", "run_id": run_id, "message": "SSH 连接失败，请检查目标配置"}
                _update_run(run_id, finished=True, error="SSH 连接失败")
                return
            yield {"event": "progress", "run_id": run_id, "done": 0, "total": len(modules),
                   "current": "上传检测模块..."}

        executor.setup()
        executor.upload_assets()

        for idx, module in enumerate(modules):
            # 检查取消
            run = _pop_run(run_id)
            if run and run["cancel"].is_set():
                yield {"event": "cancelled", "run_id": run_id}
                _update_run(run_id, finished=True)
                return

            _update_run(run_id, current=module.name, done=idx)
            yield {"event": "progress", "run_id": run_id, "done": idx,
                   "total": len(modules), "current": module.name}

            raw = executor.run_module(module, config_text, timeout=90)
            items = parse_module_output(raw, module.id, module.category)
            summary = summarize(items)

            module_result = {
                "id": module.id,
                "name": module.name,
                "category": module.category,
                "items": [{"key": it.key, "name": it.name, "value": it.value,
                           "threshold": it.threshold, "status": it.status}
                          for it in items],
                "summary": summary,
            }
            modules_result.append(module_result)
            _update_run(run_id, done=idx + 1)

            yield {"event": "module_done", "run_id": run_id,
                   "module": module.name, "items_count": len(items), "summary": summary}

        report_id = report_service.save(target, suite, modules_result, started_at, _now())
        _update_run(run_id, finished=True)

        # 汇总
        total_summary = {k: 0 for k in ("正常", "警告", "严重", "total")}
        for m in modules_result:
            ms = m["summary"]
            for k in total_summary:
                total_summary[k] += ms.get(k, 0)

        yield {"event": "finished", "run_id": run_id, "report_id": report_id,
               "summary": total_summary}

    except Exception as e:
        logger.error("巡检执行异常: %s\n%s", e, traceback.format_exc())
        _update_run(run_id, finished=True, error=str(e))
        yield {"event": "error", "run_id": run_id, "message": f"巡检执行异常: {e}"}
    finally:
        if executor:
            try:
                executor.cleanup()
            except Exception as e:
                logger.warning("执行器清理失败: %s", e)
        # 延迟移除运行记录，给客户端最后取状态的机会
        time.sleep(0)


def run_inspection_sync(target_id: str, suite: str) -> dict:
    """同步执行巡检（定时任务调用），返回结果摘要。

    与 run_inspection() 生成器共享核心逻辑，但不 yield SSE 事件。
    返回: {"success": bool, "report_id": str, "summary": dict, "finished_at": str}
    """
    from datetime import datetime as _dt
    run_id = _dt.now().strftime("%Y%m%d_%H%M%S_") + "sched"

    target = target_service.get_target(target_id)
    if not target:
        return {"success": False, "report_id": "", "summary": {}, "finished_at": _now()}

    modules = get_suite(suite)
    if not modules:
        return {"success": False, "report_id": "", "summary": {}, "finished_at": _now()}

    _register_run(run_id, target_id, suite, len(modules))
    started_at = _now()

    # 准备 config（含解密后的凭据）
    target_for_render = dict(target)
    target_for_render["credential_overrides"] = target_service.resolve_credentials(target)
    config_text = render_config(target_for_render)

    executor = None
    modules_result = []
    try:
        executor = target_service.make_executor(target, run_id)

        # 连通性检查（远程）
        if not executor.is_local:
            if not executor.test_connection():
                _update_run(run_id, finished=True, error="SSH 连接失败")
                return {"success": False, "report_id": "", "summary": {}, "finished_at": _now()}

        executor.setup()
        executor.upload_assets()

        for idx, module in enumerate(modules):
            # 检查取消
            run = _pop_run(run_id)
            if run and run["cancel"].is_set():
                _update_run(run_id, finished=True)
                return {"success": False, "report_id": "", "summary": {}, "finished_at": _now()}

            _update_run(run_id, current=module.name, done=idx)
            raw = executor.run_module(module, config_text, timeout=90)
            items = parse_module_output(raw, module.id, module.category)
            summary = summarize(items)

            module_result = {
                "id": module.id,
                "name": module.name,
                "category": module.category,
                "items": [{"key": it.key, "name": it.name, "value": it.value,
                           "threshold": it.threshold, "status": it.status}
                          for it in items],
                "summary": summary,
            }
            modules_result.append(module_result)
            _update_run(run_id, done=idx + 1)

        report_id = report_service.save(target, suite, modules_result, started_at, _now())
        _update_run(run_id, finished=True)

        # 汇总
        total_summary = {k: 0 for k in ("正常", "警告", "严重", "total")}
        for m in modules_result:
            ms = m["summary"]
            for k in total_summary:
                total_summary[k] += ms.get(k, 0)

        return {
            "success": True,
            "report_id": report_id,
            "summary": total_summary,
            "finished_at": _now(),
        }

    except Exception as e:
        logger.error("同步巡检执行异常: %s\n%s", e, traceback.format_exc())
        _update_run(run_id, finished=True, error=str(e))
        return {"success": False, "report_id": "", "summary": {}, "finished_at": _now()}
    finally:
        if executor:
            try:
                executor.cleanup()
            except Exception as e:
                logger.warning("执行器清理失败: %s", e)
