#!/usr/bin/env python3
"""
crawl_scheduler_v2.py - 智能爬虫调度器 v2
==========================================
核心思路：将所有爬虫按运行状态分组 → 健康脚本优先跑 → 同类问题集中报告

分组逻辑：
  - 正常群: 上次运行成功 → 全部执行
  - 新脚本: 从未运行过 → 全部执行  
  - 超时群: 统一超时 → 集中报告
  - 站点不可达: 连接被拒/DNS/网络不可达 → 集中报告
  - 脚本异常群: Python错误/Traceback → 集中报告
  - 其他失败群: 未知错误 → 集中报告

工作流：
  1. 读取 daily_crawl_config.json + run_status.json
  2. 将脚本按状态分组
  3. 执行正常群(并行 3 个 worker, 含重试)
  4. 执行新脚本(并行 3 个 worker)
  5. 输出分组报告 → 用户修复问题群

用法:
  python3 crawl_scheduler_v2.py              # 执行所有脚本（智能调度）
  python3 crawl_scheduler_v2.py --dry-run     # 仅分组预览，不执行
  python3 crawl_scheduler_v2.py --group 江苏   # 只跑指定分组
  python3 crawl_scheduler_v2.py --parallel 5  # 调整并行数
"""

import json, os, sys, time, subprocess, concurrent.futures, re, socket, sqlite3
from datetime import datetime
from collections import defaultdict

# ==================== 配置 ====================
BASE_DIR = "/root/gov_crawler"
CONFIG_PATH = os.path.join(BASE_DIR, "daily_crawl_config.json")
STATUS_PATH = os.path.join(BASE_DIR, "run_status.json")
DB_PATH = "/root/search.db"
DEFAULT_PARALLEL = 3
SCRIPT_TIMEOUT = 600
RETRY_FAILED = False  # 失败脚本不自动重试（让用户排查）

# ==================== 工具函数 ====================
def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)

def load_config():
    with open(CONFIG_PATH, encoding="utf-8") as f:
        cfg = json.load(f)
        # 兼容 flat-list 格式 (search_app 同款逻辑)
        if isinstance(cfg, list):
            cfg = {"crawlers": cfg}
        return cfg["crawlers"]


def cname(c):
    """兼容 name 缺失的历史条目: name -> config_name -> site_name"""
    return c.get("name") or c.get("config_name") or c.get("site_name") or c.get("script", "?")

def load_status():
    try:
        with open(STATUS_PATH, encoding="utf-8") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}

def save_status(status_map):
    with open(STATUS_PATH, "w", encoding="utf-8") as f:
        json.dump(status_map, f, ensure_ascii=False, indent=2)

