#!/usr/bin/env python3
"""chatgpt-vision — analyze images using your ChatGPT subscription, no API key.

Drives your logged-in ChatGPT browser via chrome-use, uploads an image,
sends a prompt, and returns ChatGPT's text response.

Usage:
    chatgpt-vision "what's in this image?" -i photo.jpg
    chatgpt-vision "read the text" -i screenshot.png
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any

AB_BIN_CANDIDATES = ("chrome-use", "agent-browser", "agent-browser-stealth", "abs")

JsonDict = dict[str, Any]

# ---------- chrome-use helpers ----------

def _find_ab() -> str:
    for name in AB_BIN_CANDIDATES:
        found = shutil.which(name)
        if found:
            return found
    sys.exit(
        "chrome-use not found on PATH.\n"
        "Install: curl -fsSL https://raw.githubusercontent.com/leeguooooo/chrome-use/main/install.sh | sh"
    )


def _ab(ab: str, *ab_args: str, session: str = "chatgpt-vision",
        timeout: float = 30.0, profile: str | None = None) -> str:
    """Run one chrome-use subcommand, return stdout. Raises RuntimeError on failure.
    
    ``profile`` (a Chrome profile name) is passed as chrome-use's global ``--profile``,
    which must precede the subcommand. Only the launching ``open`` call needs it.
    """
    cmd = [ab]
    if profile:
        cmd += ["--profile", profile]
    cmd += [*ab_args, "--session", session]
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=max(5.0, timeout))
    except subprocess.TimeoutExpired:
        raise RuntimeError(f"chrome-use timed out: {' '.join(ab_args[:2])}")
    if proc.returncode != 0:
        tail = (proc.stderr or proc.stdout or "").strip().splitlines()[-3:]
        raise RuntimeError(
            f"chrome-use {ab_args[0] if ab_args else ''} failed "
            f"(exit {proc.returncode}): {' / '.join(tail)[:300]}"
        )
    return proc.stdout


def _ab_eval(ab: str, js: str, session: str = "chatgpt-vision",
             timeout: float = 20.0) -> Any:
    """Run JS in the page and decode the returned value.
    
    chrome-use wraps eval output in outer JSON quotes, so we double-parse:
    outer layer = the JSON string chrome-use prints,
    inner layer = the JSON value the JS returned via JSON.stringify().
    """
    raw = _ab(ab, "eval", js, session=session, timeout=timeout).strip()
    outer = json.loads(raw)  # chrome-use's wrapping
    if isinstance(outer, str):
        # The JS result was a JSON-stringified string or object — parse the inner JSON
        return json.loads(outer)
    return outer


# ---------- JS snippets (from chatgpt-imagegen's proven patterns) ----------

# Poll state: is streaming? any new images? any assistant text? rate-limited?
JS_STATE = r"""(() => {
  const stop = !!document.querySelector(
    'button[data-testid="stop-button"], button[aria-label*="Stop" i]'
  );
  const a = document.querySelectorAll('[data-message-author-role="assistant"]');
  const lastA = a[a.length - 1];
  const dlg = [...document.querySelectorAll('[role="dialog"]')]
    .map(d => d.textContent || '').join(' ');
  return JSON.stringify({
    stop,
    assistant: a.length > 0,
    acount: a.length,
    atext: lastA ? lastA.textContent.trim() : "",
    limited: /too many requests|requests too quickly/i.test(dlg)
  });
})()"""

JS_COMPOSER_EMPTY = r"""(() => {
  const t = (document.querySelector('#prompt-textarea') || {}).textContent || '';
  return JSON.stringify(t.trim().length === 0);
})()"""


# ---------- image reference helpers ----------

def _resolve_ref(path_or_url: str) -> tuple[str, str | None]:
    """Return (local_file_path, cleanup_path_or_None)."""
    if path_or_url.startswith(("http://", "https://", "data:")):
        resp = urllib.request.urlopen(path_or_url, timeout=30)
        suffix = Path(urllib.parse.urlparse(path_or_url).path).suffix or ".jpg"
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
        tmp.write(resp.read())
        tmp.close()
        return tmp.name, tmp.name
    return path_or_url, None


# ---------- main ----------

def analyze(ab: str, image_path: str, prompt: str, session: str = "chatgpt-vision",
            profile: str | None = None, keep_tab: bool = False) -> str:
    """Upload image, send prompt, return ChatGPT's text analysis."""

    deadline = time.monotonic() + 180  # 3 min total budget

    def remaining() -> float:
        return max(1.0, deadline - time.monotonic())

    # 1. Open ChatGPT
    print("   opening ChatGPT chat...", file=sys.stderr)
    _ab(ab, "open", "https://chatgpt.com/", session=session, profile=profile, timeout=30)
    time.sleep(5)

    if remaining() < 10:
        raise RuntimeError("timeout before submitting prompt")

    # 2. Check page loaded (no rate-limit dialog on initial load)
    st = _ab_eval(ab, JS_STATE, session=session, timeout=15)
    if st.get("limited"):
        raise RuntimeError("ChatGPT rate-limited before we even started. Try again later.")

    # 3. Upload the reference image
    print(f"   uploading image...", file=sys.stderr)
    _ab(ab, "upload", 'input[accept="image/*"]', image_path,
        session=session, timeout=min(90.0, remaining()))
    time.sleep(2)

    # 4. Focus composer and type the prompt
    print(f"   sending prompt...", file=sys.stderr)
    _ab(ab, "click", "#prompt-textarea", session=session, timeout=min(20, remaining()))
    _ab(ab, "keyboard", "type", prompt, session=session, timeout=min(60, remaining()))

    # 5. Submit — Enter first, send-button fallback
    sent = False
    for _ in range(6):
        _ab(ab, "press", "Enter", session=session, timeout=min(20, remaining()))
        time.sleep(1.5)
        try:
            empty = _ab_eval(ab, JS_COMPOSER_EMPTY, session=session, timeout=10)
        except RuntimeError:
            empty = False
        if empty:
            sent = True
            break
        try:
            _ab(ab, "click", 'button[data-testid="send-button"]',
                session=session, timeout=min(15, remaining()))
        except RuntimeError:
            pass
        time.sleep(1.5)
        try:
            empty = _ab_eval(ab, JS_COMPOSER_EMPTY, session=session, timeout=10)
        except RuntimeError:
            empty = False
        if empty:
            sent = True
            break

    if not sent:
        raise RuntimeError("failed to send prompt — composer never cleared")

    # 6. Poll for response — wait for text to stabilize (no changes for ~6s)
    print(f"   waiting for response...", file=sys.stderr)
    last_text = ""
    stable_count = 0
    STABLE_THRESHOLD = 3  # consecutive polls with same text = done
    first_text_seen = False

    while time.monotonic() < deadline:
        time.sleep(2.0)
        try:
            st = _ab_eval(ab, JS_STATE, session=session, timeout=15)
        except RuntimeError:
            continue

        if not isinstance(st, dict):
            continue

        if st.get("limited"):
            raise RuntimeError("ChatGPT rate-limited mid-response. Try again later.")

        acount = st.get("acount", 0)
        atext = st.get("atext", "")

        if acount > 0 and atext:
            if not first_text_seen:
                first_text_seen = True
                print("   receiving response...", file=sys.stderr)

            if atext != last_text:
                # Text changed — reset stability counter
                last_text = atext
                stable_count = 0
            else:
                # Same text as last poll
                stable_count += 1

            if stable_count >= STABLE_THRESHOLD:
                print("   response complete!", file=sys.stderr)
                return last_text
        else:
            # No assistant text yet — still waiting
            stable_count = 0

    if last_text:
        return last_text
    raise RuntimeError(f"timed out waiting for ChatGPT's response")


