"""
Gemini 3 Flash Vision Browser Agent
Think -> Act -> Observe loop using Playwright + Gemini 3 Flash agentic vision
Connects to Elestio Browserless OR runs local headless Playwright

Usage:
    python3 core/gemini_vision_agent.py --task "your task here" [--headless] [--max-steps 50]
    python3 core/gemini_vision_agent.py --godaddy --headless
    python3 core/gemini_vision_agent.py --elestio --headless
"""

import os
import sys
import json
import base64
import asyncio
import argparse
import logging
import re
import time
from datetime import datetime
from pathlib import Path

# Load secrets.env before anything else
def load_secrets():
    secrets_path = Path("/mnt/e/genesis-system/config/secrets.env")
    if secrets_path.exists():
        with open(secrets_path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, val = line.partition("=")
                    key = key.strip()
                    val = val.strip().strip('"').strip("'")
                    if key and val:
                        os.environ.setdefault(key, val)

load_secrets()

from google import genai as new_genai
from google.genai import types as genai_types
from playwright.async_api import async_playwright, Page, Browser, BrowserContext

try:
    from playwright_stealth import Stealth
    STEALTH_AVAILABLE = True
except ImportError:
    STEALTH_AVAILABLE = False

# ── Config ─────────────────────────────────────────────────────────────────────
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY_NEW") or os.environ.get("GEMINI_API_KEY")
BROWSERLESS_URL = os.environ.get("BROWSERLESS_URL")
BROWSERLESS_TOKEN = os.environ.get("BROWSERLESS_TOKEN")
VISION_MODEL = "gemini-3-flash-preview"

LOG_DIR = Path("/mnt/e/genesis-system/logs")
SCREENSHOT_DIR = LOG_DIR / "vision_screenshots"
LOG_FILE = LOG_DIR / "vision_agent.log"
LOG_DIR.mkdir(parents=True, exist_ok=True)
SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)

# ── Logging ────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(LOG_FILE, encoding="utf-8"),
        logging.StreamHandler(sys.stdout),
    ],
)
log = logging.getLogger("vision_agent")

# ── Prompt Template ────────────────────────────────────────────────────────────
VISION_PROMPT = """You are controlling a web browser to complete a task.
Current task: {task}
Previous actions (last 10): {history}
Step: {step}/{max_steps}
Current URL: {url}

Look at the screenshot carefully. Decide the SINGLE best next action.
Respond with ONLY a valid JSON object — no markdown, no code fences, no explanation.

AVAILABLE ACTIONS:
- Navigate:   {{"action": "navigate", "url": "https://...", "description": "reason"}}
- Click:      {{"action": "click", "x": 640, "y": 400, "description": "what you're clicking"}}
- Type:       {{"action": "type", "text": "text to type", "description": "reason"}}  (types into focused element)
- Fill field: {{"action": "fill", "selector": "#id or input[name=x]", "text": "value", "description": "reason"}}
- Press key:  {{"action": "key", "key": "Enter", "description": "reason"}}
- Scroll:     {{"action": "scroll", "direction": "down", "amount": 400, "description": "reason"}}
- Wait:       {{"action": "wait", "ms": 2000, "description": "waiting for page"}}
- Done:       {{"action": "complete", "result": "full description of what was accomplished"}}
- Blocked:    {{"action": "blocked", "reason": "exact description of what is blocking (2FA, captcha, bot check, login failed)"}}

IMPORTANT RULES:
- If you see a bot check / unusual browser warning that persists after retry, use "blocked"
- If login fails with wrong credentials, use "blocked" immediately
- If you've been on the same page doing the same action for 5+ steps, try a different approach
- When adding DNS records, confirm each record is saved before adding the next
- Never type into a field without first clicking it or using fill with selector
"""


