import os
from playwright.sync_api import sync_playwright
from PIL import Image
from fpdf import FPDF

OUTPUT_DIR = "static/outputs"

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)

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

        def wait_for_layout_stable():
            page.evaluate("""
                (() => {
                    let stableCount = 0;
                    let lastHeight = document.body.scrollHeight;
                    return 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);
                    });
                })();
            """)

        def fix_layout_dynamically():
            try:
                page.evaluate("""
                    (() => {
                        console.log('🔧 Fixing layout...');
                        
                        // Reset all layout elements
                        ['.top-bar, header, .app-header', 
                         'aside, .sidebar, nav', 
                         'main, .main-content, .content, .page-content'].forEach(selector => {
                            document.querySelectorAll(selector).forEach(el => {
                                if (el) {
                                    ['position','top','left','width','height','zIndex','overflow',
                                     'transition','margin','padding','boxSizing'].forEach(prop => {
                                        el.style[prop] = '';
                                    });
                                }
                            });
                        });
                        
                        // Get accurate element dimensions
                        const getDims = (el) => {
                            if (!el || !el.isConnected) return {width: 240, height: 60};
                            const rect = el.getBoundingClientRect();
                            const style = window.getComputedStyle(el);
                            return {
                                width: rect.width + (parseFloat(style.marginLeft) || 0) + (parseFloat(style.marginRight) || 0),
                                height: rect.height + (parseFloat(style.marginTop) || 0) + (parseFloat(style.marginBottom) || 0)
                            };
                        };
                        
                        // Fix topbar
                        const topbar = document.querySelector('.top-bar, header, .app-header, .navbar');
                        let topOffset = 60;
                        
                        if (topbar) {
                            const dims = getDims(topbar);
                            topOffset = Math.round(dims.height);
                            
                            Object.assign(topbar.style, {
                                position: 'fixed',
                                top: '0',
                                left: '0',
                                width: '100vw',
                                zIndex: '10000',
                                transition: 'none',
                                boxSizing: 'border-box'
                            });
                            
                            document.body.style.marginTop = topOffset + 'px';
                            console.log('✅ Topbar height:', topOffset, 'px');
                        }
                        
                        // Fix sidebar
                        const sidebar = document.querySelector('aside, .sidebar, nav');
                        let sidebarWidth = 240;
                        
                        if (sidebar) {
                            const dims = getDims(sidebar);
                            sidebarWidth = Math.round(dims.width);
                            
                            Object.assign(sidebar.style, {
                                position: 'fixed',
                                top: topOffset + 'px',
                                left: '0',
                                height: `calc(100vh - ${topOffset}px)`,
                                width: sidebarWidth + 'px',
                                zIndex: '9999',
                                overflow: 'auto',
                                transition: 'none',
                                boxSizing: 'border-box'
                            });
                            
                            console.log('✅ Sidebar width:', sidebarWidth, 'px');
                        }
                        
                        // Fix main content using padding (more reliable than margin)
                        const content = document.querySelector('main, .main-content, .content, .page-content');
                        if (content) {
                            content.style.paddingLeft = (sidebarWidth + 20) + 'px'; // +20px for breathing room
                            content.style.position = 'relative';
                            content.style.boxSizing = 'border-box';
                            console.log('✅ Content padding-left:', (sidebarWidth + 20), 'px');
                        }
                        
                        // Hide any duplicate fixed elements that might overlap
                        document.querySelectorAll('.title-bar, .page-title, [class*="header"]').forEach(el => {
                            if (el && el !== topbar && window.getComputedStyle(el).position === 'fixed') {
                                el.style.display = 'none';
                            }
                        });
                        
                        window.scrollTo(0, 0);
                        console.log('✅ Layout fix completed');
                    })();
                """)
            except Exception as e:
                print(f"[WARN] Layout fix failed: {e}")

        # Wait for sidebar and stabilize
        page.wait_for_selector("aside, .sidebar, nav", timeout=5000)
        wait_for_layout_stable()
        fix_layout_dynamically()
        page.wait_for_timeout(2000)  # Extra time for CSS transitions to complete

        # [Rest of the code remains the same as previous version]
        # ... (JS logic for capturing workflows, screenshot loop, PDF generation)

        # JS logic for capturing workflows
        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=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 = 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;
          };
        })();
        """
        page.evaluate(js_logic)
        page.wait_for_timeout(1000)

        screenshots = []
        index = 0

        # Capture all screens
        while True:
            result = 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}")

            wait_for_layout_stable()
            fix_layout_dynamically()
            page.wait_for_timeout(1000)
            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 = 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}")
                page.evaluate(f"window.__clickSubScreen('{sub}')")
                wait_for_layout_stable()
                fix_layout_dynamically()
                page.wait_for_timeout(1000)

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

            index += 1

        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