import asyncio
import json
import logging
from typing import Any

import httpx
import websockets

from gsearch_mcp_cli.constants import (
    DEFAULT_MAX_TABS,
    TIMEOUT_JS_EVAL,
    TIMEOUT_PAGE_LOAD,
    TIMEOUT_TAB_CLEANUP,
)
from gsearch_mcp_cli.exceptions import CDPConnectionError, GSearchError, GSearchTimeoutError

logger = logging.getLogger(__name__)


class CDPConnection:
    """Async CDP connection with message dispatcher."""

    def __init__(self):
        self._ws = None
        self._msg_id = 0
        self._pending: dict[int, asyncio.Future] = {}
        self._reader_task: asyncio.Task | None = None

    async def connect(self, ws_url: str) -> None:
        self._ws = await websockets.connect(ws_url, max_size=None, close_timeout=5)
        self._reader_task = asyncio.create_task(self._read_loop())

    async def _read_loop(self):
        try:
            async for raw in self._ws:
                msg = json.loads(raw)
                msg_id = msg.get("id")
                if msg_id is not None and msg_id in self._pending:
                    self._pending[msg_id].set_result(msg)
        except websockets.exceptions.ConnectionClosed:
            pass
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.debug(f"CDP read loop error: {e}")

    async def send(self, method: str, params: dict | None = None, timeout: float = 30) -> dict:
        if self._ws is None:
            raise CDPConnectionError("Not connected")
        self._msg_id += 1
        msg_id = self._msg_id
        future = asyncio.get_event_loop().create_future()
        self._pending[msg_id] = future
        payload = {"id": msg_id, "method": method}
        if params:
            payload["params"] = params
        await self._ws.send(json.dumps(payload))
        try:
            return await asyncio.wait_for(future, timeout=timeout)
        except asyncio.TimeoutError:
            raise GSearchTimeoutError(f"CDP command {method} timed out after {timeout}s")
        finally:
            self._pending.pop(msg_id, None)

    async def close(self) -> None:
        if self._reader_task:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()


class BrowserCDP:
    """Manages browser-level CDP and tab lifecycle."""

    def __init__(self, max_tabs: int = DEFAULT_MAX_TABS):
        self._conn = CDPConnection()
        self._active_tabs: set[str] = set()
        self._max_tabs = max_tabs
        self._port: int = 0

    async def connect(self, port: int) -> None:
        self._port = port
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"http://localhost:{port}/json/version", timeout=5)
            data = resp.json()
            ws_url = data["webSocketDebuggerUrl"]
        await self._conn.connect(ws_url)

    async def is_healthy(self) -> bool:
        try:
            async with httpx.AsyncClient() as client:
                resp = await client.get(f"http://localhost:{self._port}/json/version", timeout=2)
                return resp.status_code == 200
        except Exception:
            return False

    def create_tab(self, url: str = "about:blank") -> "Tab":
        if len(self._active_tabs) >= self._max_tabs:
            raise GSearchError(f"Max concurrent tabs ({self._max_tabs}) reached")
        return Tab(self, url)

    async def _create_target(self, url: str) -> str:
        resp = await self._conn.send("Target.createTarget", {"url": url})
        target_id = resp["result"]["targetId"]
        self._active_tabs.add(target_id)
        return target_id

    async def _close_target(self, target_id: str) -> None:
        try:
            await self._conn.send("Target.closeTarget", {"targetId": target_id}, timeout=TIMEOUT_TAB_CLEANUP)
        except Exception:
            pass
        self._active_tabs.discard(target_id)

    async def _get_tab_ws_url(self, target_id: str) -> str:
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"http://localhost:{self._port}/json", timeout=5)
            pages = resp.json()
        for page in pages:
            if page.get("id") == target_id:
                return page["webSocketDebuggerUrl"]
        raise CDPConnectionError(f"Tab {target_id} not found in /json list")

    async def close(self) -> None:
        for target_id in list(self._active_tabs):
            await self._close_target(target_id)
        await self._conn.close()


class Tab:
    """Page-level CDP with context manager for guaranteed cleanup."""

    def __init__(self, browser: BrowserCDP, url: str = "about:blank"):
        self._browser = browser
        self._url = url
        self._conn: CDPConnection | None = None
        self._target_id: str | None = None

    async def __aenter__(self) -> "Tab":
        self._target_id = await self._browser._create_target(self._url)
        ws_url = await self._browser._get_tab_ws_url(self._target_id)
        self._conn = CDPConnection()
        await self._conn.connect(ws_url)
        await self._conn.send("Page.enable")
        await self._conn.send("Runtime.enable")
        await self._conn.send("Network.enable")
        await self.override_ua()
        return self

    async def __aexit__(self, *exc) -> None:
        try:
            await asyncio.wait_for(self._cleanup(), timeout=TIMEOUT_TAB_CLEANUP)
        except (Exception, asyncio.CancelledError):
            pass

    async def _cleanup(self) -> None:
        if self._conn:
            await self._conn.close()
        if self._target_id:
            await self._browser._close_target(self._target_id)

    async def navigate(self, url: str, timeout: float = TIMEOUT_PAGE_LOAD) -> None:
        await self._conn.send("Page.navigate", {"url": url}, timeout=timeout)
        await self.wait_for_load(timeout=timeout)

    async def evaluate(self, expression: str, timeout: float = TIMEOUT_JS_EVAL) -> Any:
        resp = await self._conn.send(
            "Runtime.evaluate",
            {"expression": expression, "returnByValue": True},
            timeout=timeout,
        )
        result = resp.get("result", {}).get("result", {})
        if result.get("type") == "undefined":
            return None
        return result.get("value")

    async def wait_for_load(self, timeout: float = TIMEOUT_PAGE_LOAD) -> None:
        deadline = asyncio.get_event_loop().time() + timeout
        while asyncio.get_event_loop().time() < deadline:
            state = await self.evaluate("document.readyState")
            if state in ("complete", "interactive"):
                return
            await asyncio.sleep(0.3)
        raise GSearchTimeoutError(f"Page did not load within {timeout}s")

    async def override_ua(self) -> None:
        ua = await self.evaluate("navigator.userAgent")
        if ua and "HeadlessChrome" in ua:
            new_ua = ua.replace("HeadlessChrome", "Chrome")
            await self._conn.send("Network.setUserAgentOverride", {"userAgent": new_ua})

    async def get_cookies(self) -> list[dict]:
        resp = await self._conn.send("Network.getAllCookies")
        return resp.get("result", {}).get("cookies", [])
