import asyncio
from typing import Annotated, Optional, List, Literal
from playwright.async_api import (
    async_playwright,
    Playwright,
    Browser,
    BrowserContext,
    Page,
    CDPSession,
)
from starlette.websockets import WebSocket
from pydantic import Field
from suzent.tools.base import Tool, ToolGroup, ToolErrorCode, ToolResult
from suzent.logger import get_logger

logger = get_logger(__name__)


class BrowserSessionManager:
    _instance = None

    def __init__(self):
        self._playwright: Optional[Playwright] = None
        self._browser: Optional[Browser] = None
        self._context: Optional[BrowserContext] = None
        self._page: Optional[Page] = None
        self._client: Optional[CDPSession] = None
        self._websockets: List[WebSocket] = []
        self._streaming = False
        self._init_lock: Optional[asyncio.Lock] = None

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = BrowserSessionManager()
        return cls._instance

    def set_main_loop(self, loop):
        """Set the main application event loop for thread safety."""
        self._main_loop = loop
        # Initialize lock on the main loop
        self._init_lock = asyncio.Lock()

    async def _run_on_main_loop(self, coro):
        """Execute a coroutine on the main loop, handling cross-thread awaiting."""
        # If no main loop set (e.g. testing), try to use current or assume safety
        if not hasattr(self, "_main_loop") or self._main_loop is None:
            # Only warn once to avoid log spam
            if not getattr(self, "_warned_loop", False):
                logger.warning(
                    "No main loop set for BrowserSessionManager. Assuming current loop."
                )
                self._warned_loop = True
            return await coro

        # If we are already on the main loop, just await
        try:
            current_loop = asyncio.get_running_loop()
            if current_loop is self._main_loop:
                logger.debug("Already on main loop, awaiting coroutine directly.")
                return await coro
        except RuntimeError:
            logger.debug("No running loop in current thread.")
            pass

        # Otherwise, we are in a worker thread/loop -> dispatch to main
        logger.debug("Dispatching coroutine to main loop via threadsafe future.")
        future = asyncio.run_coroutine_threadsafe(coro, self._main_loop)

        try:
            # Add a timeout to avoid infinite hangs in the synchronization layer
            return await asyncio.wait_for(asyncio.wrap_future(future), timeout=60.0)
        except asyncio.TimeoutError:
            logger.error("Timeout waiting for main loop coroutine to complete.")
            raise RuntimeError("Browser action timed out during loop synchronization.")

    async def ensure_session(self, headless: bool = True):
        # We define the inner async function that does the actual work on the main loop
        async def _launch():
            logger.debug("Entering _launch coroutine.")
            if not self._init_lock:
                logger.warning("No init lock available, potential race condition.")
                # Create logical fallback lock if needed, but set_main_loop should be called.
                self._init_lock = asyncio.Lock()

            logger.debug("Waiting for init lock...")
            async with self._init_lock:
                logger.debug("Acquired init lock.")
                # Health Check: If page exists, verify it is still functional
                if self._page:
                    logger.debug("Page exists, performing health check...")
                    try:
                        # Check if browser and page are still alive
                        if not self._page.is_closed() and self._browser.is_connected():
                            logger.debug("Health check passed.")
                            return
                        logger.info(
                            "Browser session disconnected or page closed. Re-initializing..."
                        )
                    except Exception as e:
                        logger.warning(f"Health check failed: {e}. Re-initializing...")

                    # Cleanup old session if it failed health check
                    await self.close_session()

                logger.info("Starting Browser Session...")
                self._playwright = await async_playwright().start()
                logger.debug("Playwright started.")

                args = [
                    "--no-sandbox",
                    "--disable-setuid-sandbox",
                ]

                try:
                    logger.debug(f"Launching Chromium (headless={headless})...")
                    # Add a timeout to the launch call itself
                    try:
                        self._browser = await asyncio.wait_for(
                            self._playwright.chromium.launch(
                                headless=headless, args=args
                            ),
                            timeout=30.0,
                        )
                    except asyncio.TimeoutError:
                        logger.warning(
                            "Chromium launch timed out with args. Retrying without custom args..."
                        )
                        self._browser = await asyncio.wait_for(
                            self._playwright.chromium.launch(headless=headless),
                            timeout=30.0,
                        )
                    logger.debug("Chromium launched.")
                except Exception as e:
                    if (
                        "executable doesn't exist" in str(e).lower()
                        or "not installed" in str(e).lower()
                    ):
                        logger.info(
                            "Chromium not installed. Installing now (this may take a few minutes)..."
                        )
                        import sys

                        # Use asyncio.create_subprocess_exec to avoid blocking the loop
                        process = await asyncio.create_subprocess_exec(
                            sys.executable,
                            "-m",
                            "playwright",
                            "install",
                            "chromium",
                            stdout=asyncio.subprocess.PIPE,
                            stderr=asyncio.subprocess.PIPE,
                        )
                        stdout, stderr = await process.communicate()

                        if process.returncode == 0:
                            logger.info("Chromium installed successfully.")
                            self._browser = await self._playwright.chromium.launch(
                                headless=headless, args=args
                            )
                        else:
                            error_msg = (
                                stderr.decode().strip() or stdout.decode().strip()
                            )
                            raise RuntimeError(
                                f"Failed to install Chromium: {error_msg}"
                            )
                    else:
                        logger.error(f"Failed to launch browser: {e}")
                        raise

                logger.debug("Creating browser context...")
                # Create context with video size tailored for sidebar
                self._context = await self._browser.new_context(
                    viewport={"width": 1280, "height": 800}
                )
                self._page = await self._context.new_page()
                logger.debug("Browser context and page created.")

                # Connect CDP for low-level control
                self._client = await self._context.new_cdp_session(self._page)
                logger.debug("CDP session connected.")

                # Setup Screencast
                self._client.on("Page.screencastFrame", self._on_screencast_frame)

                # If we have waiting clients, start streaming immediately
                if self._websockets:
                    logger.debug("Starting streaming for existing clients...")
                    await self.start_streaming()

        # Execute on main loop
        logger.debug("Calling ensure_session...")
        await self._run_on_main_loop(_launch())
        logger.debug("ensure_session completed.")

    # --- Wrapper methods for Thread Safety ---

    async def goto(self, url: str):
        async def _fn():
            return await self._page.goto(
                url, wait_until="domcontentloaded", timeout=15000
            )

        return await self._run_on_main_loop(_fn())

    async def click(self, x: int, y: int):
        async def _fn():
            return await self._page.mouse.click(x, y)

        return await self._run_on_main_loop(_fn())

    async def scroll(self, dx: int, dy: int):
        async def _fn():
            return await self._page.mouse.wheel(dx, dy)

        return await self._run_on_main_loop(_fn())

    async def back(self):
        async def _fn():
            return await self._page.go_back(wait_until="domcontentloaded")

        return await self._run_on_main_loop(_fn())

    async def forward(self):
        async def _fn():
            return await self._page.go_forward(wait_until="domcontentloaded")

        return await self._run_on_main_loop(_fn())

    async def reload(self):
        async def _fn():
            return await self._page.reload(wait_until="domcontentloaded")

        return await self._run_on_main_loop(_fn())

    # --- Native Semantic Helpers ---
    _selector_map = {}

    async def get_snapshot(self, interactive_only: bool = True) -> ToolResult:
        """Generate a semantic snapshot and populate selector map."""

        async def _snap():
            # clear previous map
            self._selector_map.clear()

            # Simple heuristic script to find interactive elements
            js_script = """
            () => {
                const elements = Array.from(document.querySelectorAll('a, button, input, textarea, select, [role="button"], [role="link"]'));
                const visibleElements = elements.filter(el => {
                    const rect = el.getBoundingClientRect();
                    return rect.width > 0 && rect.height > 0 && window.getComputedStyle(el).visibility !== 'hidden';
                });
                
                return visibleElements.map((el, index) => {
                    // Generate a simple unique selector if possible, or use index strategy
                    el.setAttribute('data-agent-ref', index);
                    
                    let label = el.innerText || el.getAttribute('aria-label') || el.getAttribute('placeholder') || el.value || '';
                    label = label.substring(0, 50).replace(/\\n/g, ' ');
                    
                    return {
                        index: index,
                        tagName: el.tagName.toLowerCase(),
                        type: el.getAttribute('type'),
                        label: label,
                        href: el.getAttribute('href')
                    };
                });
            }
            """
            try:
                items = await self._page.evaluate(js_script)
            except Exception as e:
                logger.error(f"Snapshot script failed: {e}")
                return ToolResult.error_result(
                    ToolErrorCode.EXECUTION_FAILED,
                    "Error generating snapshot.",
                )

            output_lines = []
            for item in items:
                idx = item["index"]
                ref = f"@e{idx}"
                self._selector_map[ref] = f"[data-agent-ref='{idx}']"

                line = f"{ref}: <{item['tagName']}"
                if item["type"]:
                    line += f" type='{item['type']}'"
                if item["href"]:
                    line += f" href='{item['href']}'"
                line += f"> {item['label']}"
                output_lines.append(line)

            if output_lines:
                return ToolResult.success_result(
                    "\n".join(output_lines),
                    metadata={
                        "interactive_only": interactive_only,
                        "element_count": len(output_lines),
                    },
                )
            return ToolResult.success_result(
                "No interactive elements found.",
                metadata={"interactive_only": interactive_only, "element_count": 0},
            )

        return await self._run_on_main_loop(_snap())

    async def interact(self, action: str, ref: str, value: str = None) -> ToolResult:
        """Interact with an element by ref or selector."""

        async def _act():
            # Resolve ref if it exists
            selector = self._selector_map.get(ref, ref)

            try:
                if action == "click":
                    await self._page.click(selector)
                    return ToolResult.success_result(
                        f"Clicked {ref}", metadata={"action": action, "ref": ref}
                    )
                elif action == "dblclick":
                    await self._page.dblclick(selector)
                    return ToolResult.success_result(
                        f"Double-clicked {ref}", metadata={"action": action, "ref": ref}
                    )
                elif action == "fill":
                    await self._page.fill(selector, value or "")
                    return ToolResult.success_result(
                        f"Filled {ref} with '{value}'",
                        metadata={"action": action, "ref": ref},
                    )
                elif action == "type":
                    await self._page.type(selector, value or "")
                    return ToolResult.success_result(
                        f"Typed '{value}' into {ref}",
                        metadata={"action": action, "ref": ref},
                    )
                elif action == "hover":
                    await self._page.hover(selector)
                    return ToolResult.success_result(
                        f"Hovered {ref}", metadata={"action": action, "ref": ref}
                    )
                elif action == "press":
                    # value is key name
                    await self._page.press(selector, value)
                    return ToolResult.success_result(
                        f"Pressed '{value}' on {ref}",
                        metadata={"action": action, "ref": ref},
                    )
                else:
                    return ToolResult.error_result(
                        ToolErrorCode.INVALID_ARGUMENT,
                        f"Unknown action {action}",
                        metadata={"action": action, "ref": ref},
                    )
            except Exception as e:
                return ToolResult.error_result(
                    ToolErrorCode.EXECUTION_FAILED,
                    f"Interaction failed: {e}",
                    metadata={"action": action, "ref": ref},
                )

        return await self._run_on_main_loop(_act())

    async def _on_screencast_frame(self, params):
        """Handle incoming CDP screencast frames."""
        try:
            # Acknowledge the frame so CDP keeps sending them
            await self._client.send(
                "Page.screencastFrameAck", {"sessionId": params.get("sessionId")}
            )

            data = params.get("data")  # Base64 string
            metadata = params.get("metadata")

            if not data or not self._websockets:
                return

            # Broadcast to all connected websockets
            # We send raw bytes to avoid base64 overhead in WS if possible,
            # but for simplicity JSON wrapping might be safer initially.
            # Let's send a JSON message with the image.

            message = {
                "type": "frame",
                "data": data,
                "timestamp": metadata.get("timestamp"),
            }

            # Broadcast loop - use a copy to avoid modification during iteration
            disconnected = []
            for ws in list(self._websockets):
                try:
                    await ws.send_json(message)
                except Exception:
                    disconnected.append(ws)

            # Cleanup disconnected - safe removal (ws may already be removed by remove_client)
            for ws in disconnected:
                if ws in self._websockets:
                    self._websockets.remove(ws)

        except Exception as e:
            logger.error(f"Error handling screencast frame: {e}")

    async def start_streaming(self):
        # No-op if browser not initialized yet (lazy init)
        if self._streaming or not self._client:
            return
        logger.info("Starting CDP Screencast...")
        await self._client.send(
            "Page.startScreencast",
            {
                "format": "jpeg",
                "quality": 60,
                "maxWidth": 1280,
                "maxHeight": 800,
                "everyNthFrame": 1,  # Send every frame
            },
        )
        self._streaming = True

    async def stop_streaming(self):
        if not self._streaming or not self._client:
            return
        try:
            await self._client.send("Page.stopScreencast")
        except Exception as e:
            # Ignore if target is closed (browser/page already gone)
            logger.debug(f"Error stopping screencast (likely closed): {e}")
        self._streaming = False

    async def add_client(self, websocket: WebSocket):
        """Accept WebSocket client without launching browser (lazy init)."""
        await websocket.accept()
        self._websockets.append(websocket)
        # If browser already running, start streaming for this client
        if self._client and not self._streaming:
            await self.start_streaming()
        # Otherwise, browser will be launched lazily when needed

    async def remove_client(self, websocket: WebSocket):
        if websocket in self._websockets:
            self._websockets.remove(websocket)
        # If no clients left, maybe stop streaming to save resources?
        if not self._websockets:
            await self.stop_streaming()

    def _get_mouse_coords(self, message: dict) -> tuple[float | None, float | None]:
        """Extract and validate mouse coordinates from a message."""
        return message.get("x"), message.get("y")

    async def handle_client_message(self, message: dict):
        """Process interaction events from the frontend."""
        action = message.get("type")

        # Lazy init: launch browser on navigate command
        if action == "navigate":
            url = message.get("url")
            if url:
                await self.ensure_session()
                await self.start_streaming()
                await self._page.goto(url)
            return

        # Other actions require browser to be already running
        if not self._page:
            return

        try:
            # Mouse actions with coordinate validation
            if action in ("click", "mousedown", "mouseup", "mousemove"):
                x, y = self._get_mouse_coords(message)
                if x is None or y is None:
                    return

                if action == "click":
                    await self._page.mouse.click(x, y)
                elif action == "mousedown":
                    await self._page.mouse.move(x, y)
                    await self._page.mouse.down()
                elif action == "mouseup":
                    await self._page.mouse.move(x, y)
                    await self._page.mouse.up()
                elif action == "mousemove":
                    await self._page.mouse.move(x, y)

            elif action == "type":
                text = message.get("text")
                if text:
                    await self._page.keyboard.type(text)

            elif action == "key":
                key = message.get("key")
                if key:
                    await self._page.keyboard.press(key)

            elif action == "scroll":
                dx, dy = message.get("dx", 0), message.get("dy", 0)
                await self._page.mouse.wheel(dx, dy)

        except Exception as e:
            logger.error(f"Error handling client browser interaction: {e}")

    async def close_session(self):
        """Clean up browser resources."""
        logger.info("Closing Browser Session...")
        await self.stop_streaming()

        if self._context:
            try:
                await self._context.close()
            except Exception as e:
                logger.debug(f"Ignored error closing context: {e}")
            self._context = None

        if self._browser:
            try:
                await self._browser.close()
            except Exception as e:
                # Common race condition on Ctrl+C: driver dies before we close
                logger.debug(
                    f"Ignored error closing browser (likely already closed): {e}"
                )
            self._browser = None

        if self._playwright:
            try:
                await self._playwright.stop()
            except Exception as e:
                logger.debug(f"Ignored error stopping playwright: {e}")
            self._playwright = None

        self._page = None
        self._client = None
        logger.info("Browser Session Closed.")


