#!/usr/bin/env python3
"""
REAL Gemini Deep Think Workflow - v2
Uses Windows Chrome executable with real user profile for authenticated access
"""

import os
import sys
import time
import json
from datetime import datetime
from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeoutError

SCREENSHOTS_DIR = "/mnt/e/genesis-system/deep_think_results/screenshots"
# Windows Chrome paths accessible from WSL
CHROME_EXE = "/mnt/c/Program Files/Google/Chrome/Application/chrome.exe"
CHROME_USER_DATA = "/mnt/c/Users/P3/AppData/Local/Google/Chrome/User Data"

os.makedirs(SCREENSHOTS_DIR, exist_ok=True)

step_log = []

def log(msg, level="INFO"):
    ts = datetime.now().strftime("%H:%M:%S")
    line = f"[{ts}] [{level}] {msg}"
    print(line)
    step_log.append(line)

def screenshot(page, name):
    path = f"{SCREENSHOTS_DIR}/{name}.png"
    try:
        page.screenshot(path=path, full_page=False)
        log(f"Screenshot saved: {path}")
    except Exception as e:
        log(f"Screenshot failed for {name}: {e}", "WARN")
    return path

def run():
    log("=" * 60)
    log("GEMINI DEEP THINK BROWSER AUTOMATION v2")
    log(f"Started: {datetime.now().isoformat()}")
    log("=" * 60)

    workflow_steps = []

    with sync_playwright() as p:
        log("Launching Chrome with Windows user profile...")
        log(f"Chrome exe: {CHROME_EXE}")
        log(f"User data: {CHROME_USER_DATA}")

        browser = None
        page = None

        try:
            # Use persistent context with Windows Chrome executable
            browser = p.chromium.launch_persistent_context(
                user_data_dir=CHROME_USER_DATA,
                executable_path=CHROME_EXE,
                headless=False,
                slow_mo=300,
                args=[
                    "--profile-directory=Default",
                    "--no-first-run",
                    "--no-default-browser-check",
                    "--disable-features=ChromeWhatsNewUI",
                ],
                viewport={"width": 1280, "height": 900},
            )
            log("Chrome launched with persistent context (real profile)")
            workflow_steps.append({
                "step": 1,
                "action": "launch_browser",
                "detail": "chromium.launch_persistent_context with Windows Chrome + Default profile",
                "code": f'browser = p.chromium.launch_persistent_context(user_data_dir="{CHROME_USER_DATA}", executable_path="{CHROME_EXE}", headless=False)'
            })

            # Get or create page
            if browser.pages:
                page = browser.pages[0]
                log(f"Using existing page. Current URL: {page.url}")
            else:
                page = browser.new_page()
                log("Created new page")

        except Exception as e:
            log(f"Chrome launch failed: {e}", "ERROR")
            log("Falling back to Chromium without profile...")
            if browser:
                try:
                    browser.close()
                except:
                    pass
            browser_obj = p.chromium.launch(headless=False, slow_mo=300)
            browser = browser_obj
            page = browser.new_page()
            workflow_steps.append({
                "step": 1,
                "action": "launch_browser_fallback",
                "detail": "Chromium headless=False (no profile - will need login)",
            })

        try:
            # STEP 1: Navigate to Gemini
            log("\nSTEP 1: Navigate to gemini.google.com")
            page.goto("https://gemini.google.com", timeout=30000, wait_until="domcontentloaded")
            time.sleep(5)  # Wait for JS to load
            screenshot(page, "01_initial_load")

            current_url = page.url
            page_title = page.title()
            log(f"URL: {current_url}")
            log(f"Title: {page_title}")

            workflow_steps.append({
                "step": 2,
                "action": "navigate",
                "url": "https://gemini.google.com",
                "code": 'page.goto("https://gemini.google.com", timeout=30000, wait_until="domcontentloaded")',
                "result": f"URL={current_url}, Title={page_title}"
            })

            # STEP 2: Check login state
            log("\nSTEP 2: Checking login state...")

            # Wait a bit more for full page load
            time.sleep(3)
            screenshot(page, "02_after_wait")

            # Check for sign in vs signed in state
            page_text = page.evaluate("() => document.body.innerText")
            is_logged_in = "Sign in" not in page_text[:500]

            log(f"Login detection - 'Sign in' in page: {'YES - NOT LOGGED IN' if 'Sign in' in page_text[:500] else 'NO - LIKELY LOGGED IN'}")

            # More specific login check
            signin_btn = page.locator('button:has-text("Sign in"), a:has-text("Sign in")').count()
            profile_area = page.locator('[aria-label*="Google Account"], [aria-label*="profile"]').count()

            log(f"Sign-in button count: {signin_btn}")
            log(f"Profile elements: {profile_area}")

            if signin_btn > 0 and profile_area == 0:
                log("NOT LOGGED IN - Chrome profile cookies not transferred to Playwright", "WARN")
                log("Need to handle login or use cookie injection", "WARN")

                # Try to click sign in and see what happens
                screenshot(page, "02b_not_logged_in")

                # Get current page structure for documentation
                log("Documenting page structure for workflow reference...")
            else:
                log("LOGGED IN - Session active!")
                screenshot(page, "02b_logged_in")

            # STEP 3: Map full UI structure regardless of login state
            log("\nSTEP 3: Mapping Gemini UI structure...")

            # Get all interactive elements
            ui_map = page.evaluate("""
                () => {
                    const els = document.querySelectorAll(
                        'button, [role="button"], a[href], [role="switch"], [role="checkbox"], ' +
                        'select, [role="combobox"], [aria-haspopup], textarea, [contenteditable="true"]'
                    );

                    return Array.from(els).map(el => {
                        const rect = el.getBoundingClientRect();
                        return {
                            tag: el.tagName,
                            text: el.textContent?.trim()?.substring(0, 80),
                            aria: el.getAttribute('aria-label'),
                            role: el.getAttribute('role'),
                            class: el.className?.toString()?.substring(0, 80),
                            id: el.id || null,
                            'data-test-id': el.getAttribute('data-test-id'),
                            visible: rect.width > 0 && rect.height > 0,
                            position: {
                                x: Math.round(rect.x),
                                y: Math.round(rect.y),
                                w: Math.round(rect.width),
                                h: Math.round(rect.height)
                            }
                        };
                    }).filter(el => el.visible);
                }
            """)

            log(f"Total visible interactive elements: {len(ui_map)}")

            # Save UI map for analysis
            ui_map_path = "/mnt/e/genesis-system/deep_think_results/ui_map.json"
            with open(ui_map_path, 'w') as f:
                json.dump(ui_map, f, indent=2)
            log(f"UI map saved to: {ui_map_path}")

            # Find model selector
            model_elements = [el for el in ui_map if any(
                x in (el.get('text') or '').lower() + (el.get('aria') or '').lower()
                for x in ['model', '2.5', '2.0', '1.5', 'pro', 'flash', 'ultra', 'gemini', 'fast']
            )]

            log(f"\nModel-related elements ({len(model_elements)}):")
            for el in model_elements:
                log(f"  [{el['tag']}] text='{el.get('text', '')}' aria='{el.get('aria', '')}' pos={el['position']}")

            # Find Deep Think / thinking elements
            think_elements = [el for el in ui_map if any(
                x in (el.get('text') or '').lower() + (el.get('aria') or '').lower()
                for x in ['deep think', 'think', 'reason', 'extended', 'advanced']
            )]

            log(f"\nThinking-related elements ({len(think_elements)}):")
            for el in think_elements:
                log(f"  [{el['tag']}] text='{el.get('text', '')}' aria='{el.get('aria', '')}' pos={el['position']}")

            # STEP 4: Try to interact with model picker
            log("\nSTEP 4: Interacting with model/mode selector...")

            # The "Fast" button in the output above is the mode picker
            # aria="Open mode picker" - this is the key!
            mode_picker = page.locator('[aria-label="Open mode picker"]')
            if mode_picker.count() > 0:
                log("[FOUND] Mode picker button (aria='Open mode picker')")
                mode_picker_text = mode_picker.text_content()
                log(f"Current mode: '{mode_picker_text}'")
                screenshot(page, "04_before_mode_click")

                mode_picker.click()
                time.sleep(2)
                screenshot(page, "04b_mode_picker_open")
                log("Mode picker clicked - checking dropdown options...")

                # Get dropdown options
                dropdown_items = page.evaluate("""
                    () => {
                        const items = document.querySelectorAll(
                            '[role="menuitem"], [role="option"], [role="listitem"], ' +
                            'li, [class*="option"], [class*="item"]'
                        );
                        return Array.from(items)
                            .filter(el => {
                                const rect = el.getBoundingClientRect();
                                return rect.width > 0 && rect.height > 0;
                            })
                            .map(el => ({
                                text: el.textContent?.trim()?.substring(0, 100),
                                aria: el.getAttribute('aria-label'),
                                role: el.getAttribute('role'),
                                class: el.className?.toString()?.substring(0, 60),
                            }));
                    }
                """)

                log(f"Dropdown items found: {len(dropdown_items)}")
                for item in dropdown_items:
                    if item.get('text') or item.get('aria'):
                        log(f"  Option: text='{item.get('text', '')}' aria='{item.get('aria', '')}'")

                workflow_steps.append({
                    "step": 4,
                    "action": "open_mode_picker",
                    "selector": '[aria-label="Open mode picker"]',
                    "code": 'page.locator(\'[aria-label="Open mode picker"]\').click()',
                    "result": f"Dropdown opened with {len(dropdown_items)} items"
                })

                # Look for 2.5 Pro or Deep Think in dropdown
                for item in dropdown_items:
                    item_text = (item.get('text') or '').lower()
                    if '2.5' in item_text or 'pro' in item_text or 'deep' in item_text:
                        log(f"[TARGET] Found: '{item.get('text')}'")

                # Try to click on Deep Think or 2.5 Pro option
                deep_think_option = page.locator('[role="menuitem"]:has-text("Deep Think"), [role="option"]:has-text("Deep Think")').first
                pro_25_option = page.locator('[role="menuitem"]:has-text("2.5 Pro"), [role="option"]:has-text("2.5 Pro")').first

                if deep_think_option.count() > 0:
                    log("[OK] Found Deep Think option in dropdown!")
                    deep_think_option.click()
                    time.sleep(2)
                    screenshot(page, "04c_deep_think_selected")
                    workflow_steps.append({
                        "step": 5,
                        "action": "select_deep_think",
                        "selector": '[role="menuitem"]:has-text("Deep Think")',
                        "code": 'page.locator(\'[role="menuitem"]:has-text("Deep Think")\').click()',
                    })
                elif pro_25_option.count() > 0:
                    log("[OK] Found 2.5 Pro option in dropdown!")
                    pro_25_option.click()
                    time.sleep(2)
                    screenshot(page, "04c_25pro_selected")
                    workflow_steps.append({
                        "step": 5,
                        "action": "select_25_pro",
                        "selector": '[role="menuitem"]:has-text("2.5 Pro")',
                        "code": 'page.locator(\'[role="menuitem"]:has-text("2.5 Pro")\').click()',
                    })
                else:
                    log("[INFO] Specific options not found - closing dropdown")
                    page.keyboard.press("Escape")
                    time.sleep(1)

            else:
                log("[WARN] Mode picker not found via aria label")

                # Try the "Tools" button which was found
                tools_btn = page.locator('button:has-text("Tools")')
                if tools_btn.count() > 0:
                    log("[FOUND] Tools button")
                    tools_btn.click()
                    time.sleep(2)
                    screenshot(page, "04_tools_open")

            screenshot(page, "05_after_mode_selection")

            # STEP 5: Send the test prompt
            log("\nSTEP 5: Sending test prompt...")

            TEST_PROMPT = "What is the single most important strategic decision for a new AI voice widget product targeting Australian digital agencies in 2026?"

            # Find input area
            input_area = page.locator('div[contenteditable="true"]').first
            if input_area.count() > 0:
                log("[OK] Found contenteditable input area")
                input_area.click()
                time.sleep(0.5)
                input_area.fill(TEST_PROMPT)
                time.sleep(1)
                screenshot(page, "06_prompt_entered")
                log(f"Prompt entered: {TEST_PROMPT[:50]}...")

                workflow_steps.append({
                    "step": 6,
                    "action": "enter_prompt",
                    "selector": 'div[contenteditable="true"]',
                    "prompt": TEST_PROMPT,
                    "code": f'page.locator(\'div[contenteditable="true"]\').first.fill("{TEST_PROMPT[:50]}...")',
                })

                # Submit prompt
                send_btn = page.locator('[aria-label="Send message"]')
                if send_btn.count() > 0:
                    log("[OK] Found send button")
                    send_btn.click()
                    log("Prompt submitted!")
                    workflow_steps.append({
                        "step": 7,
                        "action": "submit_prompt",
                        "selector": '[aria-label="Send message"]',
                        "code": 'page.locator(\'[aria-label="Send message"]\').click()',
                    })
                else:
                    log("[FALLBACK] Using Enter key to submit")
                    input_area.press("Enter")
                    workflow_steps.append({
                        "step": 7,
                        "action": "submit_prompt",
                        "method": "Enter key",
                        "code": 'page.locator(\'div[contenteditable="true"]\').first.press("Enter")',
                    })

                time.sleep(3)
                screenshot(page, "07_after_submit")

                # STEP 6: Wait for response
                log("\nSTEP 6: Waiting for Deep Think response (up to 120 seconds)...")

                response_text = ""
                for i in range(40):  # 40 x 3s = 120s
                    time.sleep(3)

                    # Check if loading indicator present
                    loading_indicators = page.locator(
                        '[aria-label*="Loading"], [aria-label*="loading"], '
                        '.loading, [class*="loading"], [class*="spinner"], '
                        '[aria-busy="true"]'
                    ).count()

                    # Try to get response text
                    response_containers = page.evaluate("""
                        () => {
                            // Try multiple selectors for response content
                            const selectors = [
                                'model-response',
                                '[data-response-index]',
                                '.response-content',
                                'message-content',
                                '[class*="response"]',
                                'p',
                            ];

                            for (const sel of selectors) {
                                const els = document.querySelectorAll(sel);
                                if (els.length > 0) {
                                    const texts = Array.from(els)
                                        .map(el => el.textContent?.trim())
                                        .filter(t => t && t.length > 50);
                                    if (texts.length > 0) {
                                        return {selector: sel, texts: texts.slice(-5)};
                                    }
                                }
                            }
                            return null;
                        }
                    """)

                    if i % 5 == 0:
                        log(f"  Waiting... {i*3}s elapsed, loading={loading_indicators}")
                        screenshot(page, f"08_waiting_{i:02d}")

                    if loading_indicators == 0 and i > 3:
                        log(f"[OK] Response complete after ~{i*3}s")
                        if response_containers:
                            log(f"Response found via selector: {response_containers['selector']}")
                            response_text = "\n".join(response_containers['texts'])
                        break

                if not response_text:
                    # Fallback extraction
                    response_text = page.evaluate("""
                        () => {
                            const main = document.querySelector('main, [role="main"], .conversation-container');
                            return (main || document.body).innerText?.substring(0, 10000) || '';
                        }
                    """)

                screenshot(page, "09_final_response")
                log(f"\nRESPONSE PREVIEW (first 300 chars):\n{response_text[:300]}")

            else:
                log("[WARN] Could not find input area!")
                response_text = "Input area not found"

            # Save workflow documentation
            workflow_doc = {
                "title": "Gemini Deep Think Real Browser Workflow",
                "date": datetime.now().isoformat(),
                "url": "https://gemini.google.com",
                "steps": workflow_steps,
                "ui_elements_found": {
                    "mode_picker": '[aria-label="Open mode picker"]',
                    "chat_input": 'div[contenteditable="true"]',
                    "send_button": '[aria-label="Send message"]',
                    "tools_button": 'button:has-text("Tools")',
                    "write_button": 'button[aria-label="Write, button, tap to use tool"]',
                    "research_button": 'button[aria-label="Research, button, tap to use tool"]',
                },
                "response_text": response_text,
                "screenshots": sorted(os.listdir(SCREENSHOTS_DIR)),
                "step_log": step_log,
            }

            workflow_doc_path = "/mnt/e/genesis-system/deep_think_results/workflow_doc.json"
            with open(workflow_doc_path, 'w') as f:
                json.dump(workflow_doc, f, indent=2, default=str)
            log(f"\nWorkflow documentation saved: {workflow_doc_path}")

            return {
                "status": "success",
                "response_text": response_text,
                "workflow_steps": workflow_steps,
                "screenshots": sorted(os.listdir(SCREENSHOTS_DIR)),
            }

        except Exception as e:
            log(f"Main execution error: {e}", "ERROR")
            import traceback
            traceback.print_exc()
            try:
                screenshot(page, "ERROR_state")
            except:
                pass
            return {"status": "error", "error": str(e)}

        finally:
            log("Closing browser...")
            try:
                browser.close()
            except:
                pass


if __name__ == "__main__":
    result = run()
    print(f"\n[FINAL RESULT STATUS] {result.get('status')}")
    if result.get('response_text'):
        print(f"\n[RESPONSE]\n{result['response_text'][:1000]}")