from __future__ import annotations
"""SalmAlm Browser Automation — Chrome DevTools Protocol (CDP) over WebSocket.

Pure stdlib. No Playwright/Selenium/Puppeteer needed.
Connects to Chrome's debug port via CDP WebSocket to:
  - Navigate pages
  - Take screenshots (base64 PNG)
  - Extract page content (DOM → text/markdown)
  - Execute JavaScript
  - Click elements, fill forms
  - Capture console logs
  - Generate PDFs

Requirements:
  Chrome/Chromium must be running with: --remote-debugging-port=9222
  Launch: google-chrome --remote-debugging-port=9222 --headless=new

Usage:
  from salmalm.browser import browser
  await browser.connect()
  await browser.navigate("https://example.com")
  text = await browser.get_text()
  screenshot = await browser.screenshot()  # base64 PNG
"""


import asyncio
import base64
import json
import struct
import hashlib
import time
import urllib.request
from typing import Any, Dict, List, Optional

from .crypto import log

# CDP WebSocket frame helpers (reuse ws.py logic but as client)
_WS_MAGIC = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"


class CDPConnection:
    """Low-level CDP WebSocket connection to Chrome."""

    def __init__(self):
        self._reader: Optional[asyncio.StreamReader] = None
        self._writer: Optional[asyncio.StreamWriter] = None
        self._connected = False
        self._msg_id = 0
        self._pending: Dict[int, asyncio.Future] = {}
        self._events: List[dict] = []
        self._event_handlers: Dict[str, list] = {}
        self._read_task: Optional[asyncio.Task] = None

    async def connect(self, ws_url: str) -> bool:
        """Connect to CDP WebSocket endpoint."""
        try:
            # Parse ws://host:port/path
            url = ws_url.replace("ws://", "")
            if "/" in url:
                hostport, path = url.split("/", 1)
                path = "/" + path
            else:
                hostport, path = url, "/"
            host, port_str = hostport.split(":")
            port = int(port_str)

            self._reader, self._writer = await asyncio.open_connection(host, port)

            # WebSocket handshake (as client)
            import os
            key = base64.b64encode(os.urandom(16)).decode()
            request = (
                f"GET {path} HTTP/1.1\r\n"
                f"Host: {hostport}\r\n"
                f"Upgrade: websocket\r\n"
                f"Connection: Upgrade\r\n"
                f"Sec-WebSocket-Key: {key}\r\n"
                f"Sec-WebSocket-Version: 13\r\n"
                f"\r\n"
            )
            self._writer.write(request.encode())
            await self._writer.drain()

            # Read response
            while True:
                line = await asyncio.wait_for(self._reader.readline(), timeout=5)
                if line == b'\r\n' or not line:
                    break

            self._connected = True
            self._read_task = asyncio.create_task(self._read_loop())
            return True

        except Exception as e:
            log.error(f"CDP connect failed: {e}")
            return False

    async def disconnect(self):
        self._connected = False
        if self._read_task:
            self._read_task.cancel()
        if self._writer:
            try:
                self._writer.close()
            except Exception:
                pass
        for f in self._pending.values():
            if not f.done():
                f.set_exception(ConnectionError("Disconnected"))
        self._pending.clear()

    async def send(self, method: str, params: dict = None, timeout: float = 30) -> dict:
        """Send CDP command and wait for response."""
        if not self._connected:
            raise ConnectionError("Not connected to Chrome")
        self._msg_id += 1
        msg_id = self._msg_id
        msg = {"id": msg_id, "method": method}
        if params:
            msg["params"] = params

        future = asyncio.get_event_loop().create_future()
        self._pending[msg_id] = future

        await self._send_frame(json.dumps(msg).encode('utf-8'))

        try:
            result = await asyncio.wait_for(future, timeout=timeout)
            return result
        except asyncio.TimeoutError:
            self._pending.pop(msg_id, None)
            raise TimeoutError(f"CDP command timeout: {method}")

    def on_event(self, method: str, handler):
        """Register event handler."""
        if method not in self._event_handlers:
            self._event_handlers[method] = []
        self._event_handlers[method].append(handler)

    async def _send_frame(self, data: bytes):
        """Send masked WebSocket frame (client must mask)."""
        import os
        mask_key = os.urandom(4)
        masked = bytes(b ^ mask_key[i % 4] for i, b in enumerate(data))

        length = len(data)
        header = bytes([0x81])  # FIN + TEXT
        if length < 126:
            header += bytes([0x80 | length])  # MASK bit set
        elif length < 65536:
            header += bytes([0x80 | 126]) + struct.pack('!H', length)
        else:
            header += bytes([0x80 | 127]) + struct.pack('!Q', length)

        self._writer.write(header + mask_key + masked)
        await self._writer.drain()

    async def _read_loop(self):
        """Background task reading CDP responses/events."""
        try:
            while self._connected:
                frame = await self._recv_frame()
                if frame is None:
                    break
                try:
                    msg = json.loads(frame)
                except json.JSONDecodeError:
                    continue

                if "id" in msg:
                    # Response to a command
                    fut = self._pending.pop(msg["id"], None)
                    if fut and not fut.done():
                        if "error" in msg:
                            fut.set_exception(
                                RuntimeError(f"CDP error: {msg['error'].get('message', '')}"))
                        else:
                            fut.set_result(msg.get("result", {}))
                elif "method" in msg:
                    # Event
                    method = msg["method"]
                    self._events.append(msg)
                    if len(self._events) > 200:
                        self._events = self._events[-100:]
                    for handler in self._event_handlers.get(method, []):
                        try:
                            handler(msg.get("params", {}))
                        except Exception:
                            pass
        except (asyncio.CancelledError, ConnectionError, OSError):
            pass
        self._connected = False

    async def _recv_frame(self) -> Optional[str]:
        """Read one WebSocket frame (server frames are unmasked)."""
        try:
            b0, b1 = await self._reader.readexactly(2)
            opcode = b0 & 0x0F
            length = b1 & 0x7F

            if length == 126:
                data = await self._reader.readexactly(2)
                length = struct.unpack('!H', data)[0]
            elif length == 127:
                data = await self._reader.readexactly(8)
                length = struct.unpack('!Q', data)[0]

            if length > 64 * 1024 * 1024:
                return None

            payload = await self._reader.readexactly(length) if length > 0 else b''

            if opcode == 0x1:  # Text
                return payload.decode('utf-8', errors='replace')
            elif opcode == 0x8:  # Close
                return None
            elif opcode == 0x9:  # Ping → Pong
                await self._send_frame(payload)  # Should be pong opcode but simplified
                return await self._recv_frame()
            return None
        except (asyncio.IncompleteReadError, ConnectionError, OSError):
            return None


