scripts/diagnose_ui_flow.py

from __future__ import annotations

import json
import os
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from playwright.sync_api import Page, expect, sync_playwright

BASE_URL = os.environ.get("KINK_AUDIT_BASE_URL", "http://127.0.0.1:8011").rstrip("/")
ARTIFACT_DIR = Path(os.environ.get("KINK_AUDIT_ARTIFACT_DIR", "/tmp/kink_diagnostics"))
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
STARTER_TARGET = 12


@dataclass(frozen=True)
class Credentials:
    user_id: str
    private_token: str


def auth_headers(creds: Credentials) -> dict[str, str]:
    return {"x-private-token": creds.private_token}


def json_request(
    path: str,
    *,
    method: str = "GET",
    payload: dict | None = None,
    headers: dict[str, str] | None = None,
    timeout: int = 30,
) -> dict:
    request_headers = dict(headers or {})
    data = None
    if payload is not None:
        request_headers["Content-Type"] = "application/json"
        data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(f"{BASE_URL}{path}", data=data, headers=request_headers, method=method)
    with urllib.request.urlopen(request, timeout=timeout) as response:
        body = response.read().decode("utf-8")
        return json.loads(body) if body else {}


def create_user() -> Credentials:
    payload = json_request("/users", method="POST", payload={})
    return Credentials(user_id=payload["id"], private_token=payload["private_token"])


def save_play(creds: Credentials, kink_id: str, interest_state: str = "like") -> None:
    json_request(
        f"/users/{urllib.parse.quote(creds.user_id)}/plays",
        method="POST",
        payload={"kink_id": kink_id, "interest_state": interest_state, "directions": ["together"]},
        headers=auth_headers(creds),
    )


def seed_completed_user() -> Credentials:
    creds = create_user()
    recs = json_request(
        f"/users/{urllib.parse.quote(creds.user_id)}/recommendations?limit={STARTER_TARGET}",
        headers=auth_headers(creds),
        timeout=60,
    ).get("items", [])
    if len(recs) < STARTER_TARGET:
        raise AssertionError(f"expected at least {STARTER_TARGET} recommendations, got {len(recs)}")
    for item in recs[:STARTER_TARGET]:
        save_play(creds, item["kink"]["id"])
    return creds


def capture_dom_state(page: Page) -> dict:
    return page.evaluate(
        """() => ({
          hasRoot: !!document.querySelector('#root'),
          viewTestIds: [...document.querySelectorAll('[data-testid$="-view"], [data-testid="starter-gate"], [data-testid^="onboarding-"]')]
            .map((el) => el.getAttribute('data-testid')).slice(0, 20),
          bodyText: document.body.innerText.slice(0, 700),
          buttons: [...document.querySelectorAll('button')]
            .map((button) => (button.textContent || '').trim()).filter(Boolean).slice(0, 25),
          inputs: [...document.querySelectorAll('input')].map((input) => input.placeholder || input.type || '').slice(0, 12)
        })""",
    )


def install_auth(page: Page, creds: Credentials | None) -> None:
    payload = json.dumps({"userId": creds.user_id, "token": creds.private_token}) if creds else ""
    script = (
        f"localStorage.setItem('kink_auth', JSON.stringify({payload}));"
        if creds
        else "localStorage.removeItem('kink_auth');"
    )
    page.add_init_script(script)