def main():
    parser = argparse.ArgumentParser(
        description="Analyze images using your ChatGPT subscription — no API key needed.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  chatgpt-vision \"what's in this image?\" -i photo.jpg\n"
            "  chatgpt-vision \"read the text\" -i screenshot.png\n"
        ),
    )
    parser.add_argument("prompt", nargs="?", help="Analysis prompt. Reads from stdin if omitted.")
    parser.add_argument("-i", "--image", "--ref", required=True, dest="image",
                        help="Path or URL to the image to analyze")
    parser.add_argument("--keep-tab", action="store_true",
                        help="Leave the ChatGPT tab open after analysis")
    parser.add_argument("--profile", default=None,
                        help="Chrome profile name (only needed for the first open call)")
    args = parser.parse_args()

    prompt = args.prompt
    if not prompt:
        prompt = sys.stdin.read().strip()
    if not prompt:
        print("error: a prompt is required", file=sys.stderr)
        parser.print_help()
        sys.exit(1)

    img_path, cleanup = _resolve_ref(args.image)
    if img_path != args.image:
        print(f"   downloaded {args.image}", file=sys.stderr)

    ab = _find_ab()
    session = "chatgpt-vision"

    try:
        result = analyze(ab, img_path, prompt, session=session,
                         profile=args.profile, keep_tab=args.keep_tab)
        print(result)  # stdout = the analysis
    finally:
        if cleanup:
            try:
                os.unlink(cleanup)
            except OSError:
                pass
        if not args.keep_tab:
            try:
                _ab(ab, "close", "--all", session=session, timeout=10)
            except RuntimeError:
                pass


if __name__ == "__main__":
    main()