class BrowserController:
    """High-level browser automation API over CDP."""

    def __init__(self, debug_host: str = "127.0.0.1", debug_port: int = 9222):
        self.debug_host = debug_host
        self.debug_port = debug_port
        self._cdp: Optional[CDPConnection] = None
        self._console_logs: List[str] = []

    @property
    def connected(self) -> bool:
        return self._cdp is not None and self._cdp._connected

    async def connect(self, tab_index: int = 0) -> bool:
        """Connect to Chrome's first tab via CDP."""
        try:
            # Get available targets
            url = f"http://{self.debug_host}:{self.debug_port}/json"
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=5) as resp:
                targets = json.loads(resp.read())

            # Find a page target
            pages = [t for t in targets if t.get("type") == "page"]
            if not pages:
                log.error("No browser tabs found")
                return False

            target = pages[min(tab_index, len(pages) - 1)]
            ws_url = target.get("webSocketDebuggerUrl", "")
            if not ws_url:
                log.error("No WebSocket URL for target")
                return False

            self._cdp = CDPConnection()
            ok = await self._cdp.connect(ws_url)
            if not ok:
                return False

            # Enable domains
            await self._cdp.send("Page.enable")
            await self._cdp.send("Runtime.enable")
            await self._cdp.send("DOM.enable")
            await self._cdp.send("Network.enable")

            # Capture console
            self._cdp.on_event("Runtime.consoleAPICalled", self._on_console)

            log.info(f"🌐 Browser connected: {target.get('title', '?')}")
            return True

        except Exception as e:
            log.error(f"Browser connect failed: {e}")
            return False

    async def disconnect(self):
        if self._cdp:
            await self._cdp.disconnect()
            self._cdp = None

    def _on_console(self, params: dict):
        """Capture console.log messages."""
        args = params.get("args", [])
        texts = [a.get("value", str(a.get("description", ""))) for a in args]
        line = f"[{params.get('type', 'log')}] {' '.join(str(t) for t in texts)}"
        self._console_logs.append(line)
        if len(self._console_logs) > 500:
            self._console_logs = self._console_logs[-250:]

    async def navigate(self, url: str, wait_load: bool = True) -> dict:
        """Navigate to URL. Returns {frameId, loaderId}."""
        if not self.connected:
            return {"error": "Not connected"}
        result = await self._cdp.send("Page.navigate", {"url": url})
        if wait_load:
            try:
                await asyncio.sleep(1)  # Simple wait; could listen for Page.loadEventFired
            except Exception:
                pass
        return result

    async def screenshot(self, full_page: bool = False,
                         format: str = "png", quality: int = 80) -> str:
        """Take screenshot, return base64 encoded image."""
        if not self.connected:
            return ""
        params = {"format": format}
        if format == "jpeg":
            params["quality"] = quality
        if full_page:
            # Get full page dimensions
            metrics = await self._cdp.send("Page.getLayoutMetrics")
            content = metrics.get("contentSize", {})
            params["clip"] = {
                "x": 0, "y": 0,
                "width": content.get("width", 1280),
                "height": content.get("height", 720),
                "scale": 1,
            }
        result = await self._cdp.send("Page.captureScreenshot", params)
        return result.get("data", "")

    async def get_text(self) -> str:
        """Extract page text content."""
        if not self.connected:
            return ""
        result = await self.evaluate(
            "document.body ? document.body.innerText : document.documentElement.textContent || ''")
        return result.get("value", "")

    async def get_html(self) -> str:
        """Get page HTML."""
        if not self.connected:
            return ""
        result = await self._cdp.send("DOM.getDocument", {"depth": -1})
        root = result.get("root", {})
        node_id = root.get("nodeId", 0)
        if node_id:
            html = await self._cdp.send("DOM.getOuterHTML", {"nodeId": node_id})
            return html.get("outerHTML", "")
        return ""

    async def evaluate(self, expression: str, return_by_value: bool = True) -> dict:
        """Execute JavaScript and return result."""
        if not self.connected:
            return {"error": "Not connected"}
        result = await self._cdp.send("Runtime.evaluate", {
            "expression": expression,
            "returnByValue": return_by_value,
            "awaitPromise": True,
        })
        remote = result.get("result", {})
        if remote.get("type") == "undefined":
            return {"value": None}
        return {"value": remote.get("value", remote.get("description", str(remote)))}

    async def click(self, selector: str) -> bool:
        """Click an element by CSS selector."""
        if not self.connected:
            return False
        # Use JS to find and click
        result = await self.evaluate(f"""
            (function() {{
                var el = document.querySelector({json.dumps(selector)});
                if (!el) return 'NOT_FOUND';
                el.click();
                return 'OK';
            }})()
        """)
        return result.get("value") == "OK"

    async def type_text(self, selector: str, text: str) -> bool:
        """Type text into an input element."""
        if not self.connected:
            return False
        result = await self.evaluate(f"""
            (function() {{
                var el = document.querySelector({json.dumps(selector)});
                if (!el) return 'NOT_FOUND';
                el.focus();
                el.value = {json.dumps(text)};
                el.dispatchEvent(new Event('input', {{bubbles: true}}));
                el.dispatchEvent(new Event('change', {{bubbles: true}}));
                return 'OK';
            }})()
        """)
        return result.get("value") == "OK"

    async def get_tabs(self) -> List[dict]:
        """List all browser tabs."""
        try:
            url = f"http://{self.debug_host}:{self.debug_port}/json"
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=5) as resp:
                targets = json.loads(resp.read())
            return [{"id": t["id"], "title": t.get("title", ""),
                     "url": t.get("url", ""), "type": t.get("type", "")}
                    for t in targets if t.get("type") == "page"]
        except Exception as e:
            return [{"error": str(e)}]

    async def new_tab(self, url: str = "about:blank") -> dict:
        """Open a new tab."""
        try:
            api_url = f"http://{self.debug_host}:{self.debug_port}/json/new?{url}"
            req = urllib.request.Request(api_url, method="PUT")
            with urllib.request.urlopen(req, timeout=5) as resp:
                return json.loads(resp.read())
        except Exception as e:
            return {"error": str(e)}

    async def pdf(self) -> str:
        """Generate PDF of current page, return base64."""
        if not self.connected:
            return ""
        result = await self._cdp.send("Page.printToPDF", {
            "printBackground": True,
            "preferCSSPageSize": True,
        })
        return result.get("data", "")

    def get_console_logs(self, limit: int = 50) -> List[str]:
        """Get captured console logs."""
        return self._console_logs[-limit:]

    def get_status(self) -> dict:
        return {
            "connected": self.connected,
            "host": f"{self.debug_host}:{self.debug_port}",
            "console_logs": len(self._console_logs),
        }


# Module-level instance
browser = BrowserController()
