Update playwright_model.py

playwright_model.py

@@ -13,16 +13,17 @@ async def capture_workflows(public_url: str, pdf_filename: str = "workflow_scree
     async with async_playwright() as p:
         browser = await p.chromium.launch(headless=True)
         page = await browser.new_page()
-        print(f"Opening page: {public_url}")
+        print(f"🌐 Opening page: {public_url}")
         await page.goto(public_url, wait_until="load")
 
-        # ✅ Inject workflow traversal logic (copied from your frontend)
+        # ✅ Inject enhanced logic to capture main and sub screens dynamically
         js_logic = """
         (function(){
           const menus=[...document.querySelectorAll('.menu-item')];
           const screens=[...document.querySelectorAll('.screen')];
           window.__ordered=[];
           const seen=new Set();
+
           for(const m of menus){
             let id=m.dataset.screen||m.dataset.target;
             if(!id){
@@ -33,65 +34,103 @@ async def capture_workflows(public_url: str, pdf_filename: str = "workflow_scree
             if(s && !seen.has(s)){seen.add(s);window.__ordered.push({menu:m,screen:s});}
           }
           for(const s of screens) if(!seen.has(s)) window.__ordered.push({menu:null,screen:s});
+
           window.__visitedWorkflows=[];
           window.__currentIndex=0;
           window.__done=false;
+
+          async function getSubScreens(screen){
+            const tabs=[...screen.querySelectorAll('.tab, .nav-link, [role="tab"], [data-tab]')];
+            const results=[];
+            for(const t of tabs){
+              const subName=t.textContent?.trim().replace(/\\s+/g,'-').toLowerCase();
+              if(subName) results.push({tab:t, name:subName});
+            }
+            return results;
+          }
+
           window.__captureNext=async function(){
             if(window.__done) return false;
             if(window.__currentIndex>=window.__ordered.length){window.__done=true;return false;}
             const pair=window.__ordered[window.__currentIndex];
             const {menu,screen}=pair;
             if(!screen){window.__done=true;return false;}
+
             const wfName=screen.id || screen.getAttribute('data-name') || ('screen_'+window.__currentIndex);
-            if(window.__visitedWorkflows.includes(wfName)){window.__done=true;return false;}
+            if(window.__visitedWorkflows.includes(wfName)){
+              window.__currentIndex++;
+              return window.__captureNext();
+            }
+
             window.__visitedWorkflows.push(wfName);
             document.querySelectorAll('.screen').forEach(s=>s.classList.remove('active'));
             document.querySelectorAll('.menu-item').forEach(m=>m.classList.remove('active'));
             screen.classList.add('active');
             if(menu) menu.classList.add('active');
             screen.scrollIntoView({behavior:'smooth',block:'center'});
+
+            const mainTitle = screen.querySelector('.page-title, .header-title, h2')?.textContent?.trim() || wfName;
+            document.title = mainTitle;
+
+            const subs = await getSubScreens(screen);
             window.__currentIndex++;
-            return {screenName:wfName};
+            return {screenName: wfName, subScreens: subs.map(x=>x.name)};
+          };
+
+          window.__clickSubScreen = async function(mainName, subName){
+            const screen = document.getElementById(mainName);
+            if(!screen) return false;
+            const tab=[...screen.querySelectorAll('.tab, .nav-link, [role="tab"], [data-tab]')]
+              .find(t=>t.textContent.trim().toLowerCase().includes(subName));
+            if(tab){
+              tab.click();
+              await new Promise(r=>setTimeout(r,700));
+              const title = screen.querySelector('.page-title, .header-title, h2')?.textContent?.trim() || mainName;
+              document.title = title + ' - ' + subName;
+              return true;
+            }
+            return false;
           };
         })();
         """
+
         await page.evaluate(js_logic)
         await asyncio.sleep(1)
 
         screenshots = []
         index = 0
 
-        # ✅ Loop through each workflow
+        # ✅ Loop through main workflows and their sub-screens dynamically
         while True:
             result = await page.evaluate("window.__captureNext()")
             if not result:
                 break
 
             screen_name = result.get("screenName", f"screen_{index}")
+            sub_screens = result.get("subScreens", [])
+
+            # Capture main screen
             screenshot_path = os.path.join(OUTPUT_DIR, f"{screen_name}.png")
-            print(f"📸 Capturing {screen_name} -> {screenshot_path}")
-
-            # ✅ Dynamically update visible and tab title before screenshot
-            await page.evaluate(
-                """(name) => {
-                    // Update visible title text
-                    const titleEl = document.querySelector('.page-title, .header-title, main h2');
-                    if (titleEl) titleEl.textContent = name;
-
-                    // Update browser tab title
-                    document.title = name;
-                }""",
-                screen_name
-            )
-
-            await asyncio.sleep(1.2)  # wait for animations/UI updates
+            print(f"📸 Capturing main screen: {screen_name}")
+
+            await asyncio.sleep(1.2)
             await page.screenshot(path=screenshot_path, full_page=True)
             screenshots.append(screenshot_path)
+
+            # Capture sub-screens/tabs
+            for sub in sub_screens:
+                print(f"  ↳ Capturing sub-screen: {sub}")
+                await page.evaluate(f"window.__clickSubScreen('{screen_name}','{sub}')")
+                await asyncio.sleep(1.0)
+                sub_path = os.path.join(OUTPUT_DIR, f"{screen_name}_{sub}.png")
+                await page.screenshot(path=sub_path, full_page=True)
+                screenshots.append(sub_path)
+
             index += 1
 
         await browser.close()
 
-        # ✅ Combine all screenshots into a single PDF
+        # ✅ Combine all screenshots into PDF
         if not screenshots:
             raise RuntimeError("No screenshots captured — check if .screen elements exist!")