def write_run_log(script_name, config_name, group_name, status, new_count, elapsed, error_detail=""):
    """写入运行日志到 search.db run_logs 表"""
    today = datetime.now().strftime("%Y-%m-%d")
    now_time = datetime.now().strftime("%H:%M:%S")
    try:
        conn = sqlite3.connect(DB_PATH, timeout=60)
        conn.execute("INSERT OR REPLACE INTO run_logs "
            "(run_date, run_time, script_name, config_name, group_name, status, new_count, elapsed_seconds, error_detail) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (today, now_time, script_name, config_name, group_name, status, new_count, elapsed, error_detail[:500]))
        conn.commit()
        conn.close()
    except Exception as e:
        print(f"  [WARN] run_log写入失败: {e}")

def run_script(script, args="", timeout=SCRIPT_TIMEOUT):
    """执行单个脚本，返回 (output, exit_code, elapsed)"""
    t0 = time.time()
    if isinstance(args, list):
        args = " ".join(str(a) for a in args)
    cmd = f'cd "{BASE_DIR}" && python3 "{script}" {args}'
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
        elapsed = time.time() - t0
        out = (r.stdout or "") + ("\n" + r.stderr if r.stderr else "")
        return out.strip(), r.returncode, round(elapsed, 1)
    except subprocess.TimeoutExpired:
        elapsed = time.time() - t0
        return "[TIMEOUT]", -1, round(elapsed, 1)
    except Exception as e:
        elapsed = time.time() - t0
        return f"[EXCEPTION] {e}", -1, round(elapsed, 1)

def classify_error(out, code):
    """根据输出和退出码分类错误类型"""
    if code == 0:
        return ""
    if code == -1 and "TIMEOUT" in out.upper():
        return "timeout"
    err_lower = out.lower()
    if any(x in err_lower for x in ["connection refused", "connection reset", "连接被拒绝", "connect failed", "errno 111"]):
        return "connection_refused"
    if any(x in err_lower for x in ["name or service not known", "temporary failure in name resolution", "无法解析", "resolve", "dns", "errno -2"]):
        return "dns_resolve"
    if any(x in err_lower for x in ["no route to host", "network is unreachable", "无法访问", "网络不可达", "errno 101", "errno 113", "连接超时"]):
        return "unreachable"
    if "read timed out" in err_lower or "timeout" in err_lower:
        return "read_timeout"
    if "traceback" in err_lower or "error" in err_lower:
        return "script_error"
    return "exit_error"

def extract_new_count(output):
    """从脚本输出中提取新增条数 - 从最后一行往前搜，避免匹配到标题里的新增XX万字样"""
    lines = output.strip().split('\n')
    for line in reversed(lines):
        line = line.strip()
        # "新增: N" or "新增：N"
        m = re.search(r'新增\s*[：:]\s*(\d+)', line)
        if m:
            return int(m.group(1))
        # "Inserted: N" or "New: N" (英文格式)
        m = re.search(r'(?:Inserted|New)\s*[：:]\s*(\d+)', line, re.IGNORECASE)
        if m:
            return int(m.group(1))
        # "新增 N 条" or "新增N条" (无冒号格式)
        m = re.search(r'新增\s*(\d+)\s*条', line)
        if m:
            return int(m.group(1))
        # "保存 N 条"
        m = re.search(r'保存\D*(\d+)', line)
        if m:
            return int(m.group(1))
    return 0

def get_status_label(entry):
    """从 run_status.json entry 获取状态标签"""
    if not entry:
        return "新脚本"
    s = entry.get("status", "")
    if s.startswith("成功"):
        return "正常"
    return s

def get_error_type(entry):
    """获取错误类型"""
    if not entry:
        return ""
    return entry.get("error_type", "")

# ==================== 分组逻辑 ====================
def group_crawlers(crawlers, status_map):
    """
    将爬虫分组:
      - normal: 上次成功 or 从未运行
      - timeout: 超时
      - site_down: 站点不可达类
      - script_error: 脚本异常
      - other_failed: 其他失败
      - disabled: 未启用
    """
    groups = {
        "正常": [],        # 上次成功
        "未运行": [],      # 从未运行
        "超时": [],
        "站点不可达": [],  # connection_refused, dns_resolve, unreachable, read_timeout
        "脚本异常": [],    # script_error
        "其他失败": [],    # exit_error, 其他
        "未启用": [],      # disabled
    }
    
    for c in crawlers:
        name = cname(c)
        script = c.get("script", "")
        enabled = c.get("enabled", True)
        
        if not enabled:
            groups["未启用"].append(c)
            continue
        
        entry = status_map.get(script) or status_map.get(name)
        
        if not entry:
            groups["未运行"].append(c)
            continue
        
        s = entry.get("status", "")
        error_type = entry.get("error_type", "")
        
        if s.startswith("成功"):
            groups["正常"].append(c)
        elif error_type == "timeout":
            groups["超时"].append(c)
        elif error_type in ("connection_refused", "dns_resolve", "unreachable", "read_timeout"):
            groups["站点不可达"].append(c)
        elif error_type == "script_error":
            groups["脚本异常"].append(c)
        else:
            groups["其他失败"].append(c)
    
    return groups

def print_group_report(groups, indent="  "):
    """打印分组报告"""
    total = sum(len(v) for k, v in groups.items())
    print(f"\n{'='*60}")
    print(f"  智能调度分组报告 — 总计 {total} 个爬虫")
    print(f"{'='*60}")
    
    order = ["正常", "未运行", "超时", "站点不可达", "脚本异常", "其他失败", "未启用"]
    emojis = {
        "正常": "✅", "未运行": "🆕", "超时": "⏰",
        "站点不可达": "🌐", "脚本异常": "🐛", "其他失败": "❓", "未启用": "⏹️"
    }
    
    for grp_name in order:
        items = groups.get(grp_name, [])
        if not items:
            continue
        emoji = emojis.get(grp_name, "📋")
        print(f"\n{indent}{emoji} {grp_name} ({len(items)}个)")
        # 按省份分组显示
        by_province = defaultdict(list)
        for c in items:
            by_province[c.get("group", "其他")].append(cname(c))
        for prov, names in sorted(by_province.items()):
            print(f"{indent}  [{prov}] {len(names)}个")
            if grp_name not in ("正常", "未运行", "未启用"):
                for n in names[:5]:
                    print(f"{indent}    · {n}")
                if len(names) > 5:
                    print(f"{indent}    ... 还有 {len(names)-5} 个")
    
    # 总结行 — 全部执行，不跳过
    print(f"\n{indent}{'='*50}")
    print(f"{indent}  全部 {total} 个均会执行（含之前失败的，站点恢复后自动归队）")
    return total, 0

# ==================== 执行器 ====================
def execute_worker(crawler):
    """单个worker执行任务，返回结果dict"""
    script = crawler["script"]
    args = crawler.get("args", "")
    name = cname(crawler)
    group_name = crawler.get("group", "")
    task_timeout = crawler.get("timeout") or SCRIPT_TIMEOUT
    
    log(f"▶ {name} ({script}) [timeout={task_timeout}s]")
    out, code, elapsed = run_script(script, args, timeout=task_timeout)
    
    if code == 0:
        new_count = extract_new_count(out)
        log(f"  ✅ 完成 +{new_count}条 [{elapsed}s]")
        return {
            "name": name, "script": script, "success": True,
            "new_count": new_count, "elapsed": elapsed,
            "error_type": "", "error": "",
            "group_name": group_name,
        }
    else:
        error_type = classify_error(out, code)
        log(f"  ❌ {error_type} [{elapsed}s]")
        return {
            "name": name, "script": script, "success": False,
            "new_count": 0, "elapsed": elapsed,
            "error_type": error_type, "error": out[:300],
            "group_name": group_name,
        }

def run_group(crawlers, label, parallel=DEFAULT_PARALLEL):
    """并行执行一组爬虫"""
    if not crawlers:
        return []
    
    n = len(crawlers)
    log(f"\n{'='*50}")
    log(f"  执行「{label}」组: {n} 个爬虫 (并行 {parallel})")
    log(f"{'='*50}")
    
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=parallel) as executor:
        futures = {executor.submit(execute_worker, c): c for c in crawlers}
        for future in concurrent.futures.as_completed(futures):
            try:
                result = future.result()
                results.append(result)
            except Exception as e:
                c = futures[future]
                results.append({
                    "name": cname(c), "script": c["script"],
                    "success": False, "new_count": 0, "elapsed": 0,
                    "error_type": "exception", "error": str(e)[:300],
                    "group_name": c.get("group", ""),
                })
    
    success = sum(1 for r in results if r["success"])
    failed = n - success
    log(f"  📊 {label}: 成功 {success} / 失败 {failed}")
    return results

