"""WeChat login session for https://3d.hunyuan.tencent.com via Playwright.

CLI 友好：同步阻塞扫码登录（非后台线程）。
未登录 / 凭证失效 / 额度满换号 时弹出有头 Edge 供用户扫码。
"""
from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any, Callable, Optional

HUNYUAN_URL = "https://3d.hunyuan.tencent.com/"
LOGIN_TIMEOUT_SEC = 300

UpdateFn = Callable[..., None]


def _tools_root() -> Path:
    return Path(__file__).resolve().parent.parent


def session_dir() -> Path:
    override = os.environ.get("HUNYUAN_SESSION_DIR", "").strip()
    if override:
        return Path(override).expanduser().resolve()
    return _tools_root() / "data" / "hunyuan"


def storage_state_path() -> Path:
    return session_dir() / "storage_state.json"


def status_file_path() -> Path:
    return session_dir() / "status.json"


def jobs_dir() -> Path:
    path = session_dir() / "jobs"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _browser_channel() -> str:
    return os.environ.get("HUNYUAN_BROWSER_CHANNEL", "msedge").strip() or "msedge"


def _write_status(**values: Any) -> dict:
    directory = session_dir()
    directory.mkdir(parents=True, exist_ok=True)
    payload = status()
    payload.update(values, updated_at=time.time())
    temporary = status_file_path().with_suffix(".tmp")
    temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    temporary.replace(status_file_path())
    return payload


def status() -> dict:
    path = status_file_path()
    if path.is_file():
        try:
            data = json.loads(path.read_text("utf-8"))
            if storage_state_path().is_file() and data.get("state") not in {
                "waiting_scan",
                "error",
                "timeout",
            }:
                data["logged_in"] = True
            return data
        except (ValueError, OSError):
            pass
    logged = storage_state_path().is_file()
    return {
        "logged_in": logged,
        "state": "ready" if logged else "logged_out",
        "message": "已登录混元3D" if logged else "尚未微信登录",
        "session_dir": str(session_dir()),
    }


def is_logged_in() -> bool:
    path = storage_state_path()
    if not path.is_file() or path.stat().st_size < 8:
        return False
    try:
        with path.open("rb") as fp:
            if fp.read(1) != b"{":
                return False
        json.loads(path.read_text("utf-8"))
        return True
    except (OSError, ValueError, UnicodeDecodeError):
        return False


def logout() -> dict:
    for path in (storage_state_path(), status_file_path()):
        if path.is_file():
            path.unlink()
    return _write_status(logged_in=False, state="logged_out", message="已退出混元3D 登录")


def _click_first(page, labels, exact=False) -> bool:
    for label in labels:
        candidates = [
            page.get_by_role("button", name=label),
            page.get_by_role("link", name=label),
            page.get_by_text(label, exact=exact),
        ]
        for locator in candidates:
            try:
                if locator.count() and locator.first.is_visible(timeout=800):
                    locator.first.click(timeout=3000)
                    page.wait_for_timeout(600)
                    return True
            except Exception:
                continue
    return False


def _looks_logged_in(page) -> bool:
    try:
        if page.get_by_text("登录", exact=True).count():
            login = page.get_by_text("登录", exact=True).first
            if login.is_visible(timeout=300):
                return False
    except Exception:
        pass
    cookies = page.context.cookies()
    interesting = [
        cookie
        for cookie in cookies
        if any(
            token in cookie["name"].lower()
            for token in ("token", "openid", "uin", "skey", "hy3d", "session", "userid")
        )
    ]
    try:
        storage_keys = page.evaluate(
            "() => Object.keys({...localStorage, ...sessionStorage})"
        )
    except Exception:
        storage_keys = []
    storage_hit = any(
        any(token in key.lower() for token in ("token", "user", "auth", "openid", "uin"))
        for key in storage_keys
    )
    has_generate = False
    try:
        has_generate = bool(page.get_by_text("立即生成").count())
    except Exception:
        pass
    # 有生成入口，且 cookie/storage 任一命中，或页面上已无「登录」
    if has_generate and (interesting or storage_hit):
        return True
    if has_generate and not _page_shows_login_cta(page):
        return True
    return bool(interesting and storage_hit)


def _page_shows_login_cta(page) -> bool:
    try:
        if page.get_by_text("登录", exact=True).count():
            return page.get_by_text("登录", exact=True).first.is_visible(timeout=300)
    except Exception:
        pass
    return False


def _try_site_logout(page) -> None:
    """尽量在官网侧退出，方便额度满后换号扫码。"""
    for label in ("退出登录", "退出", "登出", "注销"):
        if _click_first(page, (label,), exact=False):
            page.wait_for_timeout(800)
            break
    # 常见头像/用户菜单再点一次退出
    try:
        avatar = page.locator(
            '[class*="avatar"], [class*="user-info"], [class*="UserInfo"]'
        ).first
        if avatar.count() and avatar.is_visible(timeout=500):
            avatar.click(timeout=2000)
            page.wait_for_timeout(400)
            _click_first(page, ("退出登录", "退出", "登出"), exact=False)
    except Exception:
        pass