def run_browser_scenario(
    name: str,
    action: Callable[[Page, dict], None],
    *,
    setup_auth: Credentials | None = None,
) -> dict:
    events: list[dict] = []
    page_errors: list[str] = []
    console_errors: list[str] = []
    started = time.perf_counter()

    def rec(evt: dict) -> None:
        item = dict(evt)
        item["t_ms"] = round((time.perf_counter() - started) * 1000, 1)
        events.append(item)

    with sync_playwright() as playwright:
        browser = playwright.chromium.launch(headless=True, args=["--no-sandbox", "--disable-http-cache"])
        context = browser.new_context(viewport={"width": 1440, "height": 1000})
        page = context.new_page()
        install_auth(page, setup_auth)
        page.on("pageerror", lambda exc: page_errors.append(str(exc)))
        page.on("console", lambda msg: console_errors.append(f"{msg.type}: {msg.text}") if msg.type == "error" else None)
        page.on("request", lambda req: rec({"kind": "request", "method": req.method, "url": req.url}))
        page.on("response", lambda res: rec({"kind": "response", "status": res.status, "url": res.url}))
        scenario = {
            "name": name,
            "ok": True,
            "checkpoints": [],
            "page_errors": page_errors,
            "console_errors": console_errors,
        }
        try:
            page.goto(BASE_URL, wait_until="domcontentloaded", timeout=120_000)
            scenario["checkpoints"].append({"label": "boot", "state": capture_dom_state(page)})
            action(page, scenario)
            scenario["checkpoints"].append({"label": "after_action", "state": capture_dom_state(page)})
        except Exception as exc:
            scenario["ok"] = False
            scenario["failure"] = str(exc)
        finally:
            screenshot = ARTIFACT_DIR / f"{name}.png"
            page.screenshot(path=str(screenshot), full_page=True)
            scenario["screenshot"] = str(screenshot)
            context.close()
            browser.close()

    app_events = [event for event in events if event.get("url", "").startswith(BASE_URL)]
    summary: dict[str, int] = {}
    for event in app_events:
        path = event["url"].split(BASE_URL, 1)[-1].split("?", 1)[0] or "/"
        key = f"{event['kind']} {event.get('method', event.get('status', ''))} {path}"
        summary[key] = summary.get(key, 0) + 1
    scenario["request_summary"] = dict(sorted(summary.items()))
    scenario["event_counts"] = {"total": len(events), "app_only": len(app_events)}
    return scenario


def scenario_boot_unauth(page: Page, _scenario: dict) -> None:
    expect(page.get_by_test_id("onboarding-welcome")).to_be_visible(timeout=30_000)
    expect(page.get_by_test_id("onboarding-create-profile")).to_be_visible(timeout=30_000)


def scenario_boot_stored_auth(page: Page, _scenario: dict) -> None:
    expect(page.get_by_test_id("discover-view")).to_be_visible(timeout=30_000)
    expect(page.get_by_test_id("desktop-nav")).to_be_visible(timeout=30_000)


def scenario_login_existing(page: Page, _scenario: dict) -> None:
    creds = create_user()
    page.get_by_test_id("onboarding-have-profile").click()
    expect(page.get_by_test_id("onboarding-login")).to_be_visible(timeout=15_000)
    page.get_by_test_id("onboarding-login-user-id").fill(creds.user_id)
    page.get_by_test_id("onboarding-login-private-token").fill(creds.private_token)
    page.get_by_test_id("onboarding-open-profile").click()
    expect(page.get_by_test_id("starter-gate")).to_be_visible(timeout=60_000)


def scenario_starter_gate(page: Page, _scenario: dict) -> None:
    expect(page.get_by_test_id("starter-gate")).to_be_visible(timeout=60_000)
    expect(page.get_by_test_id("starter-discover-view")).to_be_visible(timeout=60_000)
    expect(page.get_by_test_id("discover-action-bar")).to_be_visible(timeout=60_000)


def main() -> int:
    completed = seed_completed_user()
    starter = create_user()
    scenarios = [
        run_browser_scenario("boot_unauth", scenario_boot_unauth),
        run_browser_scenario("boot_stored_auth", scenario_boot_stored_auth, setup_auth=completed),
        run_browser_scenario("login_existing", scenario_login_existing),
        run_browser_scenario("starter_gate", scenario_starter_gate, setup_auth=starter),
    ]
    report = {"generated_at": time.time(), "scenarios": scenarios}
    out_path = ARTIFACT_DIR / "diagnose_ui_flow.json"
    out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))
    return 0 if all(scenario["ok"] for scenario in scenarios) else 1


if __name__ == "__main__":
    raise SystemExit(main())