"""
GEO 持续采集 Worker v2 — 多进程架构

每个平台独立进程池, 主进程拉任务→分发→收结果→提交。
各子进程独立加载 Client / WASM / cookie, 无 GIL 竞争。

用法:
  python main.py
  python main.py --config worker_config.json
  python main.py --once
  python main.py --dry-run
"""

import argparse
import json
import os
import re
import signal
import sys
import time
from datetime import datetime
from multiprocessing import Process, Queue, Event
from typing import Any

sys.stdout.reconfigure(encoding="utf-8")

import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

from config import (
    load_config, normalize_platform_name, API_TO_INTERNAL,
    SUPPORTED_API_PLATFORMS, _parse_bool, _default_switches,
)
from api_client import GeoApiClient
from payload import build_submit_payload
from worker import WORKER_FUNCS

_stop_event = Event()


def is_stopping():
    return _stop_event.is_set()


def _handle_signal(sig, frame):
    if _stop_event.is_set():
        print("\n[FORCE] 二次 Ctrl+C, 立即退出")
        os._exit(1)
    _stop_event.set()
    print("\n[stop] 收到停止信号, 完成当前任务后退出...")


# ---------- 平台开关 ----------
def _merge_switches(base, override):
    if not override or not isinstance(override, dict):
        return dict(base)
    merged = dict(base)
    for k, v in override.items():
        plat = normalize_platform_name(str(k))
        if plat and plat in merged:
            try:
                merged[plat] = _parse_bool(v)
            except ValueError:
                pass
    return merged


class PlatformSwitchBoard:
    def __init__(self, switches, hot_reload=True, config_path=""):
        self._switches = dict(switches)
        self.hot_reload = hot_reload
        self.config_path = config_path
        self._last_mtime = 0.0

    def is_enabled(self, platform):
        plat = normalize_platform_name(platform)
        return bool(self._switches.get(plat))

    def enabled_list(self):
        return [p for p in SUPPORTED_API_PLATFORMS if self._switches.get(p)]

    def query_platforms(self):
        return ",".join(self.enabled_list())

    def summary(self):
        on = [p for p in SUPPORTED_API_PLATFORMS if self._switches.get(p)]
        off = [p for p in SUPPORTED_API_PLATFORMS if not self._switches.get(p)]
        return f"开启=[{','.join(on) or '-'}] 关闭=[{','.join(off) or '-'}]"

    def maybe_reload(self) -> bool:
        if not self.hot_reload or not os.path.exists(self.config_path):
            return False
        try:
            mtime = os.path.getmtime(self.config_path)
            if mtime <= self._last_mtime:
                return False
            self._last_mtime = mtime
            with open(self.config_path, "r", encoding="utf-8") as f:
                file_cfg = json.load(f)
            new_switches = _merge_switches(_default_switches(), file_cfg.get("platform_switches"))
            if new_switches != self._switches:
                old = self.summary()
                self._switches = new_switches
                print(f"[hot-reload] 平台开关更新: {old} -> {self.summary()}")
                return True
        except Exception as e:
            print(f"[hot-reload] 失败: {e}")
        return False


# ---------- 工具 ----------
FAILED_DIR = ""


