"""Exercise the production experience and save desktop/mobile review images."""
import json
import os
import re
from pathlib import Path
from playwright.sync_api import sync_playwright, expect

OUT = Path('artifacts')
OUT.mkdir(exist_ok=True)
BASE = os.environ.get('PREVIEW_URL', 'http://127.0.0.1:7860')
with sync_playwright() as p:
    browser = p.chromium.launch(executable_path=os.environ.get('CHROMIUM_PATH'), args=['--no-sandbox', '--use-gl=angle', '--use-angle=swiftshader', '--enable-unsafe-swiftshader'])
    page = browser.new_page(viewport={'width': 1440, 'height': 1000}, device_scale_factor=1)
    errors = []
    page.on('pageerror', lambda e: errors.append(str(e)))
    page.goto(BASE, wait_until='networkidle')
    page.evaluate('document.fonts.ready')
    expect(page.locator('[data-render-state="ready"]')).to_have_count(2, timeout=30000)
    page.wait_for_timeout(1200)
    page.screenshot(path=str(OUT / 'desktop-hero.png'))
    for section in ['dual-learning', 'dataset', 'petase', 'evidence', 'playground']:
        page.locator('#'+section).evaluate("e=>window.scrollTo({top:e.offsetTop-90,behavior:'instant'})")
        page.wait_for_timeout(1200)
        page.wait_for_function("id=>getComputedStyle(document.querySelector('#'+id+' .section-heading')).opacity==='1'", arg=section)
        page.screenshot(path=str(OUT / (section+'.png')))
    page.screenshot(path=str(OUT / 'desktop-full.png'), full_page=True)
    assert page.evaluate('document.documentElement.scrollWidth <= innerWidth'), 'desktop overflow'
    dual = page.locator('#dual-learning')
    for name, mode in [('Text → Protein', 't2p'), ('Protein → Text', 'p2t'), ('Joint learning', 'joint')]:
        dual.get_by_role('button', name=name, exact=True).click()
        expect(dual.locator('.dual-diagram')).to_have_class('dual-diagram mode-'+mode)
    petase = page.locator('#petase')
    for i in [1, 0]:
        petase.locator('.case-tabs button').nth(i).click()
        expect(petase.locator('[data-render-state="ready"]')).to_have_count(1)
    petase.get_by_role('button', name='Color by confidence', exact=True).click()
    expect(petase.locator('.confidence-legend')).to_be_visible()
    petase.locator('.residue-list button').nth(42).click()
    expect(petase.locator('.residue-list button').nth(42)).to_have_attribute('aria-pressed', 'true')
    page.keyboard.press('ArrowRight')
    expect(petase.locator('.residue-list button').nth(43)).to_be_focused()
    for name in ['Surface', 'Sticks', 'Cartoon']:
        petase.get_by_role('button', name=name, exact=True).click()
        page.wait_for_timeout(1200)
    with page.expect_download() as download:
        petase.get_by_role('link', name='↓ PDB', exact=True).click()
    assert download.value.suggested_filename.endswith('.pdb')
    petase.get_by_role('button', name='Try this function in the lab').click()
    lab = page.locator('#playground')
    expect(lab.locator('#model-input')).to_have_value(re.compile('Belongs to the alpha/beta hydrolase superfamily.*'))
    expect(lab.get_by_text('Original LoRA generation prompt', exact=False)).to_be_visible()
    lab.get_by_role('button', name='Use input ↗', exact=True).click()
    expect(lab.locator('.generation-result')).to_have_count(0)
    expect(lab.locator('.result-playback')).to_have_count(0)
    lab.get_by_role('button', name='Protein → Text', exact=True).click()
    lab.get_by_role('button', name='Use input ↗', exact=True).click()
    expect(lab.locator('.generation-result')).to_have_count(0)
    expect(lab.locator('#model-input')).not_to_have_value('')
    evidence = page.locator('#evidence')
    assert evidence.locator('.metric-tabs button[aria-pressed="true"]').inner_text() == 'Sequence recovery'
    for name, n in [('Keyword recovery', '476'), ('ProTrek score', '500'), ('Inverse folding recovery', '500'), ('Inverse folding SW', '500')]:
        evidence.get_by_role('button', name=name, exact=True).click()
        expect(evidence.locator('.chart-selection')).to_contain_text('n = ' + n)
        assert evidence.locator('.bar-row').count() == 7
    assert evidence.get_by_role('button', name='PAE', exact=True).count() == 0
    evidence.get_by_role('button', name='Protein → Text', exact=True).click()
    evidence.get_by_role('button', name='ROUGE-L', exact=True).click()
    evidence.locator('#benchmark-set').select_option('normal')
    assert evidence.locator('.bar-row').count() > 1
    evidence.locator('.scaling-bars button').first.click()
    expect(evidence.locator('.scaling-bars button').first).to_have_attribute('aria-pressed', 'true')
    page.set_viewport_size({'width': 390, 'height': 844})
    page.goto(BASE, wait_until='networkidle')
    page.wait_for_timeout(1200)
    page.screenshot(path=str(OUT / 'mobile-hero.png'))
    assert page.evaluate('document.documentElement.scrollWidth <= innerWidth'), 'mobile overflow'
    for section in ['dual-learning', 'dataset', 'petase', 'evidence', 'playground']:
        page.locator('#'+section).evaluate("e=>window.scrollTo({top:e.offsetTop-90,behavior:'instant'})")
        page.wait_for_timeout(1200)
        page.screenshot(path=str(OUT / ('mobile-'+section+'.png')))
    page.get_by_role('button', name='Toggle navigation').click()
    expect(page.locator('nav')).to_be_visible()
    page.locator('nav a').first.click()
    expect(page.locator('nav')).not_to_be_visible()
    (OUT / 'browser-report.json').write_text(json.dumps({'page_errors': errors}, indent=2))
    assert not errors, errors
    browser.close()
    print('Desktop/mobile layouts and both 3D viewers passed. Screenshots saved in artifacts/.')