#!/usr/bin/env python3
"""
REAL Gemini Deep Think Workflow - v3
Correct Chrome paths discovered:
- Executable: /mnt/c/Program Files (x86)/Google/Chrome/Application/chrome.exe
- User Data: /mnt/c/Users/P3/AppData/Local/Google/Chrome/User Data
"""

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"
CHROME_EXE = "/mnt/c/Program Files (x86)/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 = []
workflow_steps = []

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

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

def get_interactive_elements(page):
    """Get all visible interactive elements from the page"""
    return page.evaluate("""
        () => {
            const els = document.querySelectorAll(
                'button, [role="button"], a[href], [role="switch"], ' +
                '[role="checkbox"], [role="combobox"], [aria-haspopup], ' +
                'textarea, [contenteditable="true"], input'
            );
            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 || '',
                    'data-test-id': el.getAttribute('data-test-id') || '',
                    visible: rect.width > 0 && rect.height > 0,
                    y: Math.round(rect.y),
                    x: Math.round(rect.x),
                };
            }).filter(el => el.visible && (el.text || el.aria));
        }
    """)

def run():
    log("=" * 60)
    log("GEMINI DEEP THINK BROWSER AUTOMATION v3")
    log(f"Started: {datetime.now().isoformat()}")
    log(f"Chrome: {CHROME_EXE}")
    log(f"Profile: {CHROME_USER_DATA}")
    log("=" * 60)

    result = {
        "status": "unknown",
        "logged_in": False,
        "deep_think_found": False,
        "model_25_pro_found": False,
        "prompt_sent": False,
        "response_text": "",
        "ui_elements": {},
        "workflow_steps": [],
        "screenshots": [],
        "errors": [],
    }

    with sync_playwright() as p:

        # ---- STEP 1: Launch Chrome with real user profile ----
        log("\n[STEP 1] Launching Chrome with Windows user profile...")

        try:
            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-component-update",
                ],
                viewport={"width": 1440, "height": 900},
                ignore_default_args=["--enable-automation"],  # avoid bot detection
            )
            log("[OK] Chrome launched with real Windows profile")
            workflow_steps.append({
                "step": 1,
                "action": "Launch Chrome",
                "code": f'p.chromium.launch_persistent_context(user_data_dir="{CHROME_USER_DATA}", executable_path="{CHROME_EXE}", headless=False)',
                "note": "Uses real Chrome profile with existing session cookies"
            })

        except Exception as e:
            log(f"[ERROR] Chrome launch with profile failed: {e}", "ERROR")
            result["errors"].append(f"Chrome launch: {e}")
            log("[FALLBACK] Launching plain Chromium (no profile - no session)")
            browser = p.chromium.launch(
                headless=False,
                slow_mo=200,
                args=["--disable-blink-features=AutomationControlled"],
            )
            workflow_steps.append({
                "step": 1,
                "action": "Launch Chromium (fallback)",
                "note": "Real Chrome profile failed - using bare Chromium (will need login)"
            })

        # Get or create page
        if hasattr(browser, 'pages') and browser.pages:
            page = browser.pages[0]
        elif hasattr(browser, 'new_page'):
            page = browser.new_page()
        else:
            page = browser.new_page()

        try:
            # ---- STEP 2: Navigate to Gemini ----
            log("\n[STEP 2] Navigating to gemini.google.com...")
            page.goto("https://gemini.google.com", timeout=45000, wait_until="domcontentloaded")
            time.sleep(6)  # Allow JS to fully load
            screenshot(page, "01_gemini_loaded")

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

            workflow_steps.append({
                "step": 2,
                "action": "Navigate to Gemini",
                "url": "https://gemini.google.com",
                "code": 'page.goto("https://gemini.google.com", timeout=45000)',
                "note": "Wait 6s after domcontentloaded for full JS initialization"
            })

            # ---- STEP 3: Detect login state ----
            log("\n[STEP 3] Detecting login state...")

            # Multiple strategies to detect login
            elements = get_interactive_elements(page)
            log(f"Total interactive elements: {len(elements)}")

            # Log all elements for debugging
            for el in elements:
                log(f"  ELEM: tag={el['tag']} text='{el['text'][:50]}' aria='{el['aria'][:50]}'")

            # Check login indicators
            has_signin = any(
                'sign in' in (el.get('text') or '').lower() or
                'sign in' in (el.get('aria') or '').lower()
                for el in elements
            )

            has_profile = any(
                any(x in (el.get('aria') or '').lower()
                    for x in ['google account', 'account', 'profile', 'menu'])
                for el in elements
            )

            has_chat_input = any(
                el['tag'] in ['TEXTAREA'] or
                'contenteditable' in el.get('class', '').lower() or
                any(x in (el.get('aria') or '').lower() for x in ['message', 'chat', 'input', 'type'])
                for el in elements
            )

            log(f"Sign in button present: {has_signin}")
            log(f"Profile element present: {has_profile}")
            log(f"Chat input present: {has_chat_input}")

            # Look specifically for mode picker (appears in logged-in state)
            mode_picker_elements = [el for el in elements if
                'mode picker' in (el.get('aria') or '').lower() or
                'fast' in (el.get('text') or '').lower() or
                '2.5' in (el.get('text') or '').lower()
            ]

            is_logged_in = len(mode_picker_elements) > 0 or (has_chat_input and not has_signin)
            result["logged_in"] = is_logged_in
            log(f"Login status: {'LOGGED IN' if is_logged_in else 'NOT LOGGED IN'}")

            if not is_logged_in:
                screenshot(page, "02_NOT_LOGGED_IN")
                log("[CRITICAL] Not logged in. Chrome profile cookies didn't transfer.", "WARN")
                log("Manual intervention needed: Kinan needs to authenticate via browser", "WARN")

                # Document the state we found
                result["status"] = "needs_login"
                result["message"] = "Chrome session not active in Playwright context"
                result["workflow_steps"] = workflow_steps

                # Still continue to map UI structure for documentation
                log("Continuing to document UI structure for workflow reference...")

            else:
                log("[OK] LOGGED IN - session active!")
                screenshot(page, "02_LOGGED_IN")

            # ---- STEP 4: Map the full Gemini UI ----
            log("\n[STEP 4] Mapping Gemini UI structure...")
            screenshot(page, "03_full_ui")

            # Find key UI components
            ui_components = {
                "mode_picker": None,
                "model_selector": None,
                "deep_think_toggle": None,
                "chat_input": None,
                "send_button": None,
                "write_tool": None,
                "research_tool": None,
                "tools_button": None,
            }

            for el in elements:
                text_lower = (el.get('text') or '').lower()
                aria_lower = (el.get('aria') or '').lower()
                combined = text_lower + ' ' + aria_lower

                if 'mode picker' in aria_lower or 'fast' == text_lower:
                    ui_components["mode_picker"] = el
                    log(f"[FOUND] Mode picker: text='{el['text']}' aria='{el['aria']}'")

                if '2.5' in text_lower or 'pro' in text_lower or 'model' in aria_lower:
                    ui_components["model_selector"] = el

                if 'deep think' in combined or ('think' in combined and 'deep' in combined):
                    ui_components["deep_think_toggle"] = el
                    result["deep_think_found"] = True

                if el['tag'] in ['TEXTAREA', 'INPUT'] or 'contenteditable' in text_lower:
                    ui_components["chat_input"] = el

                if 'send message' in aria_lower:
                    ui_components["send_button"] = el

                if 'write' == text_lower or 'write, button' in aria_lower:
                    ui_components["write_tool"] = el

                if 'research' == text_lower or 'research, button' in aria_lower:
                    ui_components["research_tool"] = el

                if 'tools' == text_lower:
                    ui_components["tools_button"] = el

            result["ui_elements"] = {k: v for k, v in ui_components.items() if v}
            log(f"\nUI Components found: {[k for k,v in ui_components.items() if v]}")

            # Save UI map
            with open("/mnt/e/genesis-system/deep_think_results/ui_map_v3.json", 'w') as f:
                json.dump({"elements": elements, "components": ui_components}, f, indent=2, default=str)

            # ---- STEP 5: Click mode picker to see model options ----
            log("\n[STEP 5] Opening mode picker to see model/mode options...")

            if ui_components["mode_picker"]:
                mode_picker_el = ui_components["mode_picker"]
                log(f"Clicking mode picker: text='{mode_picker_el['text']}'")

                try:
                    # Click by position
                    page.mouse.click(
                        mode_picker_el['x'] + 10,
                        mode_picker_el['y'] + 10
                    )
                    time.sleep(2)
                    screenshot(page, "04_mode_picker_opened")

                    # Get dropdown items
                    dropdown = page.evaluate("""
                        () => {
                            // Look for visible dropdown/menu
                            const selectors = [
                                '[role="menu"] [role="menuitem"]',
                                '[role="listbox"] [role="option"]',
                                '[role="dialog"] button',
                                '.dropdown-item',
                                '[data-option-index]',
                            ];

                            for (const sel of selectors) {
                                const els = document.querySelectorAll(sel);
                                if (els.length > 0) {
                                    return {
                                        selector: sel,
                                        items: Array.from(els).map(el => ({
                                            text: el.textContent?.trim()?.substring(0, 100),
                                            aria: el.getAttribute('aria-label'),
                                            role: el.getAttribute('role'),
                                        }))
                                    };
                                }
                            }

                            // Fallback: all newly visible elements
                            const allEls = document.querySelectorAll('button, [role="button"]');
                            return {
                                selector: 'fallback',
                                items: Array.from(allEls)
                                    .filter(el => {
                                        const rect = el.getBoundingClientRect();
                                        return rect.width > 0 && rect.height > 0;
                                    })
                                    .map(el => ({
                                        text: el.textContent?.trim()?.substring(0, 80),
                                        aria: el.getAttribute('aria-label'),
                                        role: el.getAttribute('role'),
                                    }))
                                    .slice(0, 20)
                            };
                        }
                    """)

                    log(f"Dropdown selector: {dropdown.get('selector')}")
                    log(f"Dropdown items ({len(dropdown.get('items', []))}):")
                    for item in dropdown.get('items', []):
                        if item.get('text') or item.get('aria'):
                            log(f"  - text='{item.get('text', '')}' aria='{item.get('aria', '')}'")

                    # Check for Deep Think / 2.5 Pro in dropdown
                    items = dropdown.get('items', [])

                    deep_think_items = [i for i in items if
                        'deep think' in (i.get('text') or '').lower() or
                        'deep think' in (i.get('aria') or '').lower()
                    ]

                    pro_25_items = [i for i in items if
                        '2.5 pro' in (i.get('text') or '').lower() or
                        '2.5 pro' in (i.get('aria') or '').lower()
                    ]

                    if deep_think_items:
                        log(f"[FOUND] Deep Think option: {deep_think_items[0]}")
                        result["deep_think_found"] = True

                        # Click it
                        if dropdown.get('selector') != 'fallback':
                            deep_think_btn = page.locator(
                                f"{dropdown['selector']}"
                            ).filter(has_text="Deep Think").first
                            if deep_think_btn.count() > 0:
                                deep_think_btn.click()
                                time.sleep(2)
                                screenshot(page, "05_deep_think_selected")
                                log("[OK] Deep Think selected!")
                                workflow_steps.append({
                                    "step": 5,
                                    "action": "Select Deep Think",
                                    "selector": f'{dropdown["selector"]}:has-text("Deep Think")',
                                    "code": 'page.locator(\'[role="menuitem"]:has-text("Deep Think")\').click()',
                                })
                    elif pro_25_items:
                        log(f"[FOUND] 2.5 Pro option: {pro_25_items[0]}")
                        result["model_25_pro_found"] = True
                    else:
                        log("[INFO] Deep Think / 2.5 Pro not visible in dropdown")
                        # Close dropdown
                        page.keyboard.press("Escape")
                        time.sleep(0.5)

                    workflow_steps.append({
                        "step": 5,
                        "action": "Open mode picker",
                        "selector": '[aria-label="Open mode picker"]',
                        "code": 'page.locator(\'[aria-label="Open mode picker"]\').click()',
                        "found_items": [i.get('text', '') for i in items if i.get('text')],
                    })

                except Exception as e:
                    log(f"Mode picker interaction failed: {e}", "WARN")
                    result["errors"].append(f"Mode picker: {e}")
                    page.keyboard.press("Escape")

            # ---- STEP 6: Find chat input and send prompt ----
            log("\n[STEP 6] Finding chat input and sending test prompt...")
            screenshot(page, "06_before_input")

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

            # Try multiple input selectors
            input_selectors = [
                'div[contenteditable="true"]',
                'textarea',
                'rich-textarea textarea',
                '[role="textbox"]',
                '[aria-label*="message"]',
                '[aria-label*="Enter"]',
                '[placeholder*="message"]',
                'p[contenteditable="true"]',
            ]

            input_found = False
            for sel in input_selectors:
                try:
                    el = page.locator(sel).first
                    if el.count() > 0 and el.is_visible():
                        log(f"[OK] Found input via: {sel}")
                        el.click()
                        time.sleep(0.5)

                        # Use type for contenteditable, fill for textarea
                        if 'contenteditable' in sel or el.evaluate('el => el.contentEditable') == 'true':
                            page.keyboard.type(TEST_PROMPT, delay=20)
                        else:
                            el.fill(TEST_PROMPT)

                        time.sleep(1)
                        screenshot(page, "07_prompt_typed")
                        log(f"Prompt typed: '{TEST_PROMPT[:50]}...'")
                        input_found = True
                        result["prompt_sent"] = True

                        workflow_steps.append({
                            "step": 6,
                            "action": "Enter prompt",
                            "selector": sel,
                            "prompt": TEST_PROMPT,
                            "code": f'page.locator("{sel}").first.click(); page.keyboard.type(prompt)',
                        })
                        break
                except Exception as e:
                    log(f"Input selector {sel} failed: {e}", "WARN")

            if not input_found:
                log("[ERROR] Could not find chat input!", "ERROR")
                result["status"] = "no_input"
                result["errors"].append("No chat input found")
            else:
                # Submit the prompt
                log("Submitting prompt...")

                send_btn = page.locator('[aria-label="Send message"]')
                if send_btn.count() > 0 and send_btn.is_visible():
                    send_btn.click()
                    log("[OK] Clicked Send button")
                    workflow_steps.append({
                        "step": 7,
                        "action": "Submit prompt",
                        "selector": '[aria-label="Send message"]',
                        "code": 'page.locator(\'[aria-label="Send message"]\').click()',
                    })
                else:
                    page.keyboard.press("Enter")
                    log("[OK] Pressed Enter to submit")
                    workflow_steps.append({
                        "step": 7,
                        "action": "Submit prompt",
                        "method": "Enter key",
                        "code": 'page.keyboard.press("Enter")',
                    })

                time.sleep(4)
                screenshot(page, "08_after_submit")

                # ---- STEP 7: Wait for response ----
                log("\n[STEP 7] Waiting for response (up to 120 seconds)...")
                log("Note: Deep Think responses can take 60-120 seconds")

                max_wait = 120
                poll_interval = 3
                last_content = ""
                stable_count = 0

                for i in range(max_wait // poll_interval):
                    time.sleep(poll_interval)
                    elapsed = (i + 1) * poll_interval

                    # Get current page text to detect changes
                    current_content = page.evaluate("""
                        () => {
                            const response = document.querySelector(
                                'model-response, [data-response-index], .response-content, ' +
                                '[class*="response-container"]'
                            );
                            if (response) return response.innerText?.substring(0, 200);
                            return document.body.innerText?.substring(0, 500);
                        }
                    """) or ""

                    # Check for loading indicators
                    is_loading = page.evaluate("""
                        () => {
                            const loaders = document.querySelectorAll(
                                '[aria-busy="true"], [class*="loading"], [class*="spinner"], ' +
                                '.generating, [class*="generating"]'
                            );
                            return loaders.length;
                        }
                    """)

                    if current_content == last_content:
                        stable_count += 1
                    else:
                        stable_count = 0
                        last_content = current_content

                    if i % 5 == 0:  # Every 15 seconds
                        log(f"  {elapsed}s: loading_elements={is_loading}, stable_count={stable_count}")
                        screenshot(page, f"09_waiting_{i:02d}")

                    # Consider done if stable for 3 polls and not loading
                    if stable_count >= 3 and is_loading == 0 and elapsed > 15:
                        log(f"[OK] Response stable after {elapsed}s")
                        break

                screenshot(page, "10_response_final")

                # ---- STEP 8: Extract response ----
                log("\n[STEP 8] Extracting response text...")

                response_text = page.evaluate("""
                    () => {
                        // Try specific response containers first
                        const responseSelectors = [
                            'model-response',
                            '[data-response-index]',
                            '.response-content',
                            '[class*="response-container"]',
                            '[class*="model-response"]',
                            '.markdown-container',
                        ];

                        for (const sel of responseSelectors) {
                            const els = document.querySelectorAll(sel);
                            if (els.length > 0) {
                                const texts = Array.from(els)
                                    .map(el => el.innerText?.trim())
                                    .filter(t => t && t.length > 100);
                                if (texts.length > 0) {
                                    return {selector: sel, text: texts[texts.length - 1]};
                                }
                            }
                        }

                        // Fallback: get main content area
                        const main = document.querySelector('main, [role="main"]');
                        return {
                            selector: 'main fallback',
                            text: (main || document.body).innerText?.substring(0, 10000)
                        };
                    }
                """)

                if response_text and response_text.get('text'):
                    result["response_text"] = response_text['text']
                    result["response_selector"] = response_text.get('selector')
                    log(f"[OK] Response extracted via: {response_text.get('selector')}")
                    log(f"Response length: {len(response_text['text'])} chars")
                    log(f"\nPREVIEW:\n{response_text['text'][:500]}\n...")

                    sel_used = response_text.get('selector', 'model-response')
                    workflow_steps.append({
                        "step": 8,
                        "action": "Extract response",
                        "selector": sel_used,
                        "code": f'response = page.locator("{sel_used}").last.text_content()',
                    })
                    result["status"] = "success"
                else:
                    log("[WARN] Could not extract response text", "WARN")
                    result["status"] = "partial"

        except Exception as e:
            log(f"[CRITICAL ERROR] {e}", "ERROR")
            import traceback
            traceback.print_exc()
            result["errors"].append(str(e))
            result["status"] = "error"
            try:
                screenshot(page, "ERROR_state")
            except:
                pass

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

    # Save results
    result["workflow_steps"] = workflow_steps
    result["step_log"] = step_log
    result["screenshots"] = sorted(os.listdir(SCREENSHOTS_DIR))

    results_path = "/mnt/e/genesis-system/deep_think_results/v3_results.json"
    with open(results_path, 'w') as f:
        json.dump(result, f, indent=2, default=str)
    log(f"\nResults saved: {results_path}")

    return result


if __name__ == "__main__":
    result = run()
    print(f"\n{'='*60}")
    print(f"STATUS: {result.get('status')}")
    print(f"Logged in: {result.get('logged_in')}")
    print(f"Deep Think found: {result.get('deep_think_found')}")
    print(f"Prompt sent: {result.get('prompt_sent')}")
    print(f"Screenshots: {len(result.get('screenshots', []))}")
    if result.get('response_text'):
        print(f"\nRESPONSE (first 500 chars):\n{result['response_text'][:500]}")
    if result.get('errors'):
        print(f"\nErrors: {result['errors']}")
