#!/usr/bin/env python3
"""
kimi_login.py — 在 246 headless Chrome 上登录 Kimi
====================================================
用法（需要 X11 forwarding 或手动操作）:
  ssh -X hmo@192.168.1.246
  python3 kimi_login.py

或者无头模式下手动注入 cookies:
  python3 kimi_login.py --inject-cookies '{"name":"...","value":"...","domain":".kimi.com"}'

登录成功后，cookies 自动保存在 /home/hmo/chrome-kimi-profile/ 下，
后续 chrome-kimi-cdp.service 重启后仍然保持登录态。
"""
import json
import sys
import time

from playwright.sync_api import sync_playwright

CDP_URL = "http://127.0.0.1:9222"
KIMI_URL = "https://www.kimi.com/"


def check_logged_in(page):
    """Check if Kimi is logged in by looking for user avatar or quota page access."""
    try:
        page.goto("https://www.kimi.com/membership/subscription?tab=quota",
                  wait_until="networkidle", timeout=15000)
        text = page.inner_text("body") or ""
        if "login" in page.url or "passport" in page.url:
            return False
        # If we can see usage data, we're logged in
        if "%" in text:
            return True
        return False
    except Exception:
        return False


def main():
    print("[kimi-login] Connecting to CDP...")

    with sync_playwright() as p:
        try:
            browser = p.chromium.connect_over_cdp(CDP_URL)
        except Exception as e:
            print(f"[kimi-login] ERROR: Cannot connect to CDP: {e}")
            print("[kimi-login] Is chrome-kimi-cdp.service running?")
            sys.exit(1)

        context = browser.contexts[0] if browser.contexts else browser.new_context()
        page = context.new_page()

        # Check if already logged in
        print("[kimi-login] Checking current login status...")
        if check_logged_in(page):
            print("[kimi-login] ✅ Already logged in!")
            print("[kimi-login] Cookies are saved in /home/hmo/chrome-kimi-profile/")
            return

        print("[kimi-login] Not logged in. Opening Kimi login page...")
        print("[kimi-login] " + "=" * 50)
        print("[kimi-login] If running with X11 forwarding (ssh -X), a browser window should appear.")
        print("[kimi-login] Please login manually in the browser window.")
        print("[kimi-login] Waiting up to 120 seconds for login...")
        print("[kimi-login] " + "=" * 50)

        page.goto(KIMI_URL, wait_until="networkidle", timeout=30000)

        # Wait for user to login (poll every 3 seconds, up to 120s)
        for i in range(40):
            time.sleep(3)
            current_url = page.url
            print(f"[kimi-login] Checking... ({(i + 1) * 3}s) URL: {current_url[:80]}")

            if "login" not in current_url and "passport" not in current_url:
                # Might be logged in now
                if check_logged_in(page):
                    print("[kimi-login] ✅ Login successful!")
                    print("[kimi-login] Cookies saved to /home/hmo/chrome-kimi-profile/")
                    return

        print("[kimi-login] ❌ Login timeout (120s). Please try again.")
        print("[kimi-login] Make sure X11 forwarding is enabled: ssh -X hmo@192.168.1.246")


if __name__ == "__main__":
    main()
