app/google_ai_mode.py

import os
import asyncio
import json
import time
from pathlib import Path
from typing import NamedTuple

from playwright.async_api import async_playwright
from playwright.sync_api import BrowserContext, Page, TimeoutError, sync_playwright


class PageRequest(NamedTuple):
    page_number: int
    image_path: Path
    prompt: str
    screenshot_path: Path | None = None


class GoogleAiModeClient:
    def __init__(
        self,
        ai_mode_url: str,
        browser_profile_dir: Path,
        headless: bool | None = None,
        response_timeout_seconds: int = 90,
    ) -> None:
        self.ai_mode_url = ai_mode_url
        self.browser_profile_dir = browser_profile_dir
        self.headless = _browser_headless() if headless is None else headless
        self.response_timeout_seconds = response_timeout_seconds
        self._playwright = None
        self._context: BrowserContext | None = None

    def __enter__(self) -> "GoogleAiModeClient":
        self._playwright = sync_playwright().start()
        self.browser_profile_dir.mkdir(parents=True, exist_ok=True)

        launch_kwargs = {
            "user_data_dir": str(self.browser_profile_dir),
            "headless": self.headless,
            "args": [
                "--disable-blink-features=AutomationControlled",
                "--start-maximized",
            ],
        }
        proxy = _browser_proxy()
        if proxy:
            launch_kwargs["proxy"] = proxy

        try:
            self._context = self._playwright.chromium.launch_persistent_context(
                channel=os.getenv("BROWSER_CHANNEL", "chrome"),
                **launch_kwargs,
            )
        except Exception:
            self._context = self._playwright.chromium.launch_persistent_context(
                **launch_kwargs,
            )
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        if self._context:
            self._context.close()
        if self._playwright:
            self._playwright.stop()

    def run_batch(self, image_paths: list[Path], prompt: str) -> str:
        if not self._context:
            raise RuntimeError("Browser context has not started.")

        page = self._context.new_page()
        page.goto(self.ai_mode_url, wait_until="domcontentloaded", timeout=60_000)
        page.wait_for_timeout(4_000)
        self._dismiss_popups(page)
        self._attach_images(page, image_paths)
        self._enter_prompt(page, prompt)
        self._submit(page)
        return self._wait_for_json_response(page)

    def run_pages_parallel(self, requests: list[PageRequest]) -> dict[int, str]:
        if self._context:
            self._context.close()
            self._context = None
        if self._playwright:
            self._playwright.stop()
            self._playwright = None
        return asyncio.run(
            _run_pages_parallel_async(
                self.ai_mode_url,
                self.browser_profile_dir,
                self.headless,
                self.response_timeout_seconds,
                requests,
            )
        )

    def _dismiss_popups(self, page: Page) -> None:
        labels = ["I agree", "Accept all", "Got it", "Not now"]
        for label in labels:
            try:
                button = page.get_by_role("button", name=label)
                if button.count() == 1 and button.is_visible():
                    button.click(timeout=2_000)
            except Exception:
                continue

    def _attach_images(self, page: Page, image_paths: list[Path]) -> None:
        files = [str(path) for path in image_paths]
        debug_notes: list[str] = []
        page.wait_for_selector(
            'button[aria-label="More input options"], input[type="file"]',
            timeout=30_000,
        )
        file_inputs = page.locator('input[type="file"]')

        try:
            file_input_count = file_inputs.count()
            debug_notes.append(f"initial file inputs: {file_input_count}")
            if file_input_count > 0:
                self._set_files_one_by_one(page, files, file_inputs.first)
                return
        except Exception as exc:
            debug_notes.append(f"initial file input failed: {exc}")
            pass

        attach_candidates = [
            'button[aria-label*="Upload"]',
            'button[aria-label*="Attach"]',
            'button[aria-label*="Add"]',
            'div[role="button"][aria-label*="Upload"]',
            'div[role="button"][aria-label*="Attach"]',
            'div[role="button"][aria-label*="Add"]',
        ]

        for selector in attach_candidates:
            try:
                candidate = page.locator(selector)
                candidate_count = candidate.count()
                debug_notes.append(f"{selector}: {candidate_count}")
                if candidate_count < 1:
                    continue
                with page.expect_file_chooser(timeout=5_000) as chooser_info:
                    candidate.first.click(timeout=5_000)
                chooser_info.value.set_files(files)
                return
            except Exception as exc:
                debug_notes.append(f"{selector} failed: {exc}")
                continue

        menu_buttons = ["More input options", "Add files", "Add image"]
        upload_buttons = ["Upload image", "Upload file"]
        for menu_label in menu_buttons:
            try:
                menu_button = page.get_by_role("button", name=menu_label, exact=True)
                menu_count = menu_button.count()
                debug_notes.append(f"menu {menu_label}: {menu_count}")
                if menu_count < 1:
                    continue
                menu_button.first.click(timeout=5_000)
                page.wait_for_timeout(500)
                expanded_file_inputs = page.locator('input[type="file"]')
                expanded_count = expanded_file_inputs.count()
                debug_notes.append(f"expanded file inputs: {expanded_count}")
                if expanded_count > 0:
                    self._set_files_one_by_one(page, files, expanded_file_inputs.last)
                    return
                for upload_label in upload_buttons:
                    try:
                        upload_button = page.get_by_role(
                            "button",
                            name=upload_label,
                            exact=True,
                        )
                        upload_count = upload_button.count()
                        debug_notes.append(f"upload {upload_label}: {upload_count}")
                        if upload_count < 1:
                            continue
                        with page.expect_file_chooser(timeout=5_000) as chooser_info:
                            upload_button.first.click(timeout=5_000)
                        chooser_info.value.set_files(files)
                        return
                    except Exception as exc:
                        debug_notes.append(f"upload {upload_label} failed: {exc}")
                        continue
            except Exception as exc:
                debug_notes.append(f"menu {menu_label} failed: {exc}")
                continue

        raise RuntimeError(
            "Could not find an upload/attach control in AI Mode. "
            + " | ".join(debug_notes)
        )

    def _set_files_one_by_one(self, page: Page, files: list[str], file_input) -> None:
        for index, file_path in enumerate(files):
            if index > 0:
                menu_button = page.get_by_role(
                    "button",
                    name="More input options",
                    exact=True,
                )
                menu_button.first.click(timeout=5_000)
                page.wait_for_timeout(500)
                file_input = page.locator('input[type="file"]').last
            file_input.set_input_files(file_path, timeout=10_000)
            page.wait_for_timeout(1_500)

    def _enter_prompt(self, page: Page, prompt: str) -> None:
        input_candidates = [
            'textarea',
            '[contenteditable="true"]',
            'input[aria-label*="Ask"]',
            'div[role="textbox"]',
        ]

        for selector in input_candidates:
            try:
                locator = page.locator(selector)
                if locator.count() < 1:
                    continue
                target = locator.last
                target.click(timeout=5_000)
                target.fill(prompt, timeout=10_000)
                return
            except Exception:
                try:
                    page.keyboard.insert_text(prompt)
                    return
                except Exception:
                    continue

        raise RuntimeError("Could not find the AI Mode prompt input.")

    def _submit(self, page: Page) -> None:
        submit_candidates = [
            'button[aria-label*="Send"]',
            'button[aria-label*="Submit"]',
            'div[role="button"][aria-label*="Send"]',
            'div[role="button"][aria-label*="Submit"]',
        ]

        for selector in submit_candidates:
            try:
                candidate = page.locator(selector)
                if candidate.count() > 0 and candidate.last.is_enabled():
                    candidate.last.click(timeout=5_000)
                    return
            except Exception:
                continue

        page.keyboard.press("Enter")

    def _wait_for_json_response(self, page: Page) -> str:
        page.wait_for_timeout(8_000)
        deadline = time.time() + self.response_timeout_seconds
        last_text = ""

        while time.time() < deadline:
            page.wait_for_timeout(2_000)
            try:
                last_text = page.locator("body").inner_text(timeout=5_000)
            except TimeoutError:
                continue

            json_text = self._find_last_json_object(last_text)
            if json_text and not self._is_prompt_placeholder_json(json_text):
                return json_text

        raise TimeoutError(
            f"Timed out waiting for a JSON response. Last page text began: {last_text[:500]}"
        )

    @staticmethod
    def _find_last_json_object(text: str) -> str | None:
        decoder = json.JSONDecoder()
        candidates: list[str] = []
        start_positions = [index for index, char in enumerate(text) if char == "{"]
        for start in start_positions:
            try:
                parsed, end = decoder.raw_decode(text[start:])
            except json.JSONDecodeError:
                continue
            if (
                isinstance(parsed, dict)
                and parsed.get("success") is True
                and isinstance(parsed.get("columns"), list)
                and isinstance(parsed.get("data"), list)
            ):
                candidates.append(text[start : start + end])
        if candidates:
            return candidates[-1]
        return None

    @staticmethod
    def _is_prompt_placeholder_json(text: str) -> bool:
        try:
            parsed = json.loads(text)
        except json.JSONDecodeError:
            return False
        return parsed.get("data") == [
            {
                "date": "DD/MM/YYYY",
                "description": "Full transaction narration exactly as shown",
                "voucher_type": "Payment or Receipt",
                "amount": "0.00",
                "closing": "0.00",
            }
        ]


