import asyncio
from playwright.async_api import async_playwright
import os

async def main():
    html_path = os.path.abspath("app/static/bim/index.html")
    uri = f"file:///{html_path.replace(chr(92), '/')}"
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        page = await browser.new_page()
        await page.set_viewport_size({"width": 1920, "height": 1080})
        await page.goto(uri, wait_until="networkidle")
        
        # Wait a bit for 3D to render
        await asyncio.sleep(2)
        
        artifact_dir = os.path.expanduser("~/.gemini/antigravity/brain/dd9b916f-153c-4fba-b1cd-9ab49b94ab2b")
        
        # Take screenshot at day 40 (Middle of Segment 1)
        await page.evaluate("window.currentDay = 40; window.updateBim4DState(40); window.renderSiteDiaryDock(40);")
        await asyncio.sleep(1)
        await page.screenshot(path=f"{artifact_dir}/rolling_day_40.png")
        
        # Take screenshot at day 120 (Middle of Segment 2)
        await page.evaluate("window.currentDay = 120; window.updateBim4DState(120); window.renderSiteDiaryDock(120);")
        await asyncio.sleep(1)
        await page.screenshot(path=f"{artifact_dir}/rolling_day_120.png")
        
        # Take screenshot at day 210 (Middle of Segment 3)
        await page.evaluate("window.currentDay = 210; window.updateBim4DState(210); window.renderSiteDiaryDock(210);")
        await asyncio.sleep(1)
        await page.screenshot(path=f"{artifact_dir}/rolling_day_210.png")
        
        await browser.close()

if __name__ == "__main__":
    asyncio.run(main())