# ── Gemini Vision Client ───────────────────────────────────────────────────────
class GeminiVisionClient:
    def __init__(self):
        if not GEMINI_API_KEY:
            raise RuntimeError(
                "No Gemini API key. Set GEMINI_API_KEY_NEW or GEMINI_API_KEY in secrets.env"
            )
        self.client = new_genai.Client(api_key=GEMINI_API_KEY)
        log.info(f"Gemini client ready — model: {VISION_MODEL}")

    def decide_action(self, task: str, history: list[str], screenshot_bytes: bytes,
                      step: int, max_steps: int, current_url: str = "") -> dict:
        history_text = "\n".join(history[-10:]) if history else "None yet"
        prompt_text = VISION_PROMPT.format(
            task=task,
            history=history_text,
            step=step,
            max_steps=max_steps,
            url=current_url,
        )

        try:
            response = self.client.models.generate_content(
                model=VISION_MODEL,
                contents=[
                    genai_types.Part.from_bytes(data=screenshot_bytes, mime_type="image/png"),
                    genai_types.Part.from_text(text=prompt_text),
                ],
            )
            raw = response.text.strip()
        except Exception as e:
            log.error(f"Gemini API error: {e}")
            return {"action": "wait", "ms": 3000, "description": f"API error: {str(e)[:100]}"}

        log.debug(f"Gemini raw: {raw[:300]}")

        # Strip markdown fences
        cleaned = re.sub(r"^```(?:json)?\s*", "", raw)
        cleaned = re.sub(r"\s*```$", "", cleaned).strip()

        try:
            return json.loads(cleaned)
        except json.JSONDecodeError:
            m = re.search(r"\{.*\}", cleaned, re.DOTALL)
            if m:
                try:
                    return json.loads(m.group())
                except json.JSONDecodeError:
                    pass
            log.warning(f"JSON parse failed for: {raw[:200]}")
            return {"action": "wait", "ms": 2000, "description": f"parse error"}