# ==================== 主流程 ====================
def main():
    import argparse
    parser = argparse.ArgumentParser(description="智能爬虫调度器 v2")
    parser.add_argument("--dry-run", action="store_true", help="仅分组预览，不执行")
    parser.add_argument("--group", help="只运行指定省份/分组")
    parser.add_argument("--parallel", type=int, default=DEFAULT_PARALLEL, help=f"并行数 (默认 {DEFAULT_PARALLEL})")
    args = parser.parse_args()
    
    start_ts = time.time()
    log("🚀 爬虫调度器 v2 启动")
    
    crawlers = load_config()
    if args.group:
        crawlers = [c for c in crawlers if c.get("group", "") == args.group]
        log(f"过滤分组: {args.group} → {len(crawlers)} 个")
    
    status_map = load_status()
    groups = group_crawlers(crawlers, status_map)
    
    run_count, skip_count = print_group_report(groups)
    
    if args.dry_run:
        log("\n🏁 Dry-run 模式，未执行任何脚本")
        return
    
    # === 全部执行（含之前失败的）===
    all_results = []
    recovery_count = 0  # 上次失败本次成功的脚本数
    
    execution_order = ["正常", "未运行", "超时", "站点不可达", "脚本异常", "其他失败"]
    for grp_name in execution_order:
        items = groups.get(grp_name, [])
        if not items:
            continue
        results = run_group(items, grp_name, args.parallel)
        all_results.extend(results)
        # 统计恢复情况
        for r in results:
            if r["success"]:
                update_status(status_map, r)
            else:
                # 失败也更新状态（时间戳刷新）
                update_status(status_map, r)
        save_status(status_map)
    
    # 统计
    total_success = sum(1 for r in all_results if r["success"])
    total_failed = len(all_results) - total_success
    total_new = sum(r["new_count"] for r in all_results if r["success"])
    
    # 输出最终报告
    elapsed_total = time.time() - start_ts
    print(f"\n{'='*60}")
    print(f"  📊 今日运行总结")
    print(f"{'='*60}")
    print(f"  总耗时: {elapsed_total:.0f}s")
    print(f"  全部执行: {len(all_results)}/{len(crawlers)} 个 (未启用 {len(groups.get('未启用',[]))} 个跳过)")
    print(f"  今日成功: {total_success} 个 (+{total_new}条数据)")
    print(f"  今日失败: {total_failed} 个")
    
    # 失败明细
    if total_failed > 0:
        print(f"\n  📋 今日失败详情:")
        failed_by_type = defaultdict(list)
        for r in all_results:
            if not r["success"]:
                failed_by_type[r.get("error_type", "unknown")].append(r["name"])
        for etype, names in sorted(failed_by_type.items()):
            print(f"\n    {etype} ({len(names)}个):")
            for n in names[:8]:
                print(f"      · {n}")
            if len(names) > 8:
                print(f"      ... 还有 {len(names)-8} 个")
    
    if total_failed > 0:
        print(f"\n  下次日跑会再次尝试以上所有失败站点（站点恢复后自动归队）")
    
    # 固化最终状态
    save_status(status_map)
    log(f"🏁 调度完成")

