app/services/container_search.py

"""Container-friendly multi-platform AI search script.

Directly uses Playwright to start Chromium in headless mode, navigate to
AI chat platforms that support anonymous access, submit a query, and
collect the response. Designed for Hugging Face Docker Spaces.

Supported platforms (no login required):
    doubao, qwen, kimi, minimaxi
"""

import argparse
import json
import sys
import time
from pathlib import Path

from playwright.sync_api import sync_playwright

SITES = {
    "doubao": {
        "url": "https://www.doubao.com/chat/",
        "input_selectors": ['textarea', 'div[contenteditable="true"]'],
        "response_container": '[class*="message"]',
        "submit_via": "enter",
    },
    "qwen": {
        "url": "https://chat.qwen.ai/",
        "input_selectors": ['textarea', 'div[contenteditable="true"]'],
        "response_container": '[class*="message"]',
        "submit_via": "enter",
    },
    "kimi": {
        "url": "https://www.kimi.com/",
        "input_selectors": ['textarea', 'div[contenteditable="true"]'],
        "response_container": '[class*="message"]',
        "submit_via": "enter",
    },
    "minimaxi": {
        "url": "https://agent.minimaxi.com/",
        "input_selectors": ['textarea', 'div[contenteditable="true"]'],
        "response_container": '[class*="message"]',
        "submit_via": "enter",
    },
}


def _find_input(page, selectors: list[str]):
    """Find the first visible input/textarea element."""
    for sel in selectors:
        try:
            els = page.locator(sel).all()
            for el in els:
                try:
                    if el.is_visible(timeout=2000):
                        return el
                except Exception:
                    continue
        except Exception:
            continue
    return None


def _dismiss_overlays(page):
    """Try to close popups/modals/cookie banners."""
    dismiss_texts = ["知道了", "我知道了", "关闭", "OK", "确定", "Got it", "Accept",
                     "开始体验", "同意", "允许", "稍后再说", "以后再说", "跳过"]
    for text in dismiss_texts:
        try:
            btn = page.locator(f'button:has-text("{text}")').first
            if btn.is_visible(timeout=500):
                btn.click()
                time.sleep(0.5)
        except Exception:
            continue


def _extract_response_text(page, site: str) -> str:
    """Extract the latest AI response from the page."""
    # Strategy 1: Get all text blocks that look like responses
    try:
        # Wait a moment for content to render
        time.sleep(2)
        body_text = page.inner_text("body")
        return body_text
    except Exception:
        return ""


def _is_response_complete(page, prev_length: int, stable_count: int) -> tuple[bool, int, int]:
    """Check if the response has stopped growing."""
    try:
        body_text = page.inner_text("body")
        current_length = len(body_text)
        if current_length > prev_length + 20:
            return False, current_length, 0
        return stable_count >= 4, current_length, stable_count + 1
    except Exception:
        return stable_count >= 4, prev_length, stable_count + 1


def run_search(site: str, query: str, output_path: str, timeout: int = 120) -> dict:
    """Run a search on a single AI platform using headless Chromium."""
    config = SITES.get(site)
    if not config:
        return {"ok": False, "error": f"Unknown site: {site}. Available: {list(SITES.keys())}"}

    with sync_playwright() as p:
        browser = p.chromium.launch(
            headless=True,
            executable_path="/usr/bin/chromium",
            args=[
                "--no-sandbox",
                "--disable-gpu",
                "--disable-dev-shm-usage",
                "--disable-setuid-sandbox",
                "--no-first-run",
                "--no-zygote",
                "--single-process",
                "--disable-extensions",
                "--disable-software-rasterizer",
                "--disable-background-networking",
            ],
        )
        context = browser.new_context(
            viewport={"width": 1440, "height": 960},
            user_agent=(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/125.0.0.0 Safari/537.36"
            ),
            locale="zh-CN",
        )
        page = context.new_page()

        try:
            # Navigate
            print(f"[{site}] Navigating to {config['url']} ...", flush=True)
            page.goto(config["url"], wait_until="domcontentloaded", timeout=60000)
            time.sleep(3)

            # Dismiss overlays
            _dismiss_overlays(page)

            # Find input
            input_el = _find_input(page, config["input_selectors"])
            if not input_el:
                debug_path = str(Path(output_path).parent / f"{site}_debug.png")
                page.screenshot(path=debug_path)
                return {
                    "ok": False,
                    "error": f"Could not find input on {site}",
                    "debug": debug_path,
                    "title": page.title(),
                    "url": page.url,
                }

            # Record page text length before submitting
            try:
                pre_text = page.inner_text("body")
                pre_length = len(pre_text)
            except Exception:
                pre_length = 0

            # Type and submit
            print(f"[{site}] Submitting query ...", flush=True)
            input_el.click()
            time.sleep(0.5)
            input_el.fill(query)
            time.sleep(1)
            input_el.press("Enter")
            print(f"[{site}] Query sent, waiting for response ...", flush=True)

            # Wait for response to complete
            time.sleep(5)  # Initial wait
            prev_length = pre_length
            stable_count = 0

            for _ in range(timeout // 3):
                done, prev_length, stable_count = _is_response_complete(page, prev_length, stable_count)
                if done:
                    break
                time.sleep(3)

            # Extract response
            raw_text = _extract_response_text(page, site)

            # Try to extract just the new content (after the query)
            answer = raw_text
            query_pos = raw_text.rfind(query)
            if query_pos >= 0:
                after_query = raw_text[query_pos + len(query):].strip()
                if len(after_query) > 20:
                    answer = after_query

            if answer.strip():
                Path(output_path).write_text(answer.strip(), encoding="utf-8")
                print(f"[{site}] Got response ({len(answer)} chars)", flush=True)
                return {"ok": True, "answer": answer[:500], "length": len(answer)}
            else:
                debug_path = str(Path(output_path).parent / f"{site}_debug.png")
                page.screenshot(path=debug_path)
                return {"ok": False, "error": "No response captured", "debug": debug_path}

        except Exception as e:
            debug_path = str(Path(output_path).parent / f"{site}_debug.png")
            try:
                page.screenshot(path=debug_path)
            except Exception:
                pass
            return {"ok": False, "error": f"{type(e).__name__}: {e}"}
        finally:
            browser.close()


def main():
    parser = argparse.ArgumentParser(description="Container-friendly AI search (no login required)")
    parser.add_argument("--site", required=True, choices=sorted(SITES.keys()))
    parser.add_argument("--query", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--timeout", type=int, default=120)
    parser.add_argument("--test", action="store_true", help="Run browser test only")
    args = parser.parse_args()

    if args.test:
        # Diagnostic mode - just test if Playwright + Chromium work
        try:
            with sync_playwright() as p:
                browser = p.chromium.launch(
                    headless=True,
                    executable_path="/usr/bin/chromium",
                    args=["--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"],
                )
                page = browser.new_page()
                page.goto("https://example.com", timeout=15000)
                title = page.title()
                browser.close()
                print(json.dumps({"ok": True, "title": title}))
                return 0
        except Exception as e:
            print(json.dumps({"ok": False, "error": f"{type(e).__name__}: {e}"}))
            return 1

    result = run_search(args.site, args.query, args.output, args.timeout)
    print(json.dumps(result, ensure_ascii=False))
    return 0 if result.get("ok") else 1


if __name__ == "__main__":
    raise SystemExit(main())