#!/usr/bin/env python3
"""应用宝微信二维码登录（纯 Python，无浏览器自动化）。

二维码由微信 OAuth 生成，用户用微信扫码确认；本脚本轮询授权结果，完成
YYB OAuth 回调后将动态 cookies 合并到 mall-session.json，供后续纯 CK 流程使用。
"""

from __future__ import annotations

import argparse
import gzip
import json
import os
import re
import tempfile
import time
from dataclasses import dataclass
from http.cookiejar import CookieJar
from pathlib import Path
from urllib.error import HTTPError
from urllib.parse import urlencode
from urllib.request import (
    HTTPCookieProcessor,
    HTTPRedirectHandler,
    Request,
    build_opener,
)

ROOT = Path(__file__).resolve().parent.parent
OPEN_APPID = "wxd44977328b36e647"
CALLBACK = "https://yybadaccess.3g.qq.com/pc_yyb/pcyyb_oauth?login_type=WX"
OPEN_QR = "https://open.weixin.qq.com/connect/qrconnect"
POLL_QR = "https://lp.open.weixin.qq.com/connect/l/qrconnect"
USER_INFO = "https://yybadaccess.3g.qq.com/pc_yyb/pcyyb_get_user_info"
HREF = "data:text/css;base64,Ci5pbXBvd2VyQm94IC5xcmNvZGUge3dpZHRoOiAxNjBweDttYXJnaW4tdG9wOjI1cHh9"
USER_AGENT = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/150.0.0.0 Safari/537.36"
)


class NoRedirect(HTTPRedirectHandler):
    """保留 OAuth 回调的 302 和 Set-Cookie。"""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


@dataclass
class Response:
    status: int
    body: bytes

    @property
    def text(self) -> str:
        try:
            return gzip.decompress(self.body).decode("utf-8", "replace")
        except OSError:
            return self.body.decode("utf-8", "replace")


class Client:
    def __init__(self) -> None:
        self.jar = CookieJar()
        self.opener = build_opener(NoRedirect, HTTPCookieProcessor(self.jar))

    def request(
        self, url: str, *, referer: str = "", headers: dict[str, str] | None = None
    ) -> Response:
        request_headers = {"User-Agent": USER_AGENT, "Accept": "*/*"}
        if referer:
            request_headers["Referer"] = referer
        request_headers.update(headers or {})
        request = Request(url, headers=request_headers, method="GET")
        try:
            response = self.opener.open(request, timeout=30)
        except HTTPError as error:
            response = error
        return Response(response.status, response.read())

    def cookies(self) -> dict[str, str]:
        return {cookie.name: cookie.value for cookie in self.jar}


def extract_uuid(page: str) -> str:
    for pattern in (
        r"var\s+G\s*=\s*['\"]([A-Za-z0-9_-]{8,64})['\"]",
        r"connect/qrcode/([A-Za-z0-9_-]{8,64})",
        r"uuid=([A-Za-z0-9_-]{8,64})",
    ):
        match = re.search(pattern, page, re.IGNORECASE)
        if match:
            return match.group(1)
    raise ValueError("微信授权页未找到二维码 UUID")


def parse_poll(body: str) -> tuple[int, str]:
    errcode = re.search(r"wx_errcode\s*=\s*(-?\d+)", body)
    if not errcode:
        raise ValueError("二维码轮询响应缺少 wx_errcode")
    code = re.search(r"wx_code\s*=\s*['\"]([^'\"]*)['\"]", body)
    return int(errcode.group(1)), code.group(1) if code else ""


def login_type_header(value: str) -> str:
    try:
        return {"MOBILEQ": "1", "WX": "2"}[value]
    except KeyError as exc:
        raise ValueError(f"未知应用宝登录类型: {value!r}") from exc


def write_session(path: Path, cookies: dict[str, str]) -> None:
    """合并登录结果，保留 mall 的变换数据与用户已有配置。"""
    document: dict = {}
    if path.exists():
        try:
            document = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError as exc:
            raise ValueError(f"会话文件不是有效 JSON: {path}") from exc
    prior = document.get("cookies", {})
    if not isinstance(prior, dict):
        prior = {}
    document["cookies"] = {**prior, **cookies}
    document["login_type"] = cookies.get("logintype", "WX")
    document["login_updated_at"] = int(time.time())
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(
        prefix=".mall-session-", suffix=".tmp", dir=path.parent
    )
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            json.dump(document, handle, ensure_ascii=False, indent=2)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temporary, 0o600)
        os.replace(temporary, path)
        os.chmod(path, 0o600)
    except BaseException:
        Path(temporary).unlink(missing_ok=True)
        raise


