"""
Super simple client interface for BrowserFabric.

Usage:
    import browserfabric

    # Simple usage
    mcp = browserfabric.MCP("browseruse")
    session = await mcp.create_session()
    await session.navigate("https://example.com")
    await session.screenshot("page.png")

    # With persistent context
    session = await mcp.create_session(persist=True)
    await session.navigate("https://twitter.com")
    # ... user logs in ...
    ctx = await session.save_context("twitter-login")
    await session.close()

    # Restore context later
    session = await mcp.create_session(context_id=ctx["context_id"])
    # Already logged in!

    # Or even simpler
    async with browserfabric.browser() as browser:
        await browser.navigate("https://example.com")
        await browser.click("#button")
        await browser.screenshot("result.png")
"""

import asyncio
import base64
import os
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any, Dict, List, Optional

import httpx


class MCPError(Exception):
    """Base exception for MCP operations."""

    pass


class BrowserSession:
    """Represents a browser session with convenient methods."""

    def __init__(self, client: "MCPClient", session_id: str, ws_url: Optional[str] = None,
                 cdp_port: Optional[int] = None):
        self.client = client
        self.session_id = session_id
        self.ws_url = ws_url
        self.cdp_port = cdp_port
        self._closed = False

    async def navigate(self, url: str, wait_until: str = "domcontentloaded") -> Dict[str, Any]:
        """Navigate to a URL. wait_until can be: domcontentloaded, load, networkidle, commit."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "navigate", {"url": url, "wait_until": wait_until, "session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Navigation failed: {result.get('error')}")

        return result["result"]

    async def click(self, selector: str, by: str = "css") -> Dict[str, Any]:
        """Click an element."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "click_element",
            {"selector": selector, "by": by, "session_id": self.session_id},
        )

        if not result.get("success"):
            raise MCPError(f"Click failed: {result.get('error')}")

        return result["result"]

    async def type(self, selector: str, text: str, by: str = "css") -> Dict[str, Any]:
        """Type text into an element."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "type_text",
            {
                "selector": selector,
                "text": text,
                "by": by,
                "session_id": self.session_id,
            },
        )

        if not result.get("success"):
            raise MCPError(f"Type failed: {result.get('error')}")

        return result["result"]

    async def find(self, selector: str, by: str = "css") -> List[Dict[str, Any]]:
        """Find elements on the page."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "find_elements",
            {"selector": selector, "by": by, "session_id": self.session_id},
        )

        if not result.get("success"):
            raise MCPError(f"Find failed: {result.get('error')}")

        return result["result"].get("elements", [])

    async def screenshot(
        self, filename: Optional[str] = None, save_dir: str = "screenshots", full_page: bool = False
    ) -> str:
        """Take a screenshot and save it. Set full_page=True for full scrollable page."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "take_screenshot", {"full_page": full_page, "session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Screenshot failed: {result.get('error')}")

        screenshot_b64 = result["result"]["screenshot"]
        screenshot_data = base64.b64decode(screenshot_b64)

        if not filename:
            from datetime import datetime

            filename = f"screenshot_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"

        save_path = Path(save_dir)
        save_path.mkdir(exist_ok=True)

        filepath = save_path / filename
        with open(filepath, "wb") as f:
            f.write(screenshot_data)

        return str(filepath)

    async def page_info(self) -> Dict[str, Any]:
        """Get current page information."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "get_page_info", {"session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Get page info failed: {result.get('error')}")

        return result["result"]

    async def observe(self) -> Dict[str, Any]:
        """Get simplified text-based DOM tree of important visible elements with interaction paths."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "observe", {"session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Observe failed: {result.get('error')}")

        return result["result"]

    async def snapshot(self) -> str:
        """Get the accessibility tree of the page."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "snapshot", {"session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Snapshot failed: {result.get('error')}")

        return result["result"].get("snapshot", "")

    async def session_info(self) -> Dict[str, Any]:
        """Get session metadata including CDP WebSocket URL."""
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "get_session_info", {"session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Get session info failed: {result.get('error')}")

        return result["result"]

    async def wait_for(
        self, selector: str, state: str = "visible", timeout: int = 30, by: str = "css"
    ) -> Dict[str, Any]:
        """Wait for an element to reach a given state (visible, hidden, attached, detached)."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "wait_for",
            {"selector": selector, "state": state, "timeout": timeout, "by": by, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Wait failed: {result.get('error')}")
        return result["result"]

    async def scroll(self, direction: str = "down", amount: int = 500) -> Dict[str, Any]:
        """Scroll the page. Directions: down, up, top, bottom."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "scroll",
            {"direction": direction, "amount": amount, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Scroll failed: {result.get('error')}")
        return result["result"]

    async def press_key(self, key: str) -> Dict[str, Any]:
        """Press a keyboard key (Enter, Tab, Escape, ArrowDown, etc.)."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "press_key", {"key": key, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Press key failed: {result.get('error')}")
        return result["result"]

    async def select_option(self, selector: str, value: str, by: str = "css") -> Dict[str, Any]:
        """Select an option from a dropdown/select element."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "select_option",
            {"selector": selector, "value": value, "by": by, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Select failed: {result.get('error')}")
        return result["result"]

    async def hover(self, selector: str, by: str = "css") -> Dict[str, Any]:
        """Hover over an element."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "hover",
            {"selector": selector, "by": by, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Hover failed: {result.get('error')}")
        return result["result"]

    async def evaluate_js(self, expression: str) -> Any:
        """Execute JavaScript in the page and return the result."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "evaluate_js",
            {"expression": expression, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"JS eval failed: {result.get('error')}")
        return result["result"].get("result")

    async def go_back(self) -> Dict[str, Any]:
        """Navigate back in browser history."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "go_back", {"session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Go back failed: {result.get('error')}")
        return result["result"]

    async def go_forward(self) -> Dict[str, Any]:
        """Navigate forward in browser history."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "go_forward", {"session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Go forward failed: {result.get('error')}")
        return result["result"]

    async def screenshot_element(self, selector: str, by: str = "css") -> str:
        """Take a screenshot of a specific element. Returns base64 PNG."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "screenshot_element",
            {"selector": selector, "by": by, "session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Element screenshot failed: {result.get('error')}")
        return result["result"]["screenshot"]

    async def page_summary(self) -> Dict[str, Any]:
        """Get a quick summary of the current page state."""
        if self._closed:
            raise MCPError("Session is closed")
        result = await self.client._call_tool(
            "page_summary", {"session_id": self.session_id},
        )
        if not result.get("success"):
            raise MCPError(f"Page summary failed: {result.get('error')}")
        return result["result"]

    async def save_context(self, name: str) -> Dict[str, Any]:
        """Save browser context (cookies, localStorage) for later restoration.

        Args:
            name: Human-readable name for the context (e.g., "twitter-login")

        Returns:
            Dict with context_id, name, cookies_count, status
        """
        if self._closed:
            raise MCPError("Session is closed")

        result = await self.client._call_tool(
            "save_context", {"context_name": name, "session_id": self.session_id}
        )

        if not result.get("success"):
            raise MCPError(f"Save context failed: {result.get('error')}")

        return result["result"]

    async def close(self):
        """Close the browser session."""
        if self._closed:
            return

        try:
            result = await self.client._call_tool(
                "close_session", {"session_id": self.session_id}
            )

            if not result.get("success"):
                print(f"Warning: Failed to close session: {result.get('error')}")
        finally:
            self._closed = True

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()


class MCPClient:
    """Simple client for BrowserFabric services."""

    def __init__(
        self,
        service_name: str,
        api_key: Optional[str] = None,
        base_url: str = "http://localhost:9000",
        max_retries: int = 3,
        timeout: float = 60.0,
    ):
        self.service_name = service_name
        self.base_url = base_url.rstrip("/")
        self.max_retries = max_retries
        self.timeout = timeout

        # API key support
        self.headers = {}
        key = api_key or os.getenv("BROWSERFABRIC_API_KEY")
        if key:
            self.headers["Authorization"] = f"Bearer {key}"

    async def _call_tool(
        self, tool_name: str, arguments: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Call a tool on the service with automatic retry on transient errors."""
        session_id = (
            arguments.pop("session_id", None) if "session_id" in arguments else None
        )

        request_data = {"tool_name": tool_name, "arguments": arguments}
        if session_id:
            request_data["session_id"] = session_id

        last_error = None
        for attempt in range(self.max_retries):
            try:
                async with httpx.AsyncClient(timeout=self.timeout) as client:
                    response = await client.post(
                        f"{self.base_url}/api/v1/services/{self.service_name}/call",
                        headers=self.headers,
                        json=request_data,
                    )

                    # Don't retry client errors (4xx) except 429 (rate limit)
                    if response.status_code == 429:
                        retry_after = int(response.headers.get("Retry-After", 5))
                        await asyncio.sleep(retry_after)
                        continue

                    if response.status_code != 200:
                        raise MCPError(f"HTTP {response.status_code}: {response.text}")

                    return response.json()

            except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
                last_error = e
                if attempt < self.max_retries - 1:
                    wait = min(2 ** attempt, 10)  # 1s, 2s, 4s... max 10s
                    await asyncio.sleep(wait)
                    continue
                raise MCPError(f"Connection failed after {self.max_retries} attempts: {e}")

        raise MCPError(f"Failed after {self.max_retries} attempts: {last_error}")

    async def create_session(
        self,
        headless: bool = True,
        timeout: int = 30,
        context_id: Optional[str] = None,
        persist: bool = False,
    ) -> BrowserSession:
        """Create a new browser session.

        Args:
            headless: Run browser in headless mode
            timeout: Default timeout in seconds
            context_id: Optional context ID to restore cookies/localStorage from
            persist: If True, auto-save cookies/localStorage when session closes

        Returns:
            BrowserSession with ws_url for CDP connections
        """
        args = {"headless": headless, "timeout": timeout, "persist": persist}
        if context_id:
            args["context_id"] = context_id

        result = await self._call_tool("create_session", args)

        if not result.get("success"):
            raise MCPError(f"Failed to create session: {result.get('error')}")

        data = result["result"]
        return BrowserSession(
            self,
            data["session_id"],
            ws_url=data.get("ws_url"),
            cdp_port=data.get("cdp_port"),
        )

    async def list_sessions(self) -> List[Dict[str, Any]]:
        """List all active browser sessions for the current user."""
        result = await self._call_tool("list_sessions", {})
        if not result.get("success"):
            raise MCPError(f"Failed to list sessions: {result.get('error')}")
        return result["result"].get("sessions", [])

    async def list_contexts(self) -> List[Dict[str, Any]]:
        """List all saved persistent contexts."""
        result = await self._call_tool("list_contexts", {})
        if not result.get("success"):
            raise MCPError(f"Failed to list contexts: {result.get('error')}")
        return result["result"].get("contexts", [])

    async def delete_context(self, context_id: str) -> Dict[str, Any]:
        """Delete a saved persistent context."""
        result = await self._call_tool("delete_context", {"context_id": context_id})
        if not result.get("success"):
            raise MCPError(f"Failed to delete context: {result.get('error')}")
        return result["result"]

    async def health_check(self) -> Dict[str, Any]:
        """Check if the service is healthy."""
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(f"{self.base_url}/health")

            if response.status_code != 200:
                raise MCPError(f"Health check failed: HTTP {response.status_code}")

            return response.json()

    async def is_healthy(self) -> bool:
        """Quick health check. Returns True if server is reachable, False otherwise."""
        try:
            await self.health_check()
            return True
        except Exception:
            return False

    async def wait_for_ready(self, timeout: int = 30, interval: float = 1.0) -> bool:
        """Wait for the server to become ready. Returns True if ready within timeout."""
        import time
        start = time.monotonic()
        while time.monotonic() - start < timeout:
            if await self.is_healthy():
                return True
            await asyncio.sleep(interval)
        return False

    async def list_tools(self) -> List[Dict[str, Any]]:
        """List available tools for this service."""
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(
                f"{self.base_url}/api/v1/services/{self.service_name}/tools",
                headers=self.headers,
            )

            if response.status_code != 200:
                raise MCPError(f"Failed to list tools: HTTP {response.status_code}")

            result = response.json()
            return result.get("tools", [])


class MCP:
    """Main MCP interface - super simple to use!"""

    def __init__(
        self,
        service_name: str = "browseruse",
        api_key: Optional[str] = None,
        base_url: str = "http://localhost:9000",
    ):
        self.client = MCPClient(service_name, api_key, base_url)

    async def create_session(
        self,
        headless: bool = True,
        timeout: int = 30,
        context_id: Optional[str] = None,
        persist: bool = False,
    ) -> BrowserSession:
        """Create a new browser session."""
        return await self.client.create_session(headless, timeout, context_id, persist)

    async def list_contexts(self) -> List[Dict[str, Any]]:
        """List all saved persistent contexts."""
        return await self.client.list_contexts()

    async def delete_context(self, context_id: str) -> Dict[str, Any]:
        """Delete a saved persistent context."""
        return await self.client.delete_context(context_id)

    async def health_check(self) -> Dict[str, Any]:
        """Check service health."""
        return await self.client.health_check()

    async def list_tools(self) -> List[Dict[str, Any]]:
        """List available tools."""
        return await self.client.list_tools()

    async def quick_screenshot(self, url: str, filename: Optional[str] = None) -> str:
        """Quick screenshot of any URL."""
        async with await self.create_session() as session:
            await session.navigate(url)
            await asyncio.sleep(2)
            return await session.screenshot(filename)

    async def quick_navigate(self, url: str) -> BrowserSession:
        """Quick navigation - returns session for further use."""
        session = await self.create_session()
        await session.navigate(url)
        return session


@asynccontextmanager
async def browser(
    api_key: Optional[str] = None,
    base_url: str = "http://localhost:9000",
    headless: bool = True,
    context_id: Optional[str] = None,
    persist: bool = False,
):
    """
    Super convenient context manager for browser automation.

    Usage:
        async with browserfabric.browser() as browser:
            await browser.navigate("https://example.com")
            await browser.click("#button")
            await browser.screenshot("result.png")

        # With persistent context
        async with browserfabric.browser(persist=True) as browser:
            await browser.navigate("https://twitter.com")
            # ... cookies auto-saved on close
    """
    mcp = MCP("browseruse", api_key, base_url)
    session = await mcp.create_session(
        headless=headless, context_id=context_id, persist=persist
    )
    try:
        yield session
    finally:
        await session.close()


async def screenshot(
    url: str, filename: Optional[str] = None, api_key: Optional[str] = None
) -> str:
    """Take a quick screenshot of any URL."""
    mcp = MCP("browseruse", api_key)
    return await mcp.quick_screenshot(url, filename)


async def test_form(
    url: str, form_data: Dict[str, str], api_key: Optional[str] = None
) -> str:
    """Quick form testing."""
    async with browser(api_key) as session:
        await session.navigate(url)

        for selector, value in form_data.items():
            await session.type(selector, value)

        return await session.screenshot("form_test.png")


async def ensure_server_running():
    """Ensure BrowserFabric server is running."""
    try:
        mcp = MCP("browseruse")
        await mcp.health_check()
        return True
    except Exception:
        print("BrowserFabric server not running. Start it with: browserfabric serve")
        return False