def _save_failed(stage, task, detail):
    os.makedirs(FAILED_DIR, exist_ok=True)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    did = task.get("dispatchId") or task.get("id") or "unknown"
    safe_did = re.sub(r'[\\/:*?"<>|]', '_', str(did))
    path = os.path.join(FAILED_DIR, f"{ts}_{safe_did}_{stage}.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump({"stage": stage, "task": task, "detail": detail}, f, ensure_ascii=False, indent=2)
    print(f"  [saved] {path}")


# ---------- 主流程 ----------
def run_worker(config_path=None, overrides=None):
    global FAILED_DIR

    cfg = load_config(config_path)
    if overrides:
        cfg.update(overrides)
    base_dir = os.path.dirname(cfg["_config_path"])
    FAILED_DIR = os.path.join(base_dir, "worker_failed")
    os.makedirs(FAILED_DIR, exist_ok=True)

    if not cfg.get("spider_token"):
        print(f"[FATAL] 请在配置文件填写 spider_token:\n  {cfg['_config_path']}")
        return

    switches = _merge_switches(_default_switches(), cfg.get("platform_switches"))
    board = PlatformSwitchBoard(switches, hot_reload=bool(cfg.get("hot_reload", True)), config_path=cfg["_config_path"])

    once = bool(cfg.get("once"))
    dry_run = bool(cfg.get("dry_run"))
    idle_sleep = float(cfg.get("idle_sleep") or 30)
    enable_commit = bool(cfg.get("enable_task_commit", True))
    location = cfg.get("location") or ""
    conc = cfg.get("concurrency") or {}

    # 按平台创建队列和进程 (每个平台 1 个进程, 内部按 concurrency 并发)
    platform_queues = {}
    result_queue = Queue()
    workers = []

    enabled = board.enabled_list()
    conc_str_parts = []
    for plat in ("deepseek", "kimi", "tongyi", "baiduai", "doubao"):
        n = max(1, int(conc.get(plat, 1)))
        conc_str_parts.append(f"{plat}={n}")
        if plat not in enabled:
            continue
        if plat not in WORKER_FUNCS:
            continue
        q = Queue()
        platform_queues[plat] = q
        p = Process(
            target=WORKER_FUNCS[plat],
            args=(0, q, result_queue, _stop_event, cfg),
            daemon=True,
        )
        workers.append((plat, 0, p))

    print("=" * 60)
    print("GEO Task Worker v2 (多进程)")
    print(f"  配置文件 : {cfg['_config_path']}")
    print(f"  API      : {cfg['api_base']}")
    print(f"  拉题平台 : {board.query_platforms() or '(无)'}")
    print(f"  启停控制 : {board.summary()}")
    print(f"  并发(进程): {' | '.join(conc_str_parts)}")
    print(f"  count    : {cfg['batch_count']}")
    print(f"  idle     : {idle_sleep}s")
    print(f"  commit   : {enable_commit}")
    print(f"  dry_run  : {dry_run}")
    print(f"  once     : {once}")
    print(f"  进程总数 : {len(workers)}")
    print("=" * 60)

    if not platform_queues:
        print("无可用平台, 请在 worker_config.json 中开启至少一个平台")
        return

    # 启动进程
    for plat, wid, p in workers:
        p.start()
        print(f"  [{plat}-{wid}] started pid={p.pid}")

    geo = GeoApiClient(cfg["api_base"], cfg["spider_token"], cfg.get("vendor_token") or "")
    stats = {"pulled": 0, "ok": 0, "fail": 0, "skip": 0, "disabled": 0, "submitted": 0}

    def collect_results(block=False, timeout=0.1):
        """从 result_queue 中取所有可用结果并提交"""
        while True:
            try:
                r = result_queue.get(block=block, timeout=timeout)
            except Exception:
                break
            task = r["task"]
            did = task.get("dispatchId") or task.get("id") or "?"

            if not r["ok"]:
                stats["fail"] += 1
                err = r.get("error", "unknown")
                print(f"  [main] ✗ did={did} error={err}")
                _save_failed("collect", task, err)
                continue

            parsed = r["parsed"]
            payload = build_submit_payload(task, parsed, location=location)
            if not payload.get("id"):
                stats["fail"] += 1
                _save_failed("submit", task, "missing_id")
                continue

            if dry_run:
                print(f"  [main] [dry-run] did={did} 跳过提交")
                stats["ok"] += 1
                continue

            try:
                resp = geo.submit(payload)
                code = resp.get("code")
                print(f"  [main] [submit] did={did} code={code} msg={resp.get('msg')}")
                if code in (0, 200, "0"):
                    stats["ok"] += 1
                    stats["submitted"] += 1
                else:
                    stats["fail"] += 1
                    _save_failed("submit", task, resp)
            except Exception as e:
                stats["fail"] += 1
                print(f"  [main] [submit] did={did} 异常: {e}")
                _save_failed("submit", task, str(e))

    try:
        while not is_stopping():
            board.maybe_reload()
            pull_platforms = board.query_platforms()
            if not pull_platforms:
                print(f"[idle] 全部平台已关闭, {idle_sleep}s 后重试...")
                if once:
                    break
                time.sleep(idle_sleep)
                continue

            try:
                tasks = []
                for plat in pull_platforms.split(","):
                    if is_stopping():
                        break
                    plat = plat.strip()
                    if not plat:
                        continue
                    plat_tasks = geo.get_tasks(plat, cfg["batch_count"])
                    tasks.extend(plat_tasks)
            except PermissionError as e:
                print(f"[FATAL] {e}")
                break
            except Exception as e:
                if is_stopping():
                    break
                print(f"[get_task] 异常: {e}, {idle_sleep}s 后重试")
                time.sleep(idle_sleep)
                if once:
                    break
                continue

            if is_stopping():
                break

            if not tasks:
                collect_results()
                print(f"[idle] 无任务 platforms={pull_platforms}, {idle_sleep}s 后重试...")
                if once:
                    break
                time.sleep(idle_sleep)
                continue

            # 统计分布
            plat_counts = {}
            for t in tasks:
                p = (t.get("platform") or "unknown").strip().lower()
                plat_counts[p] = plat_counts.get(p, 0) + 1
            breakdown = " | ".join(f"{p}={n}" for p, n in sorted(plat_counts.items()))
            print(f"\n[pull] 取得 {len(tasks)} 条任务 (platforms={pull_platforms})")
            print(f"  分布: {breakdown}")

            # commit + 分发
            dispatched = 0
            for task in tasks:
                if is_stopping():
                    break
                stats["pulled"] += 1
                api_plat = (task.get("platform") or "").strip().lower()
                plat = normalize_platform_name(api_plat) or api_plat

                if API_TO_INTERNAL.get(plat) is None:
                    print(f"  skip 不支持平台 {api_plat} dispatchId={task.get('dispatchId')}")
                    stats["skip"] += 1
                    continue
                if not board.is_enabled(plat):
                    print(f"  skip 平台已停用 {plat}")
                    stats["disabled"] += 1
                    continue
                if plat not in platform_queues:
                    print(f"  skip 无 worker: {plat}")
                    stats["skip"] += 1
                    continue

                if enable_commit:
                    try:
                        resp = geo.task_commit(task)
                        code = resp.get("code")
                        print(f"  [commit] did={task.get('dispatchId')} code={code}")
                    except Exception as e:
                        print(f"  [commit] 异常: {e}")

                platform_queues[plat].put(task)
                dispatched += 1

            print(f"  已分发 {dispatched} 条任务到子进程")

            # 等待本轮结果 (非阻塞轮询, 直到所有分发的任务都返回)
            collected = 0
            wait_start = time.time()
            while collected < dispatched and not is_stopping():
                try:
                    r = result_queue.get(timeout=5)
                    collected += 1
                    task = r["task"]
                    did = task.get("dispatchId") or task.get("id") or "?"

                    if not r["ok"]:
                        stats["fail"] += 1
                        print(f"  [main] ✗ did={did} error={r.get('error', '?')}")
                        _save_failed("collect", task, r.get("error", "unknown"))
                        continue

                    parsed = r["parsed"]
                    payload = build_submit_payload(task, parsed, location=location)
                    if not payload.get("id"):
                        stats["fail"] += 1
                        _save_failed("submit", task, "missing_id")
                        continue

                    if dry_run:
                        print(f"  [main] [dry-run] did={did}")
                        stats["ok"] += 1
                        continue

                    try:
                        resp = geo.submit(payload)
                        code = resp.get("code")
                        print(f"  [main] [submit] did={did} code={code}")
                        if code in (0, 200, "0"):
                            stats["ok"] += 1
                            stats["submitted"] += 1
                        else:
                            stats["fail"] += 1
                            _save_failed("submit", task, resp)
                    except Exception as e:
                        stats["fail"] += 1
                        _save_failed("submit", task, str(e))
                except Exception:
                    if time.time() - wait_start > 600:
                        print("[main] 等待结果超时 (600s)")
                        break

            print(f"\n[stats] pulled={stats['pulled']} ok={stats['ok']} "
                  f"fail={stats['fail']} skip={stats['skip']} submitted={stats['submitted']}")
            print(f"[switch] {board.summary()}")
            if once or is_stopping():
                break

    finally:
        print("\n[main] 正在关闭子进程...")
        _stop_event.set()
        for plat, q in platform_queues.items():
            # 多发几个毒丸确保每个 worker 都能收到
            n = max(1, int(conc.get(plat, 1)))
            for _ in range(n):
                try:
                    q.put(None)
                except Exception:
                    pass
        for plat, wid, p in workers:
            p.join(timeout=30)
            if p.is_alive():
                print(f"  [{plat}-{wid}] 强制终止")
                p.terminate()
        print(f"\n[main] 结束 {stats}")


def main():
    parser = argparse.ArgumentParser(description="GEO 持续采集 Worker v2 (多进程)")
    parser.add_argument("--config", default=None, help="配置文件路径")
    parser.add_argument("--once", action="store_true", help="只跑一轮")
    parser.add_argument("--dry-run", action="store_true", help="不提交结果")
    args = parser.parse_args()

    signal.signal(signal.SIGINT, _handle_signal)
    signal.signal(signal.SIGTERM, _handle_signal)

    overrides = {}
    if args.once:
        overrides["once"] = True
    if args.dry_run:
        overrides["dry_run"] = True
    run_worker(args.config, overrides)


if __name__ == "__main__":
    main()