"""
九宫格验证码完整流程
① jfbym type 10 OCR雪花屏 → ② 截图九宫格 → ③ jfbym 30223 返回坐标 → ④ 点击
"""
import os, time, base64, json, random
import cv2, numpy as np
import uiautomator2 as u2
import requests
from PIL import Image

# ===== 配置 =====
DEVICE = "O7R4Y9CMPBPBU4VK"
JFBYM_TOKEN = "1nDVocTE2mJ0yLEYb2sZJ5uUY2VIEoGTkIpW44X7Kgk"
CAPTCHA_CROP = (42, 424, 1180, 575)           # 雪花屏裁剪
GRID_CROP = (38, 662, 1185, 1817)              # 九宫格裁剪
SHOT_COUNT = 20                                # 连拍张数
AREA_THRESHOLD = 100                           # 连通域面积阈值
JFBYM_URL = "https://api.jfbym.com/api/YmServer/customApi"
OUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "image", "success")
# =================

def solve_snow_captcha():
    """步骤①: 本地降噪 + jfbym type 10 OCR → 返回 (文字, extra, 设备)"""
    d = u2.connect(DEVICE)

    # 连拍 + 平均降噪
    imgs = []
    for _ in range(SHOT_COUNT):
        full = d.screenshot(format='pillow')
        imgs.append(np.array(full.crop(CAPTCHA_CROP).convert('L')))
        time.sleep(0.05)

    avg = np.mean(imgs, axis=0).astype(np.uint8)
    t = np.percentile(avg, 15)
    dark = np.where(avg < t, 0, 255).astype(np.uint8)

    _, thresh = cv2.threshold(dark, 127, 255, cv2.THRESH_BINARY_INV)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(thresh, connectivity=8)
    clean = np.full_like(dark, 255)
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] > AREA_THRESHOLD:
            clean[labels == i] = 0

    os.makedirs(OUT_DIR, exist_ok=True)
    cleaned_path = os.path.join(OUT_DIR, "cleaned.png")
    cv2.imwrite(cleaned_path, clean)

    # 发给 jfbym type 10 OCR
    with open(cleaned_path, 'rb') as f:
        b64 = base64.b64encode(f.read()).decode()

    resp = requests.post(JFBYM_URL, json={
        "token": JFBYM_TOKEN, "type": "10", "image": b64
    }, headers={"Content-Type": "application/json"}, timeout=30).json()
    print(f"  [jfbym-type10] 返回: {json.dumps(resp, ensure_ascii=False)[:200]}")

    extra = resp.get("data", {}).get("data", {})
    if isinstance(extra, list):
        extra = extra[0] if extra else {}
    if isinstance(extra, dict):
        text = extra.get("tips", "")
    elif isinstance(extra, str):
        text = extra
    else:
        text = ""
    print(f"[1] 降噪+OCR: {text}")
    return text, extra, d


def get_click_pos(d, extra):
    """步骤②③: 截图九宫格 → jfbym 30223 → 返回坐标"""
    grid_img = d.screenshot(format='pillow').crop(GRID_CROP)
    grid_path = os.path.join(OUT_DIR, "grid.png")
    grid_img.save(grid_path)
    print(f"[2] 九宫格: {grid_img.size}")

    with open(grid_path, 'rb') as f:
        grid_b64 = base64.b64encode(f.read()).decode()

    resp = requests.post(JFBYM_URL, json={
        "token": JFBYM_TOKEN,
        "type": "30223",
        "image": grid_b64,
        "extra": extra
    }, headers={"Content-Type": "application/json"}, timeout=30).json()
    print(f"  [jfbym-30223] 返回: {json.dumps(resp, ensure_ascii=False)[:200]}")

    resp_data = resp.get("data", {})
    if isinstance(resp_data, list):
        data = resp_data[0] if resp_data else {}
    else:
        data = resp_data.get("data", {})
    click_pos = data.get("click_pos", [])
    tips = data.get("tips", "")
    print(f"[3] tips={tips}, click_pos={click_pos}")
    return click_pos, tips


SUBMIT_BTN = (602, 2018)                        # 提交按钮坐标
REFRESH_WAIT = 1.5                             # 点击后等待刷新秒数(加长)
MAX_ROUNDS = 15                                 # 最大轮数, 防止死循环
CLICK_OFFSET = 18                               # 随机偏移范围（±18px）
GRID_W, GRID_H = GRID_CROP[2] - GRID_CROP[0], GRID_CROP[3] - GRID_CROP[1]  # 九宫格宽高
STALL_LIMIT = 2                                 # 连续N轮候选完全不变→判定卡滞, 提前提交
CAPTCHA_KEYWORDS = ("请依次点击", "根据提示", "没有新图片", "提交", "验证失败", "验证码错误")

_ocr_eng = None


def _captcha_still_present(d):
    """截图 + OCR 检测九宫格验证码特征词是否仍在页面上"""
    from rapidocr_onnxruntime import RapidOCR
    global _ocr_eng
    if _ocr_eng is None:
        _ocr_eng = RapidOCR()
    shot = d.screenshot(format='opencv')
    if shot is None:
        return None
    r = _ocr_eng(shot)
    if not r or not r[0]:
        return False
    texts = [item[1] for item in r[0]]
    return any(any(kw in t for kw in CAPTCHA_KEYWORDS) for t in texts)


