from functools import wraps
from typing import Any, Dict, Optional

from .client import MCPClient


def _ensure_client(func):
    """Decorator kiểm tra self.client != None trước khi gọi tool."""

    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if self.client is None:
            return {
                "ok": False,
                "error": "MCP client not connected. Call connect mcp server first.",
            }
        return func(self, *args, **kwargs)

    return wrapper


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):
        self.client = None

    # ======================================================
    # SESSION MANAGEMENT
    # ======================================================
    def connect_mcp(self, mcpUrl: str) -> Dict[str, Any]:
        # sid = self.client.create_session(mcpUrl)
        self.client = MCPClient(base_url=mcpUrl, headers=None, timeout=30, retries=2)
        return {"ok": True, "cdpUrl": "http://localhost:9222"}

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

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

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

    # ======================================================
    # NAVIGATION & DOM
    # ======================================================
    @_ensure_client
    def open_page(self, sessionId: str, url: str) -> Dict[str, Any]:
        return self.client.call_tool(
            "openPage", {"sessionId": sessionId, "url": url}
        ).get("structuredContent", {})

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

    @_ensure_client
    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]

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

    @_ensure_client
    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", {})

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

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

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

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

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

    @_ensure_client
    def upload_file(
        self,
        sessionId: str,
        selector: str,
        file_path: str,
    ) -> Dict[str, Any]:
        """
        Upload file (kể cả video lớn) vào input[type=file] theo luồng mới:
        1. Multipart upload file lên MCP server
        2. Nhận uploadId
        3. Gọi MCP tool uploadFile với uploadId

        Args:
            sessionId: MCP browser session
            selector: CSS selector, ví dụ 'input[type=file]'
            file_path: đường dẫn file local (video, pdf, doc, ...)
        """

        if not file_path:
            return {"ok": False, "error": "file_path is required"}

        # --------------------------------------------------
        # 1️⃣ Multipart upload file lên MCP server
        # --------------------------------------------------
        try:
            with open(file_path, "rb") as f:
                resp = self.client.http.post(
                    "/upload",
                    files={"file": f},
                    timeout=300,  # upload file lớn
                )
        except Exception as e:
            return {"ok": False, "error": f"upload http failed: {e}"}

        if resp.status_code != 200:
            return {
                "ok": False,
                "error": f"upload http error {resp.status_code}: {resp.text}",
            }

        data = resp.json()
        upload_id = data.get("uploadId")
        if not upload_id:
            return {"ok": False, "error": "uploadId not returned from server"}

        # --------------------------------------------------
        # 2️⃣ Gọi MCP tool uploadFile (PATH MODE)
        # --------------------------------------------------
        result = self.client.call_tool(
            "uploadFile",
            {
                "sessionId": sessionId,
                "selector": selector,
                "uploadId": upload_id,
            },
        )

        return result.get("structuredContent", {})

    @_ensure_client
    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
    # ======================================================
    @_ensure_client
    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", {})

    @_ensure_client
    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
    # ======================================================
    @_ensure_client
    def click_to_text(self, sessionId: str, text: str) -> dict:
        return self.client.call_tool(
            "clickToText", {"sessionId": sessionId, "text": text}
        ).get("structuredContent", {})

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

    @_ensure_client
    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", {})

    @_ensure_client
    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", {})

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

    @_ensure_client
    def get_dom_tree(self, sessionId, args=None):
        return self.client.call_tool(
            "getDomTree", {"sessionId": sessionId, "args": args or {}}
        )

    @_ensure_client
    def get_clickable(self, sessionId, args=None):
        return self.client.call_tool(
            "getClickable", {"sessionId": sessionId, "args": args or {}}
        )

    @_ensure_client
    def selector_map(self, sessionId, selector, args=None):
        return self.client.call_tool(
            "selectorMap",
            {"sessionId": sessionId, "selector": selector, "args": args or {}},
        )

    @_ensure_client
    def find_element_by_prompt(self, sessionId: str, prompt: str) -> Dict[str, Any]:
        """
        Gọi tool findElementByPrompt trên MCP server.
        Trả về structuredContent gồm: html, nodeId.
        """
        return self.client.call_tool(
            "findElementByPrompt", {"sessionId": sessionId, "prompt": prompt}
        ).get("structuredContent", {})

    # ======================================================
    # AI / CONTENT PARSING
    # ======================================================
    @_ensure_client
    def parse_html_by_prompt(self, html: str, prompt: str) -> Dict[str, Any]:
        """
        Parse HTML content using AI with dynamic prompt-defined structure.

        Args:
            html: Raw HTML string (client-provided)
            prompt: Instruction that defines what to extract and output structure
                    Example:
                      - "Hãy lấy nội dung bài viết, struct trả về { content }"
                      - "Hãy lấy số lượng like, share, comment, trả JSON { like, share, comment }"

        Returns:
            structuredContent (dynamic JSON defined by prompt)
        """
        return self.client.call_tool(
            "parseHTMLByPrompt",
            {
                "html": html,
                "prompt": prompt,
            },
        ).get("structuredContent", {})

    # ======================================================
    # CLEAN TEXT / READ MODE
    # ======================================================
    @_ensure_client
    def get_clean_text(self, sessionId: str) -> Dict[str, Any]:
        """
        Lấy toàn bộ visible text đã được clean trên trang hiện tại.
        - Bỏ script/style/iframe/svg/canvas
        - Chỉ text nhìn thấy (display/visibility/opacity)

        Returns:
            {
              "text": "...",
              "length": 12345
            }
        """
        return self.client.call_tool(
            "getCleanText",
            {"sessionId": sessionId},
        ).get("structuredContent", {})

    @_ensure_client
    def evaluate_stream(
        self,
        sessionId: str,
        expression: str,
        chunkSize: int = 100,
    ) -> Dict[str, Any]:
        """
        Evaluate JS expression theo chế độ STREAM.
        Dùng khi kết quả là array lớn (DOM, list, table, ...)

        Returns (init response):
        {
          "stream_id": "...",
          "total": 1234,
          "chunk_size": 100
        }
        """
        return self.client.call_tool(
            "evaluate.stream",
            {
                "sessionId": sessionId,
                "expression": expression,
                "chunkSize": int(chunkSize),
            },
        ).get("structuredContent", {})

    @_ensure_client
    def stream_pull(
        self,
        stream_id: str,
        offset: int = 0,
        limit: int = 100,
    ) -> Dict[str, Any]:
        """
        Kéo 1 chunk từ stream đã tạo bởi evaluate.stream

        Returns:
        {
          "items": [...],
          "offset": 0,
          "has_more": true
        }
        """
        return self.client.call_tool(
            "stream.pull",
            {
                "stream_id": stream_id,
                "offset": int(offset),
                "limit": int(limit),
            },
        ).get("structuredContent", {})

    @_ensure_client
    def evaluate_stream_all(
        self,
        sessionId: str,
        expression: str,
        chunkSize: int = 100,
        max_items: Optional[int] = None,
    ):
        """
        Helper: evaluate.stream + tự động pull toàn bộ dữ liệu.

        ⚠️ Chỉ dùng khi bạn THỰC SỰ cần full data.
        """
        init = self.evaluate_stream(
            sessionId=sessionId,
            expression=expression,
            chunkSize=chunkSize,
        )

        stream_id = init.get("stream_id")
        total = init.get("total", 0)

        if not stream_id:
            return []

        items = []
        offset = 0

        while True:
            chunk = self.stream_pull(
                stream_id=stream_id,
                offset=offset,
                limit=chunkSize,
            )

            part = chunk.get("items", [])
            items.extend(part)

            if max_items and len(items) >= max_items:
                return items[:max_items]

            if not chunk.get("has_more"):
                break

            offset += chunkSize

        return items

    @_ensure_client
    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["timeout"] = int(timeoutMs)

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

    @_ensure_client
    def send_keys(
        self, sessionId: str, key: str, interval: int = 100
    ) -> Dict[str, Any]:
        """
        Gửi phím: enter, tab, esc, ctrl+c, ctrl+shift+tab, ...
        """
        return self.client.call_tool(
            "sendKeys",
            {"sessionId": sessionId, "text": key, "interval": interval},
        ).get("structuredContent", {})

    @_ensure_client
    def perform(
        self,
        sessionId: str,
        action: str,
        target: Optional[str] = None,
        value: Optional[str] = None,
        x: Optional[float] = None,
        y: Optional[float] = None,
        from_point: Optional[dict] = None,
        to_point: Optional[dict] = None,
    ) -> Dict[str, Any]:
        """
        action:
        - click (x,y)
        - move / hover (x,y)
        - drag (from -> to)
        - type (target, value)
        """
        args = {
            "sessionId": sessionId,
            "action": action,
        }

        if target is not None:
            args["target"] = target
        if value is not None:
            args["value"] = value
        if x is not None:
            args["x"] = float(x)
        if y is not None:
            args["y"] = float(y)
        if from_point is not None:
            args["from"] = from_point
        if to_point is not None:
            args["to"] = to_point

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

    @_ensure_client
    def drag_and_drop(
        self,
        sessionId: str,
        from_x: float,
        from_y: float,
        to_x: float,
        to_y: float,
    ):
        return self.perform(
            sessionId=sessionId,
            action="drag",
            from_point={"x": from_x, "y": from_y},
            to_point={"x": to_x, "y": to_y},
        )

    def hover(self, sessionId: str, x: float, y: float):
        return self.perform(
            sessionId=sessionId,
            action="hover",
            x=x,
            y=y,
        )

    @_ensure_client
    def scroll(
        self,
        sessionId: str,
        *,
        x: Optional[float] = None,
        y: Optional[float] = None,
        selector: Optional[str] = None,
        position: Optional[str] = None,  # "top" | "bottom"
    ) -> Dict[str, Any]:
        """
        Scroll trang web.

        Cách dùng:
        - Scroll theo pixel:
            scroll(sessionId, y=500)
        - Scroll tới element:
            scroll(sessionId, selector="#footer")
        - Scroll top / bottom:
            scroll(sessionId, position="top")
            scroll(sessionId, position="bottom")
        """

        args: Dict[str, Any] = {
            "sessionId": sessionId,
        }

        if x is not None:
            args["x"] = float(x)
        if y is not None:
            args["y"] = float(y)
        if selector is not None:
            args["selector"] = selector
        if position is not None:
            args["position"] = position

        if len(args) == 1:
            return {
                "ok": False,
                "error": "scroll requires x/y or selector or position",
            }

        return self.client.call_tool(
            "scroll",
            args,
        ).get("structuredContent", {})
        # ======================================================
    # BROWSER RUNTIME (NO SESSION)
    # ======================================================
    @_ensure_client
    def create_browser(self, payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        Start browser runtime (idempotent).
        Không tạo session, chỉ đảm bảo browser đang RUNNING.

        payload: map cho ThirdPartyOpenRequest
        """
        return self.client.call_tool(
            "create_browser",
            payload or {},
        ).get("structuredContent", {})

    @_ensure_client
    def release_browser(self) -> Dict[str, Any]:
        """
        Release / stop browser runtime.
        Không cần payload.
        """
        return self.client.call_tool(
            "release_browser",
            {},
        ).get("structuredContent", {})
