#!/usr/bin/env python3
"""
Nami Browser Agent — WebSocket client for Nami Bridge
Controls user's real Chrome browser via Nami Bridge WebSocket relay.

Usage:
    python3 /opt/nami-army/browser_agent.py --screenshot
    python3 /opt/nami-army/browser_agent.py --navigate "https://google.com"
    python3 /opt/nami-army/browser_agent.py --eval "document.title"
    python3 /opt/nami-army/browser_agent.py --click "#search-button"
    python3 /opt/nami-army/browser_agent.py --type "#input" "search text"

Integration:
    from browser_agent import BrowserAgent
    agent = BrowserAgent()
    await agent.connect()
    await agent.navigate("https://example.com")
    screenshot = await agent.screenshot()
    await agent.close()
"""

import asyncio
import json
import sys
import time
import base64
import os
import argparse
from pathlib import Path
from datetime import datetime
from typing import Optional, Any

try:
    import websockets
except ImportError:
    print("❌ websockets not installed. Run: pip install websockets", file=sys.stderr)
    sys.exit(1)


# ─── Default Configuration ─────────────────────────────────────────────
DEFAULT_BRIDGE_URL = "ws://127.0.0.1:8888"
DEFAULT_TIMEOUT = 30  # seconds
SCREENSHOT_DIR = Path("/tmp/nami_screenshots")


class BrowserAgent:
    """High-level browser control via Nami Bridge WebSocket relay."""

    def __init__(self, bridge_url: str = DEFAULT_BRIDGE_URL, timeout: int = DEFAULT_TIMEOUT):
        self.bridge_url = bridge_url
        self.timeout = timeout
        self.ws: Optional[websockets.WebSocketClientProtocol] = None
        self.connected = False
        self.browser_connected = False

    async def connect(self) -> bool:
        """Connect to Nami Bridge and register as agent."""
        try:
            self.ws = await asyncio.wait_for(
                websockets.connect(self.bridge_url),
                timeout=5
            )
            # Send hello
            await self.ws.send(json.dumps({
                "type": "hello",
                "role": "agent",
                "name": "NamiBrowserAgent",
                "version": "1.0.0"
            }))

            # Wait for ack
            response = await asyncio.wait_for(self.ws.recv(), timeout=5)
            ack = json.loads(response)
            if ack.get("type") == "ack" and ack.get("status") == "connected":
                self.connected = True
                # Check if a browser is actually connected via ping
                self.browser_connected = await self.ping()
                return True
            return False
        except Exception as e:
            print(f"❌ Connection failed: {e}", file=sys.stderr)
            return False

    async def _send_and_wait(self, msg: dict, expected_type: str = None, timeout: int = None) -> Optional[dict]:
        """Send a message and wait for a matching response."""
        if not self.ws:
            raise ConnectionError("Not connected to Nami Bridge")

        timeout = timeout or self.timeout
        await self.ws.send(json.dumps(msg))

        try:
            response = await asyncio.wait_for(self.ws.recv(), timeout=timeout)
            result = json.loads(response)

            if result.get("type") == "error":
                raise RuntimeError(f"Browser error: {result.get('error', 'unknown')}")

            if expected_type and result.get("type") != expected_type:
                # Try one more recv for the expected type
                response = await asyncio.wait_for(self.ws.recv(), timeout=timeout)
                result = json.loads(response)

            return result
        except asyncio.TimeoutError:
            raise TimeoutError(f"Request timed out after {timeout}s")
        except Exception as e:
            raise RuntimeError(f"Request failed: {e}")

    # ─── Core Commands ───────────────────────────────────────────────

    async def ping(self) -> bool:
        """Check if browser is alive."""
        try:
            await self.ws.send(json.dumps({"type": "ping"}))
            response = await asyncio.wait_for(self.ws.recv(), timeout=5)
            return json.loads(response).get("type") == "pong"
        except:
            return False

    async def navigate(self, url: str) -> dict:
        """Navigate the active tab to a URL."""
        return await self._send_and_wait(
            {"type": "navigate", "url": url},
            expected_type="result"
        )

    async def screenshot(self, save: bool = True) -> Optional[Path]:
        """Capture a screenshot of the current tab. Returns path if saved."""
        result = await self._send_and_wait(
            {"type": "screenshot"},
            expected_type="screenshot_data"
        )

        if result and result.get("data"):
            # Decode base64 data URL
            data = result["data"]
            if data.startswith("data:image/png;base64,"):
                data = data.replace("data:image/png;base64,", "")

            image_bytes = base64.b64decode(data)

            if save:
                SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = SCREENSHOT_DIR / f"screenshot_{timestamp}.png"
                filename.write_bytes(image_bytes)
                print(f"📸 Screenshot saved: {filename}")
                return filename

            return image_bytes

        return None

    async def click(self, selector: str) -> dict:
        """Click an element by CSS selector."""
        return await self._send_and_wait(
            {"type": "click", "selector": selector},
            expected_type="result"
        )

    async def type_text(self, selector: str, text: str) -> dict:
        """Type text into an input element."""
        return await self._send_and_wait(
            {"type": "type", "selector": selector, "text": text},
            expected_type="result"
        )

    async def evaluate(self, code: str) -> Any:
        """Execute JavaScript in the active tab and return the result."""
        result = await self._send_and_wait(
            {"type": "eval", "code": code},
            expected_type="eval_result"
        )
        return result.get("result") if result else None

    async def get_html(self) -> Optional[str]:
        """Get the full HTML of the current page."""
        result = await self._send_and_wait(
            {"type": "get_html"},
            expected_type="html_result"
        )
        return result.get("html") if result else None

    async def scroll(self, amount: int = 500) -> dict:
        """Scroll the page by pixels."""
        return await self._send_and_wait(
            {"type": "scroll", "amount": amount},
            expected_type="result"
        )

    async def get_page_info(self) -> Optional[dict]:
        """Get current page URL and title."""
        return await self._send_and_wait(
            {"type": "get_page_info"},
            expected_type="page_info"
        )

    async def close(self):
        """Close the WebSocket connection."""
        if self.ws:
            await self.ws.close()
            self.connected = False


