import logging
import os
from functools import wraps
from typing import Any, Dict, List, Optional

from .client import MCPClient

logger = logging.getLogger(__name__)


def _ensure_client(func):
    @wraps(func)
    async def wrapper(self, *args, **kwargs):
        if self.client is None:
            logger.error("MCP async client not connected")
            return {
                "ok": False,
                "error": "MCP async client not connected. Provide client or call connect_mcp first.",
            }
        return await func(self, *args, **kwargs)

    return wrapper


class MCPTools:
    """Asynchronous tools wrapper for MCP HTTP server."""

    def __init__(self, client: Optional[MCPClient] = None):
        self.client = client
        logger.info("Async MCPTools initialized")

    def connect_mcp(
        self,
        mcp_url: Optional[str] = None,
        mcpUrl: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        timeout: int = 30,
        retries: int = 2,
        proxies: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        url = (mcp_url or mcpUrl or "").strip()
        if not url:
            return {"ok": False, "error": "mcp_url (or mcpUrl) is required"}

        self.client = MCPClient(
            base_url=url,
            headers=headers,
            timeout=timeout,
            retries=retries,
            proxies=proxies,
        )
        return {"ok": True, "base_url": url}

    @_ensure_client
    async def close(self) -> Dict[str, Any]:
        await self.client.aclose()
        return {"ok": True}

    @_ensure_client
    async def _call_tool(self, tool: str, arguments: Optional[Dict[str, Any]] = None):
        return await self.client.call_tool(tool, arguments or {})

    @_ensure_client
    async def _structured(self, tool: str, arguments: Optional[Dict[str, Any]] = None):
        return (await self._call_tool(tool, arguments)).get("structuredContent", {})

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

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

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

    @_ensure_client
    async def open_page(self, sessionId: str, url: str):
        return await self._structured("openPage", {"sessionId": sessionId, "url": url})

    @_ensure_client
    async def get_html(self, sessionId: str):
        return await self._structured("getHTML", {"sessionId": sessionId})

    @_ensure_client
    async def click(self, sessionId: str, selector: str):
        return await self._structured("click", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def get_bounding_box(self, sessionId: str, selector: str):
        return await self._structured("getBoundingBox", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def click_to_text(self, sessionId: str, text: str):
        return await self._structured("clickToText", {"sessionId": sessionId, "text": text})

    @_ensure_client
    async def click_by_node_id(self, sessionId: str, nodeId: int):
        return await self._structured("clickByNodeId", {"sessionId": sessionId, "nodeId": nodeId})

    @_ensure_client
    async def screenshot(self, sessionId: str):
        return await self._call_tool("screenshot", {"sessionId": sessionId})

    @_ensure_client
    async def import_cookies(self, sessionId: str, cookies: List[Dict[str, Any]]):
        return await self._structured("importCookies", {"sessionId": sessionId, "cookies": cookies})

    @_ensure_client
    async def click_bounding_box(self, sessionId: str, selector: str):
        return await self._structured("clickBoundingBox", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def upload_file_to_server(self, file_path: str):
        logger.info("upload_file_to_server.start | file_path=%s", file_path)

        if not file_path:
            logger.error("upload_file_to_server.failed | reason=file_path_missing")
            return {"ok": False, "error": "file_path is required"}

        resolved_path = os.path.abspath(file_path)

        if not os.path.exists(resolved_path):
            logger.error(
                "upload_file_to_server.failed | reason=file_not_found | file_path=%s",
                resolved_path,
            )
            return {"ok": False, "error": f"file not found: {resolved_path}"}

        if not os.path.isfile(resolved_path):
            logger.error(
                "upload_file_to_server.failed | reason=not_a_file | file_path=%s",
                resolved_path,
            )
            return {"ok": False, "error": f"path is not a file: {resolved_path}"}

        file_name = os.path.basename(resolved_path)
        file_size = os.path.getsize(resolved_path)
        logger.info(
            "upload_file_to_server.prepare | file_name=%s | size_bytes=%s | file_path=%s",
            file_name,
            file_size,
            resolved_path,
        )
        upload_url = self.client.upload_url()
        logger.info("upload_file_to_server.endpoint | upload_url=%s", upload_url)

        try:
            with open(resolved_path, "rb") as f:
                files = {"file": (file_name, f, "application/octet-stream")}
                resp = await self.client.post_upload(upload_url, files=files, timeout=300)
        except Exception as exc:
            logger.exception(
                "upload_file_to_server.failed | reason=http_exception | file_path=%s",
                resolved_path,
            )
            return {"ok": False, "error": f"upload http failed: {exc}"}

        if resp.status_code != 200:
            logger.error(
                "upload_file_to_server.failed | reason=http_status | status_code=%s | file_path=%s",
                resp.status_code,
                resolved_path,
            )
            return {"ok": False, "error": f"http {resp.status_code}: {resp.text}"}

        upload_id = resp.json().get("uploadId")

        if not upload_id:
            logger.error(
                "upload_file_to_server.failed | reason=missing_upload_id | file_path=%s",
                resolved_path,
            )
            return {"ok": False, "error": "uploadId not returned"}

        logger.info(
            "upload_file_to_server.success | uploadId=%s | file_name=%s | size_bytes=%s",
            upload_id,
            file_name,
            file_size,
        )
        return {
            "ok": True,
            "uploadId": upload_id,
            "filePath": resolved_path,
            "fileName": file_name,
            "sizeBytes": file_size,
            "message": (
                "File uploaded to MCP server successfully. "
                "Use uploadId with set_file_for_input to attach it to a file input."
            ),
        }

    @_ensure_client
    async def set_file_for_input(self, sessionId: str, selector: str, uploadId: str):
        logger.info(
            "set_file_for_input.start | sessionId=%s | selector=%s | uploadId=%s",
            sessionId,
            selector,
            uploadId,
        )

        if not sessionId:
            logger.error("set_file_for_input.failed | reason=session_id_missing")
            return {"ok": False, "error": "sessionId is required"}

        if not selector:
            logger.error(
                "set_file_for_input.failed | reason=selector_missing | sessionId=%s",
                sessionId,
            )
            return {"ok": False, "error": "selector is required"}

        if not uploadId:
            logger.error(
                "set_file_for_input.failed | reason=upload_id_missing | sessionId=%s | selector=%s",
                sessionId,
                selector,
            )
            return {"ok": False, "error": "uploadId is required"}

        try:
            result = await self._structured(
                "uploadFile",
                {"sessionId": sessionId, "selector": selector, "uploadId": uploadId},
            )
        except Exception as exc:
            logger.exception(
                "set_file_for_input.failed | reason=tool_exception | sessionId=%s | selector=%s | uploadId=%s",
                sessionId,
                selector,
                uploadId,
            )
            return {"ok": False, "error": f"uploadFile tool failed: {exc}"}

        logger.info(
            "set_file_for_input.success | sessionId=%s | selector=%s | uploadId=%s | ok=%s",
            sessionId,
            selector,
            uploadId,
            result.get("ok"),
        )
        return {
            **result,
            "sessionId": sessionId,
            "selector": selector,
            "uploadId": result.get("uploadId", uploadId),
            "message": (
                "UploadId has been attached to the target file input successfully."
                if result.get("ok")
                else "uploadFile tool returned without ok=true."
            ),
        }

    @_ensure_client
    async def upload_file(self, sessionId: str, selector: str, file_path: str):
        logger.info(
            "upload_file.start | sessionId=%s | selector=%s | file_path=%s",
            sessionId,
            selector,
            file_path,
        )

        upload_result = await self.upload_file_to_server(file_path)
        if not upload_result.get("ok"):
            logger.error(
                "upload_file.failed | step=upload_file_to_server | file_path=%s | error=%s",
                file_path,
                upload_result.get("error"),
            )
            return upload_result

        set_result = await self.set_file_for_input(
            sessionId=sessionId,
            selector=selector,
            uploadId=upload_result["uploadId"],
        )
        if not set_result.get("ok"):
            logger.error(
                "upload_file.failed | step=set_file_for_input | sessionId=%s | selector=%s | uploadId=%s | error=%s",
                sessionId,
                selector,
                upload_result["uploadId"],
                set_result.get("error"),
            )
            return set_result

        logger.info(
            "upload_file.success | sessionId=%s | selector=%s | uploadId=%s",
            sessionId,
            selector,
            upload_result["uploadId"],
        )
        return {
            **set_result,
            "filePath": upload_result.get("filePath"),
            "fileName": upload_result.get("fileName"),
            "sizeBytes": upload_result.get("sizeBytes"),
            "message": (
                "File uploaded to MCP server and attached to the target input successfully."
            ),
        }

    @_ensure_client
    async 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[str, float]] = None,
        to_point: Optional[Dict[str, float]] = None,
    ):
        args: Dict[str, Any] = {"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 await self._structured("perform", args)

    @_ensure_client
    async def wait_for_selector(
        self,
        sessionId: str,
        selector: str,
        timeout: Optional[int] = None,
        timeoutMs: Optional[int] = None,
    ):
        args: Dict[str, Any] = {"sessionId": sessionId, "selector": selector}
        resolved_timeout = timeout if timeout is not None else timeoutMs
        if resolved_timeout is not None:
            args["timeout"] = int(resolved_timeout)
        return await self._structured("waitForSelector", args)

    @_ensure_client
    async def scroll(
        self,
        sessionId: str,
        x: Optional[float] = None,
        y: Optional[float] = None,
        selector: Optional[str] = None,
        position: Optional[str] = None,
    ):
        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
        return await self._structured("scroll", args)

    @_ensure_client
    async def evaluate(self, sessionId: str, expression: str):
        return await self._structured("evaluate", {"sessionId": sessionId, "expression": expression})

    @_ensure_client
    async def evaluate_stream(self, sessionId: str, expression: str, chunkSize: int = 100):
        return await self._structured(
            "evaluate.stream",
            {"sessionId": sessionId, "expression": expression, "chunkSize": chunkSize},
        )

    @_ensure_client
    async def stream_pull(self, stream_id: str, offset: int = 0, limit: int = 100):
        return await self._structured(
            "stream.pull", {"stream_id": stream_id, "offset": offset, "limit": limit}
        )

    @_ensure_client
    async def get_cookies(self, sessionId: str):
        return await self._structured("getCookies", {"sessionId": sessionId})

    @_ensure_client
    async def parse_html_by_prompt(self, html: str, prompt: str):
        return await self._structured("parseHTMLByPrompt", {"html": html, "prompt": prompt})

    @_ensure_client
    async def viewport(self, sessionId: str, viewport: Optional[Dict[str, Any]] = None):
        args: Dict[str, Any] = {"sessionId": sessionId}
        if viewport is not None:
            args["viewport"] = viewport
        return await self._structured("viewport", args)

    # Input tools
    @_ensure_client
    async def send_keys(self, sessionId: str, text: str, interval: Optional[int] = None):
        args: Dict[str, Any] = {"sessionId": sessionId, "text": text}
        if interval is not None:
            args["interval"] = int(interval)
        return await self._structured("sendKeys", args)

    @_ensure_client
    async def type(self, sessionId: str, selector: str, text: str):
        return await self._structured("type", {"sessionId": sessionId, "selector": selector, "text": text})

    @_ensure_client
    async def send_text(self, sessionId: str, text: str, interval: Optional[int] = None):
        args: Dict[str, Any] = {"sessionId": sessionId, "text": text}
        if interval is not None:
            args["interval"] = int(interval)
        return await self._structured("sendText", args)

    @_ensure_client
    async def clear_all(self, sessionId: str):
        return await self._structured("clearAll", {"sessionId": sessionId})

    @_ensure_client
    async def send_combo(self, sessionId: str, text: str):
        return await self._structured("sendCombo", {"sessionId": sessionId, "text": text})

    # Tab tools
    @_ensure_client
    async def get_current_url(self, sessionId: str):
        return await self._structured("getCurrentUrl", {"sessionId": sessionId})

    @_ensure_client
    async def list_tabs(self, sessionId: str):
        return await self._structured("listTabs", {"sessionId": sessionId})

    @_ensure_client
    async def switch_tab(
        self,
        sessionId: str,
        tabId: Optional[str] = None,
        targetId: Optional[str] = None,
    ):
        resolved_tab_id = tabId or targetId
        if not resolved_tab_id:
            return {"ok": False, "error": "tabId (or targetId) is required"}
        return await self._structured("switchTab", {"sessionId": sessionId, "tabId": resolved_tab_id})

    @_ensure_client
    async def close_tab(self, sessionId: str, tabId: str):
        return await self._structured("closeTab", {"sessionId": sessionId, "tabId": tabId})

    @_ensure_client
    async def new_tab(self, sessionId: str, url: Optional[str] = None):
        args: Dict[str, Any] = {"sessionId": sessionId}
        if url:
            args["url"] = url
        return await self._structured("newTab", args)

    @_ensure_client
    async def switch_tab_by_title(self, sessionId: str, title: str):
        return await self._structured("switchTabByTitle", {"sessionId": sessionId, "title": title})

    @_ensure_client
    async def switch_tab_by_url(self, sessionId: str, url: str):
        return await self._structured("switchTabByUrl", {"sessionId": sessionId, "url": url})

    @_ensure_client
    async def current_tab(self, sessionId: str):
        return await self._structured("currentTab", {"sessionId": sessionId})

    # DOM tools
    @_ensure_client
    async def get_dom_tree(self, sessionId: str):
        return await self._structured("getDomTree", {"sessionId": sessionId})

    @_ensure_client
    async def get_clickable(self, sessionId: str):
        return await self._structured("getClickable", {"sessionId": sessionId})

    @_ensure_client
    async def selector_map(self, sessionId: str, selector: str):
        return await self._structured("selectorMap", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def get_clean_text(self, sessionId: str):
        return await self._structured("getCleanText", {"sessionId": sessionId})

    # Element tools
    @_ensure_client
    async def find_all(self, sessionId: str, selector: str):
        return await self._structured("findAll", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def find_element(self, sessionId: str, selector: str):
        return await self._structured("findElement", {"sessionId": sessionId, "selector": selector})

    @_ensure_client
    async def find_element_by_xpath(self, sessionId: str, xpath: str):
        return await self._structured("findElementByXPath", {"sessionId": sessionId, "xpath": xpath})

    @_ensure_client
    async def find_element_xpath(self, sessionId: str, xpath: str):
        return await self.find_element_by_xpath(sessionId, xpath)

    @_ensure_client
    async def find_element_by_text(self, sessionId: str, text: str):
        return await self._structured("findElementByText", {"sessionId": sessionId, "text": text})

    @_ensure_client
    async def find_element_by_prompt(self, sessionId: str, prompt: str):
        return await self._structured("findElementByPrompt", {"sessionId": sessionId, "prompt": prompt})

    # Browser tools
    @_ensure_client
    async def create_browser(self, payload: Optional[Dict[str, Any]] = None):
        return await self._structured("createBrowser", payload or {})

    @_ensure_client
    async def release_browser(self, pod_name: str):
        return await self._structured("releaseBrowser", {"pod_name": pod_name})

    # Compatibility aliases
    @_ensure_client
    async def cookies(self, sessionId: str):
        return await self.get_cookies(sessionId)

    @_ensure_client
    async def current_url(self, sessionId: str):
        return await self.get_current_url(sessionId)

    @_ensure_client
    async def set_viewport(
        self,
        sessionId: str,
        *,
        width: int,
        height: int,
        deviceScaleFactor: float = 1.0,
        mobile: bool = False,
    ):
        return await self.viewport(
            sessionId,
            {
                "width": int(width),
                "height": int(height),
                "deviceScaleFactor": float(deviceScaleFactor),
                "mobile": bool(mobile),
            },
        )

    @_ensure_client
    async def get_viewport(self, sessionId: str):
        return await self.viewport(sessionId)

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

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

    @_ensure_client
    async def evaluate_stream_all(
        self,
        sessionId: str,
        expression: str,
        chunkSize: int = 100,
        max_items: Optional[int] = None,
    ):
        init = await self.evaluate_stream(sessionId, expression, chunkSize)
        stream_id = init.get("stream_id")
        if not stream_id:
            return []

        items = []
        offset = 0
        while True:
            chunk = await self.stream_pull(stream_id, offset, chunkSize)
            items.extend(chunk.get("items", []))
            if max_items is not None and len(items) >= max_items:
                return items[:max_items]
            if not chunk.get("has_more"):
                break
            offset += chunkSize

        return items
