#!/usr/bin/env python3
import asyncio, base64, os, subprocess
from typing import Any
import mcp.types as types
from mcp.server import NotificationOptions, Server
from mcp.server.sse import SseServerTransport
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response
from starlette.routing import Mount, Route
import uvicorn

server = Server("android-mcp")

def adb(cmd, timeout=30):
    try:
        r = subprocess.run(["bash","-c",f"adb {cmd}"], capture_output=True, text=True, timeout=timeout)
        return ((r.stdout or "")+(r.stderr or "")).rstrip() or "(no output)"
    except subprocess.TimeoutExpired:
        return f"[TIMEOUT {timeout}s]"
    except Exception as e:
        return f"[ERROR] {e}"

def ok(t): return [types.TextContent(type="text", text=t)]

@server.list_tools()
async def list_tools():
    return [
        types.Tool(name="android_shell", description="Run adb shell command on Android.",
            inputSchema={"type":"object","properties":{"command":{"type":"string"}},"required":["command"]}),
        types.Tool(name="android_screenshot", description="Take Android screenshot (base64 PNG).",
            inputSchema={"type":"object","properties":{},"required":[]}),
        types.Tool(name="android_tap", description="Tap at X,Y on Android screen.",
            inputSchema={"type":"object","properties":{"x":{"type":"integer"},"y":{"type":"integer"}},"required":["x","y"]}),
        types.Tool(name="android_swipe", description="Swipe on Android screen.",
            inputSchema={"type":"object","properties":{"x1":{"type":"integer"},"y1":{"type":"integer"},"x2":{"type":"integer"},"y2":{"type":"integer"},"duration":{"type":"integer","default":300}},"required":["x1","y1","x2","y2"]}),
        types.Tool(name="android_type", description="Type text on Android.",
            inputSchema={"type":"object","properties":{"text":{"type":"string"}},"required":["text"]}),
        types.Tool(name="android_key", description="Press Android key (e.g. KEYCODE_HOME, KEYCODE_BACK).",
            inputSchema={"type":"object","properties":{"keycode":{"type":"string"}},"required":["keycode"]}),
        types.Tool(name="android_install_apk", description="Install APK from URL.",
            inputSchema={"type":"object","properties":{"url":{"type":"string"}},"required":["url"]}),
        types.Tool(name="android_status", description="Android emulator status.",
            inputSchema={"type":"object","properties":{},"required":[]}),
    ]

@server.call_tool()
async def call_tool(name, arguments):
    if name == "android_shell":
        return ok(adb(f"shell {arguments['command']}"))
    if name == "android_screenshot":
        adb("shell screencap -p /sdcard/screen.png")
        adb("pull /sdcard/screen.png /tmp/android_screen.png")
        try:
            import pathlib
            b64 = base64.b64encode(pathlib.Path("/tmp/android_screen.png").read_bytes()).decode()
            return ok(f"data:image/png;base64,{b64}")
        except Exception as e:
            return ok(f"[ERROR] {e}")
    if name == "android_tap":
        return ok(adb(f"shell input tap {arguments['x']} {arguments['y']}"))
    if name == "android_swipe":
        a = arguments
        return ok(adb(f"shell input swipe {a['x1']} {a['y1']} {a['x2']} {a['y2']} {a.get('duration',300)}"))
    if name == "android_type":
        text = arguments["text"].replace(" ", "%s").replace("'", "")
        return ok(adb(f"shell input text '{text}'"))
    if name == "android_key":
        return ok(adb(f"shell input keyevent {arguments['keycode']}"))
    if name == "android_install_apk":
        url = arguments["url"]
        subprocess.run(["bash","-c",f"wget -qO /tmp/app.apk {url}"])
        return ok(adb("install /tmp/app.apk", timeout=60))
    if name == "android_status":
        return ok(adb("devices") + "\n" + adb("shell getprop ro.build.version.release"))
    return ok(f"Unknown: {name}")

sse = SseServerTransport("/messages/")
async def handle_sse(request: Request):
    async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
        await server.run(streams[0], streams[1], server.create_initialization_options())

app = Starlette(routes=[
    Route("/sse", endpoint=handle_sse),
    Route("/health", endpoint=lambda r: Response("android-mcp OK")),
    Mount("/messages/", app=sse.handle_post_message),
])

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)