#!/usr/bin/env python3
"""
微信「联系客服解封」常规批量执行器

适用：以后任意设备、任意微信号被限制，编辑配置文件即可一键全自动解封。

链路：调用 SDK `POST /api/v3/account/unblock-customer-service`
    → 设备端 unblock_via_customer_service skill
    → 自动跑「安全保护中 → 联系客服 → 终止使用说明 → 联系专属客服 → 允许 → 微信安全专属客服 → AI 持续对话」

用法：
    python3 sdk/scripts/wechat_unblock_batch.py -c sdk/config/unblock_targets.yaml
    python3 sdk/scripts/wechat_unblock_batch.py -c xxx.yaml --dry-run    # 只校验不执行
    python3 sdk/scripts/wechat_unblock_batch.py -c xxx.yaml --only 192.168.110.80:5555

输出：
    sdk/tmp/unblock_batch_{ts}/
        ├── summary.md                  汇总表格
        ├── {device}_{wxid}.json        每条任务原始返回（含 turns/screenshots）
        └── {device}_{wxid}.log         每条任务日志
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sys
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

LOG = logging.getLogger("unblock_batch")


# ---------- 配置加载 ----------

def _load_yaml(path: str) -> dict:
    try:
        import yaml  # type: ignore
        with open(path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f) or {}
    except ImportError:
        # 退化：把简单 yaml 当 json-like 解析（仅支持本配置示例）
        import re
        content = open(path, "r", encoding="utf-8").read()
        content = re.sub(r"#.*$", "", content, flags=re.MULTILINE)
        # 极简退化：要求用户安装 PyYAML
        raise RuntimeError("未安装 PyYAML，请执行: pip install pyyaml") from None


@dataclass
class TargetTask:
    device_id: str
    wxid: str = ""
    phone: str = ""
    reason: str = ""
    max_rounds: int = 18
    chat_interval_sec: int = 8
    use_web_search: bool = True
    platform: str = "wechat"
    enabled: bool = True
    note: str = ""

    # 运行时
    started_at: float = 0
    finished_at: float = 0
    status: str = "pending"     # pending / running / success / reject / timeout / fallback / error / skipped
    final_message: str = ""
    rounds: int = 0
    raw_response: Optional[Dict[str, Any]] = field(default=None)


# ---------- 调用 SDK ----------

def _call_unblock(sdk_base: str, t: TargetTask, timeout: int) -> Dict[str, Any]:
    qs = urllib.parse.urlencode({
        "device_id": t.device_id,
        "platform": t.platform,
        "reason": t.reason,
        "phone": t.phone,
        "wxid": t.wxid,
        "max_rounds": t.max_rounds,
        "chat_interval_sec": t.chat_interval_sec,
        "use_web_search": "true" if t.use_web_search else "false",
    })
    url = f"{sdk_base.rstrip('/')}/api/v3/account/unblock-customer-service?{qs}"
    LOG.info(f"[{t.device_id}] POST {url}")
    req = urllib.request.Request(url, method="POST", headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = resp.read().decode("utf-8")
    try:
        return json.loads(body)
    except Exception:
        return {"code": 599, "raw_body": body[:2000]}


# ---------- 单条任务 ----------

def _run_one(sdk_base: str, t: TargetTask, out_dir: Path,
             retry_on_fallback: int) -> TargetTask:
    if not t.enabled:
        t.status = "skipped"
        LOG.info(f"[{t.device_id}] disabled -> skip")
        return t

    t.started_at = time.time()
    t.status = "running"
    log_path = out_dir / f"{_safe(t.device_id)}_{_safe(t.wxid or 'noid')}.log"
    json_path = out_dir / f"{_safe(t.device_id)}_{_safe(t.wxid or 'noid')}.json"

    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    LOG.addHandler(file_handler)

    attempts = 0
    response = None
    try:
        attempts += 1
        # 单条最长耗时：max_rounds * (chat_interval + 5) + 90 兜底
        timeout = t.max_rounds * (t.chat_interval_sec + 5) + 90
        response = _call_unblock(sdk_base, t, timeout=timeout)
        t.raw_response = response

        data = response.get("data", {}) if isinstance(response, dict) else {}
        session = data.get("session", {}) if isinstance(data, dict) else {}
        status = data.get("status") or session.get("final_status") or "unknown"
        msg = data.get("message") or data.get("error") or session.get("final_message") or ""
        rounds = session.get("rounds", 0)

        t.status = status
        t.final_message = msg[:300]
        t.rounds = int(rounds or 0)

        # 链路失败时按需重试
        while t.status == "fallback" and attempts <= retry_on_fallback:
            LOG.warning(f"[{t.device_id}] fallback，重试 {attempts}/{retry_on_fallback}")
            time.sleep(8)
            attempts += 1
            response = _call_unblock(sdk_base, t, timeout=timeout)
            t.raw_response = response
            data = response.get("data", {}) if isinstance(response, dict) else {}
            session = data.get("session", {}) if isinstance(data, dict) else {}
            t.status = data.get("status") or session.get("final_status") or "unknown"
            t.final_message = (data.get("message") or session.get("final_message") or "")[:300]
            t.rounds = int(session.get("rounds", 0) or 0)

    except Exception as exc:
        t.status = "error"
        t.final_message = f"{type(exc).__name__}: {exc}"[:300]
        LOG.exception(f"[{t.device_id}] 异常")
        response = {"error": str(exc)}
    finally:
        t.finished_at = time.time()
        try:
            json_path.write_text(
                json.dumps(response or {}, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
        except Exception:
            pass
        LOG.removeHandler(file_handler)
        file_handler.close()

    return t


def _safe(s: str) -> str:
    return "".join(c if c.isalnum() else "_" for c in s)[:48]


# ---------- 主流程 ----------

def main() -> int:
    parser = argparse.ArgumentParser(description="微信联系客服解封 · 批量执行器")
    parser.add_argument("-c", "--config", required=True, help="yaml 配置文件路径")
    parser.add_argument("--dry-run", action="store_true", help="只校验不执行")
    parser.add_argument("--only", help="只跑某个 device_id（覆盖配置）")
    parser.add_argument("--out", help="输出目录（默认 sdk/tmp/unblock_batch_{ts}）")
    parser.add_argument("--log-level", default="INFO")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s %(message)s",
    )

    cfg = _load_yaml(args.config)
    sdk_base = cfg.get("sdk_base") or os.environ.get("SDK_BASE_URL") or "http://127.0.0.1:8899"
    platform = cfg.get("platform", "wechat")
    default_max = int(cfg.get("default_max_rounds", 18))
    default_interval = int(cfg.get("default_chat_interval_sec", 8))
    default_web = bool(cfg.get("default_use_web_search", True))
    concurrency = max(1, int(cfg.get("concurrency", 1)))
    retry_on_fallback = int(cfg.get("retry_on_fallback", 1))
    between_sleep = int(cfg.get("between_task_sleep_sec", 6))

    raw_targets = cfg.get("targets") or []
    tasks: List[TargetTask] = []
    for raw in raw_targets:
        if not isinstance(raw, dict):
            continue
        dev = raw.get("device_id")
        if not dev:
            continue
        if args.only and dev != args.only:
            continue
        tasks.append(TargetTask(
            device_id=dev,
            wxid=raw.get("wxid", "") or "",
            phone=raw.get("phone", "") or "",
            reason=raw.get("reason", "") or "",
            max_rounds=int(raw.get("max_rounds", default_max)),
            chat_interval_sec=int(raw.get("chat_interval_sec", default_interval)),
            use_web_search=bool(raw.get("use_web_search", default_web)),
            platform=raw.get("platform", platform),
            enabled=bool(raw.get("enabled", True)),
            note=raw.get("note", "") or "",
        ))

    if not tasks:
        LOG.error("配置中未找到任何 target")
        return 2

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    if args.out:
        out_dir = Path(args.out)
    else:
        sdk_root = Path(__file__).resolve().parents[1]
        out_dir = sdk_root / "tmp" / f"unblock_batch_{ts}"
    out_dir.mkdir(parents=True, exist_ok=True)

    LOG.info(f"配置: {args.config}")
    LOG.info(f"SDK : {sdk_base}")
    LOG.info(f"任务数: {len(tasks)} (concurrency={concurrency})")
    LOG.info(f"输出: {out_dir}")

    if args.dry_run:
        LOG.info("--- DRY RUN ---")
        for t in tasks:
            LOG.info(f"  {t.device_id}  wxid={t.wxid}  reason={t.reason[:30]}  enabled={t.enabled}")
        return 0

    # 健康检查
    try:
        with urllib.request.urlopen(f"{sdk_base.rstrip('/')}/api/v3/connection/status", timeout=10) as resp:
            conn = json.loads(resp.read().decode())
        data = conn.get("data", {})
        LOG.info(f"SDK 健康 ws_online={data.get('online_ws_count')} adb={data.get('adb_count')}")
    except Exception as exc:
        LOG.error(f"SDK 不可达: {exc}")
        return 3

    # 执行（串行 / 并发）
    if concurrency <= 1:
        for t in tasks:
            _run_one(sdk_base, t, out_dir, retry_on_fallback)
            if between_sleep > 0:
                time.sleep(between_sleep)
    else:
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as ex:
            futures = {ex.submit(_run_one, sdk_base, t, out_dir, retry_on_fallback): t for t in tasks}
            for fut in concurrent.futures.as_completed(futures):
                t = futures[fut]
                try:
                    fut.result()
                except Exception as exc:
                    LOG.exception(f"[{t.device_id}] 调度异常: {exc}")

    # 汇总报告
    _write_summary(tasks, out_dir, sdk_base, ts)
    success = sum(1 for t in tasks if t.status == "success")
    print(f"\n=== 完成 {success}/{len(tasks)} ===\n报告: {out_dir / 'summary.md'}")
    return 0 if success == len(tasks) else 1


def _write_summary(tasks: List[TargetTask], out_dir: Path, sdk_base: str, ts: str) -> None:
    lines = [
        f"# 微信联系客服解封 · 批量执行报告 {ts}",
        "",
        f"- SDK: `{sdk_base}`",
        f"- 任务数: {len(tasks)}",
        f"- 成功: {sum(1 for t in tasks if t.status == 'success')}",
        f"- 拒绝: {sum(1 for t in tasks if t.status == 'reject')}",
        f"- 超时: {sum(1 for t in tasks if t.status == 'timeout')}",
        f"- 链路失败: {sum(1 for t in tasks if t.status == 'fallback')}",
        f"- 异常: {sum(1 for t in tasks if t.status == 'error')}",
        f"- 跳过: {sum(1 for t in tasks if t.status == 'skipped')}",
        "",
        "| # | device | wxid | 状态 | 轮次 | 耗时(s) | 客服最终回复 |",
        "|---|--------|------|------|------|---------|--------------|",
    ]
    for i, t in enumerate(tasks, 1):
        elapsed = round(t.finished_at - t.started_at, 1) if t.started_at else 0
        msg = (t.final_message or "").replace("|", "/").replace("\n", " ")[:80]
        lines.append(f"| {i} | `{t.device_id}` | `{t.wxid}` | **{t.status}** | {t.rounds} | {elapsed} | {msg} |")
    lines += [
        "",
        "## 状态说明",
        "- **success**：客服已解封 / 已恢复使用",
        "- **reject**：客服明确拒绝解封",
        "- **timeout**：达到最大轮次或客服长时间无回复",
        "- **fallback**：未能进入客服小程序（路径变化）",
        "- **error**：调用异常（网络 / 设备不在线等）",
        "- **skipped**：配置 enabled=false",
        "",
        "## 详细 JSON",
        "",
        "每条任务详见同目录 `{device_id}_{wxid}.json`（含 turns / screenshots / web_snippets）",
    ]
    (out_dir / "summary.md").write_text("\n".join(lines), encoding="utf-8")


if __name__ == "__main__":
    sys.exit(main())