# ── Browser Agent ──────────────────────────────────────────────────────────────
class GeminiVisionAgent:
    def __init__(self, headless: bool = True, max_steps: int = 50,
                 viewport_width: int = 1280, viewport_height: int = 900,
                 use_stealth: bool = True):
        self.headless = headless
        self.max_steps = max_steps
        self.viewport = {"width": viewport_width, "height": viewport_height}
        self.use_stealth = use_stealth and STEALTH_AVAILABLE
        self.gemini = GeminiVisionClient()
        self.history: list[str] = []
        self.step = 0
        self.results: list[dict] = []

    def reset(self):
        self.history = []
        self.step = 0
        self.results = []

    async def _get_screenshot(self, page: Page) -> bytes:
        try:
            return await page.screenshot(type="png", full_page=False)
        except Exception as e:
            log.warning(f"Screenshot failed: {e}")
            return b""

    async def _save_screenshot(self, screenshot: bytes, label: str = "") -> Path:
        if not screenshot:
            return Path("/dev/null")
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        label_clean = re.sub(r"[^\w\-]", "_", label)[:40]
        fname = SCREENSHOT_DIR / f"step_{self.step:03d}_{ts}_{label_clean}.png"
        fname.write_bytes(screenshot)
        log.info(f"  Screenshot: {fname.name}")
        return fname

    async def _execute_action(self, page: Page, action: dict) -> str:
        act = action.get("action", "wait")
        desc = action.get("description", "")
        log.info(f"  EXEC [{act}] {desc}")

        if act == "navigate":
            url = action.get("url", "")
            try:
                await page.goto(url, wait_until="domcontentloaded", timeout=30000)
                try:
                    await page.wait_for_load_state("networkidle", timeout=8000)
                except Exception:
                    pass
            except Exception as e:
                log.warning(f"Navigate timeout/error (continuing): {e}")
            return f"navigated to {url}"

        elif act == "click":
            x, y = int(action.get("x", 0)), int(action.get("y", 0))
            await page.mouse.click(x, y)
            await asyncio.sleep(0.8)
            return f"clicked ({x},{y})"

        elif act == "fill":
            selector = action.get("selector", "")
            text = action.get("text", "")
            try:
                await page.fill(selector, text, timeout=5000)
                return f"filled '{selector}' with '{text[:40]}'"
            except Exception as e:
                log.warning(f"Fill failed for {selector}: {e}")
                return f"fill failed: {e}"

        elif act == "type":
            text = action.get("text", "")
            await page.keyboard.type(text, delay=40)
            return f"typed '{text[:50]}'"

        elif act == "key":
            key = action.get("key", "Enter")
            await page.keyboard.press(key)
            await asyncio.sleep(0.5)
            return f"pressed {key}"

        elif act == "scroll":
            direction = action.get("direction", "down")
            amount = int(action.get("amount", 300))
            delta = amount if direction == "down" else -amount
            await page.mouse.wheel(0, delta)
            await asyncio.sleep(0.3)
            return f"scrolled {direction} {amount}px"

        elif act == "wait":
            ms = int(action.get("ms", 1000))
            await asyncio.sleep(ms / 1000)
            return f"waited {ms}ms"

        elif act in ("complete", "blocked"):
            return f"terminal: {act}"

        else:
            log.warning(f"Unknown action: {act}")
            await asyncio.sleep(1)
            return f"unknown action '{act}' skipped"

    async def run(self, task: str) -> dict:
        log.info("=" * 70)
        log.info(f"TASK: {task[:200]}")
        log.info(f"Model: {VISION_MODEL} | headless: {self.headless} | "
                 f"max_steps: {self.max_steps} | stealth: {self.use_stealth}")
        log.info("=" * 70)

        start_ts = time.time()
        result = {"task": task, "status": "incomplete", "steps": 0,
                  "result": None, "blocked_reason": None, "elapsed_s": 0,
                  "history": []}

        async with async_playwright() as pw:
            if BROWSERLESS_URL and BROWSERLESS_TOKEN:
                log.info(f"Connecting to Browserless: {BROWSERLESS_URL}")
                browser: Browser = await pw.chromium.connect_over_cdp(
                    f"{BROWSERLESS_URL}?token={BROWSERLESS_TOKEN}"
                )
                context: BrowserContext = await browser.new_context(
                    viewport=self.viewport,
                    user_agent=(
                        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                        "AppleWebKit/537.36 (KHTML, like Gecko) "
                        "Chrome/131.0.0.0 Safari/537.36"
                    ),
                    locale="en-AU",
                    timezone_id="Australia/Sydney",
                )
            else:
                log.info("Launching local Chromium headless")
                browser: Browser = await pw.chromium.launch(
                    headless=self.headless,
                    args=[
                        "--no-sandbox",
                        "--disable-dev-shm-usage",
                        "--disable-blink-features=AutomationControlled",
                        "--disable-infobars",
                    ],
                )
                context: BrowserContext = await browser.new_context(
                    viewport=self.viewport,
                    user_agent=(
                        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                        "AppleWebKit/537.36 (KHTML, like Gecko) "
                        "Chrome/131.0.0.0 Safari/537.36"
                    ),
                    locale="en-AU",
                    timezone_id="Australia/Sydney",
                )

            if self.use_stealth:
                try:
                    await Stealth().apply_stealth_async(context)
                    log.info("Stealth mode applied")
                except Exception as e:
                    log.warning(f"Stealth apply failed: {e}")

            page: Page = await context.new_page()

            screenshot = await self._get_screenshot(page)
            await self._save_screenshot(screenshot, "initial")

            for self.step in range(1, self.max_steps + 1):
                current_url = page.url
                log.info(f"\n--- Step {self.step}/{self.max_steps} | URL: {current_url[:80]} ---")

                if not screenshot:
                    screenshot = await self._get_screenshot(page)

                # Think
                try:
                    action = self.gemini.decide_action(
                        task=task,
                        history=self.history,
                        screenshot_bytes=screenshot,
                        step=self.step,
                        max_steps=self.max_steps,
                        current_url=current_url,
                    )
                except Exception as e:
                    log.error(f"Gemini error: {e}")
                    action = {"action": "wait", "ms": 3000, "description": "gemini error"}

                log.info(f"  ACTION: {json.dumps(action)}")

                # Terminal actions
                if action.get("action") == "complete":
                    result_text = action.get("result", "Task completed")
                    log.info(f"  COMPLETE: {result_text}")
                    result.update({"status": "complete", "result": result_text,
                                   "steps": self.step})
                    await self._save_screenshot(screenshot, "complete")
                    break

                if action.get("action") == "blocked":
                    reason = action.get("reason", "Unknown block")
                    log.warning(f"  BLOCKED: {reason}")
                    result.update({"status": "blocked", "blocked_reason": reason,
                                   "steps": self.step})
                    await self._save_screenshot(screenshot, "blocked")
                    break

                # Act
                try:
                    status = await self._execute_action(page, action)
                except Exception as e:
                    status = f"execution error: {e}"
                    log.error(f"  Error: {e}")

                self.history.append(
                    f"Step {self.step}: [{action.get('action')}] "
                    f"{action.get('description','')} — {status}"
                )
                self.results.append({"step": self.step, "action": action, "status": status})

                # Observe
                await asyncio.sleep(1.5)
                screenshot = await self._get_screenshot(page)

                if self.step % 5 == 0:
                    await self._save_screenshot(screenshot, f"progress_s{self.step}")

            else:
                log.warning("Max steps reached without completion")
                result["steps"] = self.max_steps

            result["elapsed_s"] = round(time.time() - start_ts, 1)
            result["history"] = self.history

            await context.close()
            await browser.close()

        log.info(f"\nRESULT: {json.dumps({k: v for k, v in result.items() if k != 'history'}, indent=2)}")
        return result