def start_wechat_login_blocking(
    timeout_sec: int = LOGIN_TIMEOUT_SEC,
    *,
    update: Optional[UpdateFn] = None,
    reason: str = "auth",
    clear_existing: bool = False,
) -> dict:
    """打开有头浏览器供微信扫码，阻塞直到成功 / 超时 / 失败。

    reason:
      - auth: 未登录或凭证失效
      - quota: 当前账号额度满，需换号
    clear_existing: 先清本地会话；quota 场景会再尝试官网退出。
    """
    from playwright.sync_api import sync_playwright

    def _emit(event: str, **payload: Any) -> None:
        if update:
            update(event=event, **payload)

    if clear_existing or reason == "quota":
        logout()

    if reason == "quota":
        message = (
            "当前混元账号额度/次数已满，请在弹出的浏览器中退出旧号并用微信扫码登录其他账号"
        )
    else:
        message = "请在弹出的浏览器窗口中用微信扫码登录混元3D"

    _write_status(
        logged_in=False,
        state="waiting_scan",
        message=message,
        reason=reason,
    )
    _emit(
        "login_waiting",
        stage=message,
        progress=5,
        timeout_sec=timeout_sec,
        reason=reason,
    )
    # 同步打到 stderr，方便 Agent/终端注意弹窗
    print(f"[hunyuan] {message}（最长等待 {timeout_sec}s）", flush=True)

    try:
        with sync_playwright() as playwright:
            browser = playwright.chromium.launch(
                channel=_browser_channel(), headless=False
            )
            context = browser.new_context(
                locale="zh-CN", viewport={"width": 1280, "height": 900}
            )
            page = context.new_page()
            page.goto(HUNYUAN_URL, wait_until="domcontentloaded", timeout=60000)
            page.wait_for_timeout(2000)

            if reason == "quota":
                _try_site_logout(page)
                page.goto(HUNYUAN_URL, wait_until="domcontentloaded", timeout=60000)
                page.wait_for_timeout(1500)

            _click_first(page, ("登录", "立即登录"), exact=True)
            page.wait_for_timeout(800)
            _click_first(page, ("微信", "微信登录", "微信扫码"), exact=False)

            deadline = time.monotonic() + timeout_sec
            while time.monotonic() < deadline:
                if _looks_logged_in(page):
                    session_dir().mkdir(parents=True, exist_ok=True)
                    context.storage_state(path=str(storage_state_path()))
                    result = _write_status(
                        logged_in=True,
                        state="ready",
                        message="微信登录成功",
                        reason=reason,
                    )
                    browser.close()
                    _emit("login_ok", stage="微信登录成功", progress=15, **result)
                    print("[hunyuan] 微信登录成功，会话已保存", flush=True)
                    return result
                page.wait_for_timeout(1500)

            result = _write_status(
                logged_in=False,
                state="timeout",
                message="微信登录超时，请重试",
                reason=reason,
            )
            browser.close()
            _emit("login_failed", stage="微信登录超时", progress=0, **result)
            raise TimeoutError("微信登录超时，请重试")
    except TimeoutError:
        raise
    except Exception as exc:
        result = _write_status(
            logged_in=False, state="error", message=f"登录失败：{exc}", reason=reason
        )
        _emit("login_failed", stage=str(exc), progress=0, **result)
        raise RuntimeError(f"登录失败：{exc}") from exc


def ensure_login(
    *,
    interactive: bool = True,
    timeout_sec: int = LOGIN_TIMEOUT_SEC,
    update: Optional[UpdateFn] = None,
    force: bool = False,
    reason: str = "auth",
) -> dict:
    """若无会话则弹微信登录；interactive=False 且未登录则抛错。

    force=True 时忽略本地凭证，强制打开登录页（换号 / 重登）。
    """
    if is_logged_in() and not force:
        current = status()
        if update:
            update(
                event="login_cached",
                stage="已有混元登录凭证",
                progress=12,
                **{k: v for k, v in current.items() if k != "event"},
            )
        return current

    if not interactive:
        raise RuntimeError("尚未微信登录混元3D，请先执行 --login 或允许交互登录。")

    if update:
        update(
            event="login_required",
            stage=(
                "当前账号额度已满，需要换号登录"
                if reason == "quota"
                else "需要微信登录混元3D"
            ),
            progress=2,
            reason=reason,
        )
    return start_wechat_login_blocking(
        timeout_sec=timeout_sec,
        update=update,
        reason=reason,
        clear_existing=force or reason == "quota",
    )