import asyncio
import os
from playwright.async_api import async_playwright
from PIL import Image
from fpdf import FPDF

OUTPUT_DIR = "static/outputs"

async def capture_workflows(public_url: str, pdf_filename: str = "workflow_screens.pdf"):
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    pdf_path = os.path.join(OUTPUT_DIR, pdf_filename)

    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        page = await browser.new_page()
        print(f"Opening page: {public_url}")
        await page.goto(public_url, wait_until="load")

        # --- Utility: wait for stable layout ---
        async def wait_for_layout_stable():
            await page.evaluate("""
                (async () => {
                    let stableCount = 0;
                    let lastHeight = document.body.scrollHeight;
                    await new Promise((resolve) => {
                        const check = setInterval(() => {
                            const current = document.body.scrollHeight;
                            if (Math.abs(current - lastHeight) < 1) {
                                stableCount++;
                                if (stableCount >= 5) {
                                    clearInterval(check);
                                    setTimeout(resolve, 300);
                                }
                            } else {
                                stableCount = 0;
                                lastHeight = current;
                            }
                        }, 100);
                    });
                })();
            """)

        # --- Fix only sidebar; keep top bar dynamic per screen ---
        async def fix_layout_dynamically():
            try:
                await page.evaluate("""
                    const sidebar = document.querySelector('aside, .sidebar, nav');
                    const topbar = document.querySelector('.top-bar, header, .app-header, .navbar, .page-header');
                    // Reset old fixed styles
                    [sidebar, topbar].forEach(el => {
                        if (!el) return;
                        el.style.position = '';
                        el.style.top = '';
                        el.style.left = '';
                        el.style.width = '';
                        el.style.height = '';
                        el.style.zIndex = '';
                        el.style.overflow = '';
                        el.style.transition = '';
                    });
                    // Fix sidebar only (topbar remains dynamic)
                    if (sidebar) {
                        const rect = sidebar.getBoundingClientRect();
                        const sidebarWidth = rect.width || sidebar.offsetWidth || 220;
                        const topOffset = topbar ? topbar.offsetHeight || 60 : 60;
                        sidebar.style.position = 'fixed';
                        sidebar.style.top = topOffset + 'px';
                        sidebar.style.left = '0';
                        sidebar.style.height = `calc(100vh - ${topOffset}px)`;
                        sidebar.style.width = sidebarWidth + 'px';
                        sidebar.style.zIndex = '1000';
                        sidebar.style.overflow = 'auto';
                        sidebar.style.transition = 'none';
                        sidebar.dataset.locked = 'true';
                    }
                    const main = document.querySelector('main, .main-content, .app-body, .content');
                    if (main && sidebar) {
                        main.style.marginLeft = sidebar.offsetWidth + 'px';
                        main.style.marginTop = topbar ? (topbar.offsetHeight || 60) + 'px' : '0';
                        main.style.position = 'relative';
                    }
                    window.scrollTo(0, 0);
                """)
            except Exception as e:
                print(f"[WARN] Layout fix failed: {e}")

        # --- Wait until sidebar exists ---
        await page.wait_for_selector(
            "aside, .sidebar, nav, .side-panel, .left-menu, .menu",
            timeout=20000
        )
        await wait_for_layout_stable()
        await fix_layout_dynamically()
        await asyncio.sleep(1)

        # --- JS Logic: detect workflows and tabs 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){
              const href=m.getAttribute('href');
              if(href && href.startsWith('#')) id=href.substring(1);
            }
            const s=screens.find(x=>x.id===id);
            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;
          window.__getSubScreens = function(screen){
            const tabs=[...screen.querySelectorAll('.tab, .nav-link, [role="tab"], [data-tab], .sub-tab, .tab-item')];
            const list=[];
            for(const t of tabs){
              const sub=t.textContent?.trim();
              if(sub && !list.includes(sub)) list.push(sub);
            }
            return list;
          };
          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.__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'});
            window.__currentIndex++;
            const subs = window.__getSubScreens(screen);
            return {screenName:wfName, subScreens:subs};
          };
          window.__clickSubScreen = async function(name){
            const tabs=[...document.querySelectorAll('.tab, .nav-link, [role="tab"], [data-tab], .sub-tab, .tab-item')];
            const t=tabs.find(x=>x.textContent.trim()===name);
            if(t){t.click(); return true;}
            return false;
          };
        })();
        """
        await page.evaluate(js_logic)
        await asyncio.sleep(1)

        screenshots = []
        index = 0

        # --- Capture all screens ---
        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", [])
            screenshot_path = os.path.join(OUTPUT_DIR, f"{screen_name}.png")
            print(f"📸 Capturing main screen: {screen_name}")

            await wait_for_layout_stable()
            await fix_layout_dynamically()
            await asyncio.sleep(1.0)
            await page.screenshot(path=screenshot_path, full_page=True)
            screenshots.append(screenshot_path)

            # --- Capture sub-tabs ---
            first_active_skipped = False
            for sub in sub_screens:
                is_active = await page.evaluate(
                    """(subText) => {
                        const tabs=[...document.querySelectorAll('.tab, .nav-link, [role="tab"], [data-tab], .sub-tab, .tab-item')];
                        const t=tabs.find(x=>x.textContent.trim()===subText);
                        if(!t) return false;
                        const cls=t.getAttribute('class')||'';
                        return cls.includes('active');
                    }""",
                    sub
                )
                if is_active and not first_active_skipped:
                    first_active_skipped = True
                    continue

                print(f"   ↳ Capturing sub-screen: {sub}")
                await page.evaluate(f"window.__clickSubScreen('{sub}')")
                await wait_for_layout_stable()
                await fix_layout_dynamically()
                await asyncio.sleep(1.0)

                sub_name_clean = sub.replace(" ", "_").lower()
                sub_path = os.path.join(OUTPUT_DIR, f"{screen_name}_{sub_name_clean}.png")
                await page.screenshot(path=sub_path, full_page=True)
                screenshots.append(sub_path)

            index += 1

        await browser.close()

        # --- Generate PDF ---
        if not screenshots:
            raise RuntimeError("No screenshots captured — check if .screen elements exist!")

        print(f"🧾 Combining {len(screenshots)} screenshots into PDF: {pdf_path}")

        pdf = FPDF()
        for img_path in screenshots:
            image = Image.open(img_path)
            w, h = image.size
            pdf_w, pdf_h = 210, 297
            aspect = h / w
            pdf.add_page()
            pdf.image(img_path, 0, 0, pdf_w, pdf_w * aspect)
        pdf.output(pdf_path, "F")

        print(f"✅ PDF generated successfully: {pdf_path}")
        return pdf_path


generate_ui_report = capture_workflows