package main

import (
	"context"
	"log"

	"github.com/electricbubble/gadb"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

type runAdbInput struct {
	Command      string `json:"command" jsonschema:"the adb shell command to run on the device"`
	DeviceSerial string `json:"deviceSerial,omitempty" jsonschema:"optional device serial; uses first connected device if empty"`
}

func runAdb(ctx context.Context, req *mcp.CallToolRequest, input runAdbInput) (*mcp.CallToolResult, any, error) {
	adbClient, err := gadb.NewClient()
	if err != nil {
		return &mcp.CallToolResult{
			IsError: true,
			Content: []mcp.Content{
				&mcp.TextContent{Text: "Failed to connect to ADB server: " + err.Error()},
			},
		}, nil, nil
	}

	devices, err := adbClient.DeviceList()
	if err != nil {
		return &mcp.CallToolResult{
			IsError: true,
			Content: []mcp.Content{
				&mcp.TextContent{Text: "Failed to list devices: " + err.Error()},
			},
		}, nil, nil
	}

	if len(devices) == 0 {
		return &mcp.CallToolResult{
			IsError: true,
			Content: []mcp.Content{
				&mcp.TextContent{Text: "No devices connected"},
			},
		}, nil, nil
	}

	var dev gadb.Device
	if input.DeviceSerial != "" {
		found := false
		for _, d := range devices {
			if d.Serial() == input.DeviceSerial {
				dev = d
				found = true
				break
			}
		}
		if !found {
			return &mcp.CallToolResult{
				IsError: true,
				Content: []mcp.Content{
					&mcp.TextContent{Text: "Device not found: " + input.DeviceSerial},
				},
			}, nil, nil
		}
	} else {
		dev = devices[0]
	}

	output, err := dev.RunShellCommand(input.Command)
	if err != nil {
		return &mcp.CallToolResult{
			IsError: true,
			Content: []mcp.Content{
				&mcp.TextContent{Text: "Command failed: " + err.Error()},
			},
		}, nil, nil
	}

	return &mcp.CallToolResult{
		Content: []mcp.Content{
			&mcp.TextContent{Text: output},
		},
	}, nil, nil
}

func main() {
	server := mcp.NewServer(&mcp.Implementation{
		Name:    "cariddi-android-adb",
		Version: "v1.0.0",
	}, nil)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "runAdb",
		Description: "Run an adb shell command on a connected Android device and return the output. If deviceSerial is omitted, uses the first connected device.",
	}, runAdb)

	if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
		log.Fatal(err)
	}
}