class BrowsingTool(Tool):
    name = "BrowsingTool"
    tool_name = "browser_action"
    group = ToolGroup.WEB

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.session_mgr = BrowserSessionManager.get_instance()

    async def _execute(self, command: str, arguments: list = None):
        arguments = arguments or []

        # Ensure session exists (this also starts CDP if needed)
        await self.session_mgr.ensure_session()

        # --- PATH A: NATIVE (System/User Logic) ---
        if command == "open":
            url = arguments[0] if arguments else "about:blank"
            if not url.startswith("http"):
                url = "https://" + url
            try:
                # Use manager wrapper for thread safety
                await self.session_mgr.goto(url)
            except Exception as e:
                return ToolResult.error_result(
                    ToolErrorCode.EXECUTION_FAILED,
                    f"Error opening {url} (partial load): {e}",
                    metadata={"command": command, "url": url},
                )
            return ToolResult.success_result(
                f"Opened {url}", metadata={"command": command, "url": url}
            )

        elif command == "back":
            await self.session_mgr.back()
            return ToolResult.success_result(
                "Navigated back.", metadata={"command": command}
            )

        elif command == "forward":
            await self.session_mgr.forward()
            return ToolResult.success_result(
                "Navigated forward.", metadata={"command": command}
            )

        elif command == "reload" or command == "refresh":
            await self.session_mgr.reload()
            return ToolResult.success_result(
                "Reloaded page.", metadata={"command": command}
            )

        elif command == "click_coords":
            x, y = int(arguments[0]), int(arguments[1])
            # Use manager wrapper for thread safety
            await self.session_mgr.click(x, y)
            return ToolResult.success_result(
                f"Clicked at {x}, {y}", metadata={"command": command, "x": x, "y": y}
            )

        elif command == "scroll":
            # arguments: [dx, dy]
            # If simplistic usage: just scroll down
            await self.session_mgr.scroll(0, 500)
            return ToolResult.success_result(
                "Scrolled down.", metadata={"command": command}
            )

        # --- PATH B: PYTHON NATIVE SEMANTIC LOGIC (Replacing Agent-Browser CLI) ---
        # The Agent uses 'snapshot' to "reason" about the page (e.g. use @e1 locators)
        # We now handle this natively via Playwright injection for reliability.

        if command == "snapshot":
            # arguments like ['-i'] are ignored as we default to interactive for now
            return await self.session_mgr.get_snapshot(interactive_only=True)

        elif command in ["click", "dblclick", "hover"]:
            if not arguments:
                return ToolResult.error_result(
                    ToolErrorCode.MISSING_REQUIRED_PARAM,
                    f"{command} requires a target ref (e.g. {command} @e1)",
                    metadata={"command": command},
                )
            ref = arguments[0]
            return await self.session_mgr.interact(command, ref)

        elif command in ["fill", "type"]:
            if not arguments:
                return ToolResult.error_result(
                    ToolErrorCode.MISSING_REQUIRED_PARAM,
                    f"{command} requires a target ref (e.g. {command} @e1 'text')",
                    metadata={"command": command},
                )
            ref = arguments[0]
            val = arguments[1] if len(arguments) > 1 else ""
            return await self.session_mgr.interact(command, ref, val)

        elif command == "press":
            if not arguments:
                return ToolResult.error_result(
                    ToolErrorCode.MISSING_REQUIRED_PARAM,
                    "press requires ref and key (e.g. press @e1 Enter)",
                    metadata={"command": command},
                )
            ref = arguments[0]
            key = arguments[1] if len(arguments) > 1 else "Enter"
            return await self.session_mgr.interact("press", ref, key)

        return ToolResult.error_result(
            ToolErrorCode.INVALID_ARGUMENT,
            f"Unknown command {command}",
            metadata={"command": command},
        )

    async def forward(
        self,
        command: Annotated[
            Literal[
                "open",
                "snapshot",
                "click",
                "dblclick",
                "hover",
                "fill",
                "type",
                "press",
                "back",
                "forward",
                "reload",
                "refresh",
                "click_coords",
                "scroll",
            ],
            Field(description="Browser command to execute."),
        ],
        arguments: Annotated[
            Optional[list[str]],
            Field(
                default=None,
                description="Optional command arguments such as a URL, ref, text, or coordinates.",
            ),
        ] = None,
    ) -> ToolResult:
        """Control a browser session to navigate and interact with web pages.

        Optimal workflow: open <url>, then snapshot to get element refs, then click/fill using refs.

        Args:
            command: The command to execute (open, snapshot, click, fill, scroll, back, forward, refresh, click_coords).
            arguments: Optional arguments for the command (e.g., url, selector, text).
        """
        return await self._execute(command, arguments)