def tap_loop(d, extra):
    """步骤④: 每次只点一个→等待刷新→重新截图识别→直到无匹配→提交
    返回 True/False: 提交后延迟一段时间再 OCR 复检验证码特征词是否消失, 以此判定真实成败
    (而不是把"识别到的提示文字"当作成败信号——jfbym的tips有时是纯字符串, 会被误判为falsy)。"""
    prev_sig = None
    stall_count = 0
    submitted = False

    for round_num in range(1, MAX_ROUNDS + 1):
        pos, tips = get_click_pos(d, extra)
        if not pos:
            print(f"[4] 第{round_num}轮无匹配，点击提交按钮")
            d.click(SUBMIT_BTN[0] + random.randint(-8, 8), SUBMIT_BTN[1] + random.randint(-5, 5))
            submitted = True
            break

        sig = (tips, tuple(sorted(pos)))
        if sig == prev_sig:
            stall_count += 1
        else:
            stall_count = 0
        prev_sig = sig
        if stall_count >= STALL_LIMIT:
            print(f"[4] 连续{stall_count + 1}轮候选完全未变化，判定检测/点击卡滞，直接提交")
            d.click(SUBMIT_BTN[0] + random.randint(-8, 8), SUBMIT_BTN[1] + random.randint(-5, 5))
            submitted = True
            break

        # 只点第一个，点完重新截图
        x, y = pos[0]
        ox = x + random.randint(-CLICK_OFFSET, CLICK_OFFSET)
        oy = y + random.randint(-CLICK_OFFSET, CLICK_OFFSET)
        ox = max(5, min(GRID_W - 5, ox))
        oy = max(5, min(GRID_H - 5, oy))
        print(f"[4] 第{round_num}轮 共{len(pos)}个匹配，先点({x},{y})→偏移({ox},{oy})")
        d.click(GRID_CROP[0] + ox, GRID_CROP[1] + oy)
        time.sleep(REFRESH_WAIT)
    else:
        print(f"[4] 超过{MAX_ROUNDS}轮，直接提交")
        d.click(SUBMIT_BTN[0] + random.randint(-8, 8), SUBMIT_BTN[1] + random.randint(-5, 5))
        submitted = True

    if not submitted:
        return False

    # 提交后不能立刻检测(页面还没刷新完成), 等一段随机间隔再 OCR 复检
    time.sleep(random.uniform(1.8, 2.6))
    still_present = _captcha_still_present(d)
    if still_present is None:
        print("[4] 提交后复检失败(截图/OCR异常)，保守判定为未通过")
        return False
    if still_present:
        print("[4] 提交后复检: 验证码特征词仍在，判定未通过")
        return False
    print("[4] 提交后复检: 验证码特征词已消失，判定通过")
    return True


def solve(driver=None):
    """九宫格验证码求解, 可传入外部driver或自动连接"""
    if driver is not None:
        d = driver
    else:
        d = u2.connect(DEVICE)

    if not JFBYM_TOKEN:
        raise ValueError("请先设置 JFBYM_TOKEN")
    text, extra, _ = solve_snow_captcha_with_driver(d)
    return tap_loop(d, extra)

def solve_snow_captcha_with_driver(d):
    """步骤①: 本地降噪 + jfbym type 10 OCR → 返回 (文字, extra, 设备)"""
    # 连拍 + 平均降噪
    imgs = []
    for _ in range(SHOT_COUNT):
        full = d.screenshot(format='pillow')
        imgs.append(np.array(full.crop(CAPTCHA_CROP).convert('L')))
        time.sleep(0.05)

    avg = np.mean(imgs, axis=0).astype(np.uint8)
    t = np.percentile(avg, 15)
    dark = np.where(avg < t, 0, 255).astype(np.uint8)

    _, thresh = cv2.threshold(dark, 127, 255, cv2.THRESH_BINARY_INV)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(thresh, connectivity=8)
    clean = np.full_like(dark, 255)
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] > AREA_THRESHOLD:
            clean[labels == i] = 0

    os.makedirs(OUT_DIR, exist_ok=True)
    cleaned_path = os.path.join(OUT_DIR, "cleaned.png")
    cv2.imwrite(cleaned_path, clean)

    # 发给 jfbym type 10 OCR
    with open(cleaned_path, 'rb') as f:
        b64 = base64.b64encode(f.read()).decode()

    resp = requests.post(JFBYM_URL, json={
        "token": JFBYM_TOKEN, "type": "10", "image": b64
    }, headers={"Content-Type": "application/json"}, timeout=30).json()
    print(f"  [jfbym-type10] 返回: {json.dumps(resp, ensure_ascii=False)[:200]}")

    extra = resp.get("data", {}).get("data", {})
    if isinstance(extra, list):
        extra = extra[0] if extra else {}
    text = extra.get("tips", "") if isinstance(extra, dict) else ""
    print(f"[1] 降噪+OCR: {text}")
    return text, extra, d

if __name__ == '__main__':
    print(f"设备: {DEVICE}\n")
    print(f"结果: {solve()}")