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

CLI 友好：同步阻塞扫码登录（非后台线程）。
"""
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 = 180

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:
    return storage_state_path().is_file()


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
    return bool((interesting or storage_hit) and has_generate)


def start_wechat_login_blocking(
    timeout_sec: int = LOGIN_TIMEOUT_SEC,
    *,
    update: Optional[UpdateFn] = None,
) -> dict:
    """打开有头浏览器供微信扫码，阻塞直到成功 / 超时 / 失败。"""
    from playwright.sync_api import sync_playwright

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

    _write_status(
        logged_in=False,
        state="waiting_scan",
        message="请在弹出的浏览器窗口中用微信扫码登录",
    )
    _emit(
        "login_waiting",
        stage="请在弹出的浏览器窗口中用微信扫码登录",
        progress=5,
        timeout_sec=timeout_sec,
    )

    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(2500)
            _click_first(page, ("登录", "立即登录"), exact=True)
            _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="微信登录成功"
                    )
                    browser.close()
                    _emit("login_ok", stage="微信登录成功", progress=15, **result)
                    return result
                page.wait_for_timeout(1500)

            result = _write_status(
                logged_in=False, state="timeout", message="微信登录超时，请重试"
            )
            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}"
        )
        _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,
) -> dict:
    """若无会话则弹微信登录；interactive=False 且未登录则抛错。"""
    if is_logged_in():
        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="需要微信登录混元3D",
            progress=2,
        )
    return start_wechat_login_blocking(timeout_sec=timeout_sec, update=update)