from typing import Any, Dict, Optional

from .client import MCPClient


class MCPTools:
    """
    Wrapper chuẩn cho Google ADK + MCP Server.
    - Tool nào trả text/html/... → dùng structuredContent
    - Tool screenshot → trả đúng content để ADK Web hiển thị image
    """

    def __init__(self, client: MCPClient):
        self.client = client

    # ======================================================
    # SESSION MANAGEMENT
    # ======================================================

    def create_session(self, cdpUrl: str) -> Dict[str, Any]:
        sid = self.client.create_session(cdpUrl)
        return {"sessionId": sid}

    def close_session(self, sessionId: str) -> Dict[str, Any]:
        ok = self.client.close_session(sessionId)
        return {"ok": bool(ok)}

    def list_sessions(self) -> Dict[str, Any]:
        return {"sessions": self.client.list_local_sessions()}

    # ======================================================
    # NAVIGATION & DOM
    # ======================================================

    def open_page(self, sessionId: str, url: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "openPage", {"sessionId": sessionId, "url": url}
        ).get("structuredContent", {})

    def get_html(self, sessionId: str) -> Dict[str, Any]:
        return self.client.call_tool("getHTML", {"sessionId": sessionId}).get(
            "structuredContent", {}
        )

    def screenshot(self, sessionId: str) -> Dict[str, Any]:
        """
        Trả về đúng phần IMAGE content:
        {
          "type": "image",
          "mimeType": "image/png",
          "data": "<base64>"
        }
        """
        full = self.client.call_tool("screenshot", {"sessionId": sessionId})
        return full["content"][0]

    def click(self, sessionId: str, selector: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "click", {"sessionId": sessionId, "selector": selector}
        ).get("structuredContent", {})

    def type(self, sessionId: str, selector: str, text: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "type", {"sessionId": sessionId, "selector": selector, "text": text}
        ).get("structuredContent", {})

    def evaluate(self, sessionId: str, expression: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "evaluate", {"sessionId": sessionId, "expression": expression}
        ).get("structuredContent", {})

    # ======================================================
    # ELEMENT UTILITIES
    # ======================================================

    def find_element(self, sessionId: str, selector: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "findElement", {"sessionId": sessionId, "selector": selector}
        ).get("structuredContent", {})

    def find_all(self, sessionId: str, selector: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "findAll", {"sessionId": sessionId, "selector": selector}
        ).get("structuredContent", {})

    def get_bounding_box(self, sessionId: str, selector: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "getBoundingBox", {"sessionId": sessionId, "selector": selector}
        ).get("structuredContent", {})

    def click_bounding_box(self, sessionId: str, selector: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "clickBoundingBox", {"sessionId": sessionId, "selector": selector}
        ).get("structuredContent", {})

    def upload_file(
        self, sessionId: str, selector: str, filename: str, base64data: str
    ) -> Dict[str, Any]:
        """
        Upload file vào input[type=file]
        Args:
            sessionId: MCP browser session
            selector: CSS selector, ví dụ 'input[type=file]'
            filename: tên file trên browser side
            base64data: dữ liệu base64 (không kèm header)
        Returns:
            structured result từ server
        """
        return self.client.call_tool(
            "uploadFile",
            {
                "sessionId": sessionId,
                "selector": selector,
                "filename": filename,
                "data": base64data,
            },
        ).get("structuredContent", {})

    def wait_for_selector(
        self, sessionId: str, selector: str, timeoutMs: Optional[int] = None
    ) -> Dict[str, Any]:
        args = {"sessionId": sessionId, "selector": selector}
        if timeoutMs is not None:
            args["timeoutMs"] = int(timeoutMs)

        return self.client.call_tool("waitForSelector", args).get(
            "structuredContent", {}
        )

    # ======================================================
    # TAB MANAGEMENT
    # ======================================================

    def new_tab(
        self, sessionId: str, url: Optional[str] = "about:blank"
    ) -> Dict[str, Any]:
        return self.client.call_tool(
            "newTab", {"sessionId": sessionId, "url": url}
        ).get("structuredContent", {})

    def switch_tab(self, sessionId: str, targetId: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "switchTab", {"sessionId": sessionId, "targetId": targetId}
        ).get("structuredContent", {})

    # ======================================================
    # ADVANCED ACTIONS
    # ======================================================

    def click_to_text(self, sessionId: str, text: str) -> dict:
        return self.client.call_tool(
            "clickToText", {"sessionId": sessionId, "text": text}
        ).get("structuredContent", {})

    def find_element_xpath(self, sessionId: str, xpath: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "findElementByXPath", {"sessionId": sessionId, "xpath": xpath}
        ).get("structuredContent", {})

    def find_element_by_text(self, sessionId: str, text: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "findElementByText", {"sessionId": sessionId, "text": text}
        ).get("structuredContent", {})

    def click_by_node_id(self, sessionId: str, nodeId: int) -> Dict[str, Any]:
        return self.client.call_tool(
            "clickByNodeId", {"sessionId": sessionId, "nodeId": nodeId}
        ).get("structuredContent", {})

    def import_cookies(self, sessionId: str, cookies: dict) -> Dict[str, Any]:
        return self.client.call_tool(
            "importCookies", {"sessionId": sessionId, "cookies": cookies}
        ).get("structuredContent", {})