# ─── CLI Interface ───────────────────────────────────────────────────

async def cli_main():
    parser = argparse.ArgumentParser(
        description="Nami Browser Agent — control browser via Nami Bridge",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 browser_agent.py --status          # Check connection
  python3 browser_agent.py --screenshot      # Capture screenshot
  python3 browser_agent.py --navigate "https://google.com"
  python3 browser_agent.py --eval "document.title"
  python3 browser_agent.py --click "#button"
  python3 browser_agent.py --type "#input" "hello"
  python3 browser_agent.py --html > page.html
  python3 browser_agent.py --url "http://vps:8888" --screenshot
        """
    )

    parser.add_argument("--url", default=DEFAULT_BRIDGE_URL,
                        help=f"Nami Bridge WebSocket URL (default: {DEFAULT_BRIDGE_URL})")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
                        help=f"Timeout in seconds (default: {DEFAULT_TIMEOUT})")

    # Actions
    parser.add_argument("--status", action="store_true", help="Check connection status")
    parser.add_argument("--screenshot", action="store_true", help="Capture and save screenshot")
    parser.add_argument("--navigate", type=str, metavar="URL", help="Navigate to URL")
    parser.add_argument("--eval", type=str, metavar="CODE", dest="js_code", help="Execute JavaScript")
    parser.add_argument("--click", type=str, metavar="SELECTOR", help="Click element by selector")
    parser.add_argument("--type", nargs=2, metavar=("SELECTOR", "TEXT"), dest="type_args",
                        help="Type text into input")
    parser.add_argument("--scroll", type=int, metavar="PX", help="Scroll page by N pixels")
    parser.add_argument("--html", action="store_true", help="Get page HTML")
    parser.add_argument("--info", action="store_true", help="Get page info (URL, title)")

    args = parser.parse_args()

    # Create agent
    agent = BrowserAgent(bridge_url=args.url, timeout=args.timeout)

    # Connect
    if not await agent.connect():
        print("❌ Cannot connect to Nami Bridge. Is it running?", file=sys.stderr)
        sys.exit(1)

    try:
        # Status check
        if args.status or (not any([
            args.screenshot, args.navigate, args.js_code, args.click,
            args.type_args, args.scroll, args.html, args.info
        ])):
            alive = await agent.ping()
            print(f"📡 Nami Bridge: {args.url}")
            print(f"🔌 Agent connected: {agent.connected}")
            print(f"🌐 Browser connected: {'✅ Yes' if alive else '❌ No — install Chrome Extension!'}")
            if not any([args.screenshot, args.navigate, args.js_code, args.click,
                        args.type_args, args.scroll, args.html, args.info]):
                return

        # Execute actions
        if args.navigate:
            print(f"🧭 Navigating to: {args.navigate}")
            result = await agent.navigate(args.navigate)
            print(f"✅ Navigated: {result.get('url', 'unknown')}")
            await asyncio.sleep(2)  # Wait for page load

        if args.js_code:
            print(f"⚡ Evaluating: {args.js_code}")
            result = await agent.evaluate(args.js_code)
            print(f"Result: {result}")

        if args.click:
            print(f"🖱️  Clicking: {args.click}")
            result = await agent.click(args.click)
            print(f"✅ Clicked")

        if args.type_args:
            selector, text = args.type_args
            print(f"⌨️  Typing '{text}' into {selector}")
            result = await agent.type_text(selector, text)
            print(f"✅ Typed")

        if args.scroll:
            print(f"📜 Scrolling {args.scroll}px")
            result = await agent.scroll(args.scroll)
            print(f"✅ Scrolled")

        if args.html:
            print("📄 Getting HTML...")
            html = await agent.get_html()
            print(html[:10000] if html else "No HTML")
            if html and len(html) > 10000:
                print(f"\n... (truncated, total {len(html)} chars)")

        if args.info:
            print("ℹ️  Getting page info...")
            info = await agent.get_page_info()
            if info:
                print(json.dumps(info, indent=2))

        if args.screenshot:
            print("📸 Capturing screenshot...")
            path = await agent.screenshot(save=True)
            if path:
                print(f"✅ Screenshot saved: {path}")
            else:
                print("❌ Screenshot failed")

    except Exception as e:
        print(f"❌ Error: {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        await agent.close()


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