async def _run_pages_parallel_async(
    ai_mode_url: str,
    browser_profile_dir: Path,
    headless: bool,
    response_timeout_seconds: int,
    requests: list[PageRequest],
) -> dict[int, str]:
    browser_profile_dir.mkdir(parents=True, exist_ok=True)
    return await _run_pages_parallel_async_once(
        ai_mode_url,
        browser_profile_dir,
        headless,
        response_timeout_seconds,
        requests,
    )


async def _run_pages_parallel_async_once(
    ai_mode_url: str,
    browser_profile_dir: Path,
    headless: bool,
    response_timeout_seconds: int,
    requests: list[PageRequest],
) -> dict[int, str]:
    async with async_playwright() as playwright:
        launch_kwargs = {
            "user_data_dir": str(browser_profile_dir),
            "headless": headless,
            "args": [
                "--disable-blink-features=AutomationControlled",
                "--start-maximized",
            ],
        }
        proxy = _browser_proxy()
        if proxy:
            launch_kwargs["proxy"] = proxy
        try:
            context = await playwright.chromium.launch_persistent_context(
                channel=os.getenv("BROWSER_CHANNEL", "chrome"),
                **launch_kwargs,
            )
        except Exception:
            context = await playwright.chromium.launch_persistent_context(**launch_kwargs)

        try:
            tasks = [
                _run_single_page_async(context, ai_mode_url, response_timeout_seconds, request)
                for request in requests
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            retry_requests = [
                request
                for request, result in zip(requests, results)
                if isinstance(result, Exception)
            ]
            if retry_requests:
                await asyncio.sleep(2)
            for request in retry_requests:
                retry_result = await _run_single_page_async(
                    context,
                    ai_mode_url,
                    response_timeout_seconds,
                    request,
                )
                results[requests.index(request)] = retry_result
        finally:
            await context.close()

    responses: dict[int, str] = {}
    errors: list[str] = []
    for request, result in zip(requests, results):
        if isinstance(result, Exception):
            errors.append(f"page {request.page_number}: {result}")
        else:
            responses[request.page_number] = result

    if errors:
        raise RuntimeError("; ".join(errors))
    return responses


async def _run_single_page_async(
    context,
    ai_mode_url: str,
    response_timeout_seconds: int,
    request: PageRequest,
) -> str:
    page = await context.new_page()
    try:
        await page.goto(ai_mode_url, wait_until="domcontentloaded", timeout=60_000)
        await page.wait_for_load_state("domcontentloaded", timeout=20_000)
        await page.wait_for_timeout(2_000)
        await _wait_for_ai_mode_controls_async(page)
        await _dismiss_popups_async(page)
        await _attach_image_async(page, request.image_path)
        await _enter_prompt_async(page, request.prompt)
        await _submit_async(page)
        await page.wait_for_timeout(5_000)
        if request.screenshot_path:
            request.screenshot_path.parent.mkdir(parents=True, exist_ok=True)
            await page.screenshot(path=str(request.screenshot_path), full_page=False)
        return await _wait_for_json_response_async(page, response_timeout_seconds)
    finally:
        await page.close()


async def _wait_for_ai_mode_controls_async(page) -> None:
    selector = (
        'textarea, [contenteditable="true"], div[role="textbox"], '
        'button[aria-label="More input options"], input[type="file"]'
    )
    deadline = time.time() + 30
    last_text = ""
    while time.time() < deadline:
        if await page.locator(selector).count() > 0:
            return
        try:
            last_text = (await page.locator("body").inner_text(timeout=2_000))[:600]
        except Exception:
            last_text = ""
        await page.wait_for_timeout(1_000)

    raise RuntimeError(
        "AI Mode controls were not available. "
        f"Visible page text: {last_text}"
    )


async def _dismiss_popups_async(page) -> None:
    labels = ["I agree", "Accept all", "Got it", "Not now"]
    for label in labels:
        try:
            button = page.get_by_role("button", name=label)
            if await button.count() == 1 and await button.is_visible():
                await button.click(timeout=2_000)
        except Exception:
            continue


async def _attach_image_async(page, image_path: Path) -> None:
    file_path = str(image_path)
    errors: list[str] = []

    for attempt in range(1, 4):
        try:
            await _focus_ai_input_async(page)
            await page.wait_for_timeout(400)

            file_inputs = page.locator('input[type="file"]')
            if await file_inputs.count() > 0:
                await file_inputs.last.set_input_files(file_path, timeout=10_000)
                await page.wait_for_timeout(1_000)
                return

            opened = await _open_more_input_options_async(page)
            if not opened:
                errors.append(f"attempt {attempt}: no clickable More input options button")
                await page.wait_for_timeout(1_500)
                continue
            await page.wait_for_timeout(1_000)

            expanded_inputs = page.locator('input[type="file"]')
            if await expanded_inputs.count() > 0:
                await expanded_inputs.last.set_input_files(file_path, timeout=10_000)
                await page.wait_for_timeout(1_000)
                return

            upload_button = page.locator(
                '[aria-label="Upload image"], [aria-label="Upload file"], '
                '[role="button"]:has-text("Upload image"), [role="button"]:has-text("Upload file"), '
                'button:has-text("Upload image"), button:has-text("Upload file")'
            )
            if await upload_button.count() > 0:
                async with page.expect_file_chooser(timeout=7_000) as chooser_info:
                    await upload_button.last.click(timeout=7_000)
                chooser = await chooser_info.value
                await chooser.set_files(file_path)
                await page.wait_for_timeout(1_000)
                return

            errors.append(f"attempt {attempt}: menu opened but no upload input/button")
        except Exception as exc:
            errors.append(f"attempt {attempt}: {exc}")
        await page.wait_for_timeout(1_500)

    try:
        visible_text = (await page.locator("body").inner_text(timeout=3_000))[:800]
    except Exception:
        visible_text = "Could not read page text."
    raise RuntimeError(
        "Could not find upload image/file button. "
        + " | ".join(errors)
        + f" | visible page text: {visible_text}"
    )


async def _open_more_input_options_async(page) -> bool:
    box = await page.evaluate(
        """
        () => {
          const targets = [
            ...document.querySelectorAll('[aria-label="More input options"]'),
            ...document.querySelectorAll('[aria-label*="input options" i]'),
            ...document.querySelectorAll('[aria-label*="Add" i]'),
            ...document.querySelectorAll('[aria-label*="Upload" i]')
          ];
          const unique = [...new Set(targets)];
          const visible = unique.find((element) => {
            const style = window.getComputedStyle(element);
            const rect = element.getBoundingClientRect();
            return (
              style.display !== "none" &&
              style.visibility !== "hidden" &&
              Number(style.opacity || "1") > 0 &&
              rect.width > 0 &&
              rect.height > 0
            );
          });
          const target = visible || unique[unique.length - 1];
          if (!target) return false;
          target.scrollIntoView({ block: "center", inline: "center" });
          const rect = target.getBoundingClientRect();
          return { x: rect.left + rect.width / 2, y: rect.top + rect.height / 2 };
        }
        """
    )
    if not box:
        return False
    await page.mouse.click(box["x"], box["y"])
    return True


async def _focus_ai_input_async(page) -> None:
    input_candidates = [
        'textarea',
        '[contenteditable="true"]',
        'input[aria-label*="Ask"]',
        'div[role="textbox"]',
    ]
    for selector in input_candidates:
        try:
            locator = page.locator(selector)
            if await locator.count() > 0:
                await locator.last.click(timeout=3_000)
                return
        except Exception:
            continue


async def _enter_prompt_async(page, prompt: str) -> None:
    input_candidates = [
        'textarea',
        '[contenteditable="true"]',
        'input[aria-label*="Ask"]',
        'div[role="textbox"]',
    ]
    for selector in input_candidates:
        try:
            locator = page.locator(selector)
            if await locator.count() < 1:
                continue
            target = locator.last
            await target.click(timeout=5_000)
            await target.fill(prompt, timeout=10_000)
            return
        except Exception:
            try:
                await page.keyboard.insert_text(prompt)
                return
            except Exception:
                continue
    raise RuntimeError("Could not find the AI Mode prompt input.")


async def _submit_async(page) -> None:
    submit_candidates = [
        'button[aria-label*="Send"]',
        'button[aria-label*="Submit"]',
        'div[role="button"][aria-label*="Send"]',
        'div[role="button"][aria-label*="Submit"]',
    ]
    for selector in submit_candidates:
        try:
            candidate = page.locator(selector)
            if await candidate.count() > 0 and await candidate.last.is_enabled():
                await candidate.last.click(timeout=5_000)
                return
        except Exception:
            continue
    await page.keyboard.press("Enter")


async def _wait_for_json_response_async(page, response_timeout_seconds: int) -> str:
    deadline = time.time() + response_timeout_seconds
    last_text = ""
    while time.time() < deadline:
        await page.wait_for_timeout(1_000)
        try:
            last_text = await page.locator("body").inner_text(timeout=5_000)
        except TimeoutError:
            continue
        json_text = GoogleAiModeClient._find_last_json_object(last_text)
        if json_text and not GoogleAiModeClient._is_prompt_placeholder_json(json_text):
            return json_text
    raise TimeoutError(
        f"Timed out waiting for a JSON response. Last page text began: {last_text[:500]}"
    )


def _browser_headless() -> bool:
    return os.getenv("BROWSER_HEADLESS", "true").strip().lower() not in {
        "0",
        "false",
        "no",
    }


def _browser_proxy() -> dict[str, str] | None:
    server = os.getenv("BROWSER_PROXY_SERVER", "").strip()
    if not server:
        return None
    proxy = {"server": server}
    username = os.getenv("BROWSER_PROXY_USERNAME", "").strip()
    password = os.getenv("BROWSER_PROXY_PASSWORD", "").strip()
    if username:
        proxy["username"] = username
    if password:
        proxy["password"] = password
    return proxy