#!/usr/bin/env python3
"""PS pension module UI audit — https://retire.krisforbes.ca in EN + FR at 320/360/390 + 1280.
Drives localStorage state (PS answers), captures screenshots + geometry at the PS quiz steps
and the results card. Checks:
- no horizontal document overflow (scrollWidth <= innerWidth + 1)
- #step-body is the scroller; #btn-next + #panel bottom cluster fully in viewport
- #ps-card table fits the viewport; sources links present
Writes audit/ps-pension/geo.json + *.png, prints PS-AUDIT-OK/FAIL.
"""
import json, os, sys
from playwright.sync_api import sync_playwright

OUT = 'audit/ps-pension'
os.makedirs(OUT, exist_ok=True)
CHROME = '/usr/bin/google-chrome'
BASE = 'https://retire.krisforbes.ca/'
LS = 'retire-planner-v1'

ANS = {
    "province": "ON", "age": 60, "retireAge": 60, "household": "single",
    "partner": {"age": 60, "income": 0, "rrsp": 0, "tfsa": 0, "nonreg": 0, "monthlyContribution": 0},
    "income": 0, "rrsp": 100000, "tfsa": 30000, "nonreg": 20000,
    "monthlyContribution": 0, "employerMatchMonthly": 0,
    "spendingAnnual": 50000, "cppStartAge": 65, "oasStartAge": 65,
    "pensionAnnual": 0, "rentalAnnual": 0, "otherAnnual": 0,
    "psPension": True, "psJoinYear": 2000, "psServiceYears": 30, "psAvgSalary": 100000,
    "psOccupation": "it", "psItLevel": "IT-02", "psItStep": 1,
    "realReturnPre": 3.5, "realReturnPost": 2.5, "inflation": 2.0, "lifeExp": 95,
    "interests": ["dock", "travel"],
}
STEPS = {  # visible step index (single household, module ON): 20 steps
    'ps-toggle': 11, 'ps-join': 12, 'ps-service': 13, 'ps-salary': 14,
    'ps-occ': 15, 'ps-it': 16, 'extra-note': 17,
}
DEVICES = [
    ('iphone14', 390, 844, 3.0),
    ('pixel',    360, 800, 2.75),
    ('se',       320, 700, 2.0),
    ('desktop', 1280, 900, 1.0),
]

GEO_JS = """() => {
  const de = document.documentElement;
  const out = {
    docOverflowX: de.scrollWidth - window.innerWidth,
    stepBodyScroll: (() => { const s = document.getElementById('step-body'); return s ? (s.scrollHeight > s.clientHeight) : null; })(),
    btnNextInView: (() => { const b = document.getElementById('btn-next'); if (!b) return null; const r = b.getBoundingClientRect(); return r.bottom <= window.innerHeight + 1 && r.top >= -1; })(),
  };
  const card = document.getElementById('ps-card');
  if (card && !card.hidden) {
    const r = card.getBoundingClientRect();
    const t = card.querySelector('table');
    out.psCard = { visible: getComputedStyle(card).display !== 'none', width: Math.round(r.width),
                   fits: (t ? t.scrollWidth : 0) <= window.innerWidth + 1,
                   links: card.querySelectorAll('a').length };
  }
  return out;
}"""

def main():
    results, fails = [], 0
    with sync_playwright() as pw:
        browser = pw.chromium.launch(executable_path=CHROME)
        for lang in ('en', 'fr'):
            for name, w, h, dpr in DEVICES:
                ctx = browser.new_context(viewport={'width': w, 'height': h}, device_scale_factor=dpr,
                                          is_mobile=(w < 800), has_touch=(w < 800))
                page = ctx.new_page()
                page.goto(BASE, wait_until='domcontentloaded')
                page.evaluate("([ls, st]) => localStorage.setItem(ls, JSON.stringify(st))",
                              [LS, {"a": ANS, "step": 0, "screen": "quiz", "lang": lang}])
                for shot, idx in STEPS.items():
                    page.evaluate("([ls, st]) => localStorage.setItem(ls, JSON.stringify(st))",
                                  [LS, {"a": ANS, "step": idx, "screen": "quiz", "lang": lang}])
                    page.reload(wait_until='domcontentloaded')
                    page.wait_for_timeout(350)
                    geo = page.evaluate(GEO_JS)
                    png = f"{OUT}/{lang}-{name}-{shot}.png"
                    page.screenshot(path=png)
                    bad = [k for k in ('docOverflowX',) if geo[k] > 1]
                    if geo.get('btnNextInView') is False: bad.append('btnNext')
                    rec = {"lang": lang, "device": name, "shot": shot, "geo": geo, "ok": not bad, "bad": bad}
                    results.append(rec); fails += 1 if bad else 0
                    if not bad: print(f"  ok  {lang} {name:9s} {shot}")
                    else: print(f"  FAIL {lang} {name:9s} {shot} {bad} {geo}")
                # results screen with PS card
                page.evaluate("([ls, st]) => localStorage.setItem(ls, JSON.stringify(st))",
                              [LS, {"a": ANS, "step": 0, "screen": "results", "lang": lang}])
                page.reload(wait_until='domcontentloaded')
                page.wait_for_timeout(900)
                geo = page.evaluate(GEO_JS)
                page.evaluate("() => { const c = document.getElementById('ps-card'); if (c) c.scrollIntoView(); }")
                page.wait_for_timeout(200)
                page.screenshot(path=f"{OUT}/{lang}-{name}-results-pscard.png")
                bad = []
                if geo.get('docOverflowX', 0) > 1: bad.append('overflowX')
                pc = geo.get('psCard')
                if not pc or not pc.get('visible'): bad.append('psCardHidden')
                elif not pc.get('fits'): bad.append('psCardWide')
                elif pc.get('links', 0) < 5: bad.append('links<5')
                rec = {"lang": lang, "device": name, "shot": "results-pscard", "geo": geo, "ok": not bad, "bad": bad}
                results.append(rec); fails += 1 if bad else 0
                if not bad: print(f"  ok  {lang} {name:9s} results-pscard")
                else: print(f"  FAIL {lang} {name:9s} results-pscard {bad} {geo}")
                ctx.close()
        browser.close()
    with open(f'{OUT}/geo.json', 'w') as f:
        json.dump(results, f, indent=1)
    print(f"{len(results)} shots, {fails} failures")
    print("PS-AUDIT-OK" if fails == 0 else "PS-AUDIT-FAIL")
    sys.exit(1 if fails else 0)

if __name__ == '__main__':
    main()