# ── Pre-built Tasks ────────────────────────────────────────────────────────────

GODADDY_DNS_TASK = """
TASK: Add DNS A records to GoDaddy for sunaiva.com and sunaivadigital.com

Login credentials:
- URL: https://sso.godaddy.com/
- Username/Customer ID: 108866004
- Password: TripleWebsites333

EXACT STEPS:
1. Navigate to https://sso.godaddy.com/
2. Fill the username field (id="username") with: 108866004
3. Fill the password field (id="password") with: TripleWebsites333
4. Click the Sign In button (id="submitBtn")
5. Wait for redirect to account dashboard
6. Navigate to https://account.godaddy.com/products
7. Look for domains listed: sunaiva.com and/or sunaivadigital.com
8. For each domain, click on DNS or Manage DNS
9. Add A record: Name=@, Value=152.53.201.221, TTL=600
10. Add A record: Name=memory, Value=152.53.201.221, TTL=600
11. Add A record: Name=api, Value=152.53.201.221, TTL=600
12. Save all changes

If login fails or bot check appears that cannot be resolved, report as BLOCKED.
Report which domains were found and which DNS records were successfully added.
"""

ELESTIO_TASK = """
TASK: Find the Browserless WebSocket URL and token from Elestio

Steps:
1. Navigate to https://app.elestio.app/
2. You may see a login page. Try navigating to: https://app.elestio.app/dashboard
3. Look at the URL - if redirected to login, the admin console might need credentials.
4. Check if there's a direct service URL for browserless
5. Try navigating to: https://browserless-genesis-u50607.vm.elestio.app/
6. Or try: https://browserless-u50607.vm.elestio.app/
7. Look for a WebSocket URL in the format: wss://browserless-XXXXX.vm.elestio.app
8. If you find a browserless service, get its health page or config page
9. Report the exact WebSocket URL and any auth token found

Also check:
- https://openclaw-genesis-u50607.vm.elestio.app/ (existing service for reference)

Report: the Browserless WebSocket URL and token, or what was found at each URL.
"""