def main() -> int:
    parser = argparse.ArgumentParser(description="应用宝微信扫码登录（纯 Python）")
    parser.add_argument(
        "--session", type=Path, default=ROOT / "config/mall-session.json"
    )
    parser.add_argument("--qr", type=Path, default=ROOT / "config/wechat-login.jpg")
    parser.add_argument("--timeout", type=float, default=300, help="二维码等待秒数")
    parser.add_argument("--interval", type=float, default=2, help="轮询间隔秒数")
    args = parser.parse_args()
    if args.timeout <= 0 or args.interval <= 0:
        raise ValueError("--timeout 和 --interval 必须为正数")

    client = Client()
    state = f"{time.time():.6f}"
    query = {
        "appid": OPEN_APPID,
        "fast_login": "0",
        "href": HREF,
        "redirect_uri": CALLBACK,
        "response_type": "code",
        "scope": "snsapi_login,snsapi_runtime_pcsdk",
        "self_redirect": "true",
        "state": state,
    }
    authorization_url = f"{OPEN_QR}?{urlencode(query)}"
    page = client.request(authorization_url, referer="https://m.yyb.qq.com/")
    if page.status != 200:
        raise RuntimeError(f"微信授权页请求失败: HTTP {page.status}")
    uuid = extract_uuid(page.text)
    image = client.request(
        f"https://open.weixin.qq.com/connect/qrcode/{uuid}", referer=authorization_url
    )
    if image.status != 200 or not image.body:
        raise RuntimeError(f"微信二维码请求失败: HTTP {image.status}")
    args.qr.parent.mkdir(parents=True, exist_ok=True)
    args.qr.write_bytes(image.body)
    print(f"微信登录二维码: {args.qr}")
    print("请用微信扫码并在手机确认；本终端将自动继续。")

    deadline = time.monotonic() + args.timeout
    last = ""
    code = ""
    while time.monotonic() < deadline:
        poll_query = {"uuid": uuid}
        if last:
            poll_query["last"] = last
        poll = client.request(
            f"{POLL_QR}?{urlencode(poll_query)}", referer=authorization_url
        )
        errcode, code = parse_poll(poll.text)
        if errcode == 405 and code:
            break
        if errcode == 405:
            # 已扫码、手机待确认：继续轮询，不能立即退出（否则 code 为空导致登录失败）。
            last = "405"
        elif errcode == 402:
            raise RuntimeError("二维码已过期，请重新执行登录")
        elif errcode == 403:
            raise RuntimeError("用户取消了扫码登录")
        else:
            last = "404"
        time.sleep(args.interval)
    else:
        raise TimeoutError("二维码轮询超时")
    if not code:
        raise RuntimeError("扫码成功响应缺少 OAuth code")

    callback_url = f"{CALLBACK}&{urlencode({'code': code, 'state': state})}"
    callback = client.request(callback_url, referer=authorization_url)
    if callback.status not in (200, 302):
        raise RuntimeError(f"YYB OAuth 回调失败: HTTP {callback.status}")
    cookies = client.cookies()
    openid = cookies.get("openid", "")
    access_token = cookies.get("accesstoken", "")
    login_type = cookies.get("logintype", "WX")
    if not openid or not access_token:
        raise RuntimeError("YYB OAuth 回调未写入 openid/accesstoken")
    info = client.request(
        USER_INFO,
        headers={
            "Ual-Access-Login-Type": login_type_header(login_type),
            "Ual-Access-Access-Token": access_token,
            "Ual-Access-Openid": openid,
            "Origin": "https://m.yyb.qq.com",
            "Referer": "https://m.yyb.qq.com/",
        },
    )
    if info.status != 200:
        raise RuntimeError(f"YYB 登录态校验失败: HTTP {info.status}")
    try:
        value = json.loads(info.text)
    except json.JSONDecodeError as exc:
        raise RuntimeError("YYB 登录态校验返回非 JSON") from exc
    if isinstance(value, dict) and value.get("ret") not in (None, 0):
        raise RuntimeError(f"YYB 登录态校验失败: ret={value.get('ret')}")
    write_session(args.session, cookies)
    print(f"登录成功，cookies 已写入: {args.session}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, TimeoutError, ValueError, RuntimeError) as error:
        print(f"登录失败: {error}")
        raise SystemExit(1) from error