def update_status(status_map, result):
    """更新单条状态到 status_map, 同时写入 run_logs"""
    ts = time.time()
    script = result.get("script", "")
    name = result.get("name", "")
    group_name = result.get("group_name", "")
    
    if result["success"]:
        status_str = f"成功 +{result['new_count']}条"
    else:
        et = result.get("error_type", "")
        if et == "timeout":
            status_str = "超时"
        elif et in ("connection_refused", "dns_resolve", "unreachable", "read_timeout"):
            status_str = "站点不可达"
        elif et == "script_error":
            status_str = "脚本异常"
        elif et == "exit_error":
            status_str = f"失败 rc={result.get('error','?')}"
        else:
            status_str = "失败"
    
    entry = {
        "status": status_str,
        "ts": ts,
        "elapsed": result.get("elapsed", 0),
        "error_type": result.get("error_type", ""),
    }
    if script:
        status_map[script] = entry
    if name and name != script:
        status_map[name] = entry
    
    # 写入 DB run_logs — 使用 config_name 作为标识
    db_script = script or name
    db_status = "成功" if result["success"] else status_str.split(" ")[0]
    write_run_log(
        script_name=db_script,
        config_name=name,
        group_name=group_name,
        status=db_status,
        new_count=result.get("new_count", 0),
        elapsed=result.get("elapsed", 0),
        error_detail=result.get("error", ""),
    )

if __name__ == "__main__":
    main()