# ── GoDaddy Direct DNS API (when we have API key) ─────────────────────────────
class GoDaddyDirectDNS:
    """Use GoDaddy REST API when API_KEY+SECRET are available."""

    BASE = "https://api.godaddy.com"

    def __init__(self, api_key: str, api_secret: str):
        self.headers = {
            "Authorization": f"sso-key {api_key}:{api_secret}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        }

    async def list_domains(self) -> list:
        import aiohttp
        async with aiohttp.ClientSession(headers=self.headers) as s:
            async with s.get(f"{self.BASE}/v1/domains?limit=50") as r:
                data = await r.json()
                return data if isinstance(data, list) else []

    async def add_a_record(self, domain: str, name: str, ip: str, ttl: int = 600) -> dict:
        import aiohttp
        record = [{"data": ip, "name": name, "ttl": ttl, "type": "A"}]
        async with aiohttp.ClientSession(headers=self.headers) as s:
            async with s.patch(
                f"{self.BASE}/v1/domains/{domain}/records/A/{name}",
                json=record
            ) as r:
                return {"status": r.status, "domain": domain, "name": name, "ip": ip}


# ── CLI Entry Point ────────────────────────────────────────────────────────────
async def main():
    parser = argparse.ArgumentParser(description="Gemini Vision Browser Agent")
    parser.add_argument("--task", type=str, default=None)
    parser.add_argument("--headless", action="store_true", default=True)
    parser.add_argument("--no-headless", dest="headless", action="store_false")
    parser.add_argument("--max-steps", type=int, default=50)
    parser.add_argument("--godaddy", action="store_true", help="GoDaddy DNS setup")
    parser.add_argument("--elestio", action="store_true", help="Elestio Browserless discovery")
    parser.add_argument("--no-stealth", action="store_true", help="Disable stealth mode")
    args = parser.parse_args()

    use_stealth = not args.no_stealth
    agent = GeminiVisionAgent(
        headless=args.headless,
        max_steps=args.max_steps,
        use_stealth=use_stealth,
    )
    all_results = []

    if args.godaddy:
        log.info("\n>>> TASK: GoDaddy DNS Setup <<<\n")
        result = await agent.run(GODADDY_DNS_TASK)
        all_results.append({"task": "godaddy_dns", "result": result})
        agent.reset()

    if args.elestio:
        log.info("\n>>> TASK: Elestio Browserless Discovery <<<\n")
        result = await agent.run(ELESTIO_TASK)
        all_results.append({"task": "elestio_browserless", "result": result})

        # Auto-save browserless URL if found
        if result.get("status") == "complete" and result.get("result"):
            rt = result["result"]
            wss_m = re.search(r"wss://[^\s\"'<>]+", rt)
            tok_m = re.search(r"token[=:\s]+([a-zA-Z0-9_\-]{20,})", rt, re.I)
            secrets_path = Path("/mnt/e/genesis-system/config/secrets.env")
            additions = []
            if wss_m:
                wss = wss_m.group(0).rstrip(".,)")
                log.info(f"Found Browserless WSS URL: {wss}")
                content = secrets_path.read_text()
                if "BROWSERLESS_URL" not in content:
                    additions.append(f"BROWSERLESS_URL={wss}")
            if tok_m:
                tok = tok_m.group(1)
                log.info(f"Found Browserless token: {tok[:10]}...")
                content = secrets_path.read_text()
                if "BROWSERLESS_TOKEN" not in content:
                    additions.append(f"BROWSERLESS_TOKEN={tok}")
            if additions:
                with open(secrets_path, "a") as f:
                    f.write("\n# Browserless (auto-discovered)\n")
                    f.write("\n".join(additions) + "\n")
                log.info("Saved Browserless credentials to secrets.env")

    if args.task:
        log.info(f"\n>>> TASK: {args.task[:80]} <<<\n")
        result = await agent.run(args.task)
        all_results.append({"task": "custom", "result": result})

    if not any([args.godaddy, args.elestio, args.task]):
        parser.print_help()
        return

    # Summary
    print("\n" + "=" * 70)
    print("FINAL SUMMARY")
    print("=" * 70)
    for item in all_results:
        r = item["result"]
        print(f"\nTask:    {item['task']}")
        print(f"Status:  {r.get('status')}")
        print(f"Steps:   {r.get('steps')}")
        print(f"Elapsed: {r.get('elapsed_s')}s")
        if r.get("result"):
            print(f"Result:  {r.get('result')}")
        if r.get("blocked_reason"):
            print(f"Blocked: {r.get('blocked_reason')}")


if __name__ == "__main__":
    asyncio.run(main())