// ARTEMIS adapter for the pinned MIT Cua Driver. No model, scheduler, or recorder.
package main

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	_ "embed"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

//go:embed tools.json
var toolsJSON []byte

const adapterVersion = "0.36.0+artemis.1"

type obj = map[string]any
type request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Method  string          `json:"method"`
	Params  obj             `json:"params"`
}
type peer struct {
	cmd   *exec.Cmd
	input io.WriteCloser
	lines chan []byte
	next  int
	mu    sync.Mutex
	dead  bool
	job   uintptr
}

func object(v any) obj {
	m, _ := v.(map[string]any)
	if m == nil {
		return obj{}
	}
	return m
}
func str(v any) string { s, _ := v.(string); return s }
func num(v any) int64 {
	switch n := v.(type) {
	case float64:
		return int64(n)
	case int64:
		return n
	case int:
		return int64(n)
	}
	return 0
}
func items(v any) []any      { a, _ := v.([]any); return a }
func textBlock(s string) obj { return obj{"type": "text", "text": s} }
func receipt(execution, side, code string) obj {
	return obj{"version": 1, "execution": execution, "side_effect": side, "retryable": false, "error_code": code}
}
func failure(msg, execution, side, code string) obj {
	return obj{"isError": true, "content": []any{textBlock(msg)}, "structuredContent": obj{"artemis": receipt(execution, side, code)}}
}
func writeTool(name string) bool {
	return name != "list_windows" && name != "get_window_state" && name != "verify_state"
}

func startPeer(exe string) (*peer, error) {
	c := exec.Command(exe, "mcp", "--direct")
	c.Env = append(os.Environ(), "DO_NOT_TRACK=1", "CUA_TELEMETRY=0", "CUA_DRIVER_RS_TELEMETRY_ENABLED=false")
	configureProcess(c)
	in, err := c.StdinPipe()
	if err != nil {
		return nil, err
	}
	out, err := c.StdoutPipe()
	if err != nil {
		return nil, err
	}
	c.Stderr = os.Stderr
	if err = c.Start(); err != nil {
		return nil, err
	}
	job, err := ownProcess(c.Process.Pid)
	if err != nil {
		_ = c.Process.Kill()
		_ = c.Wait()
		return nil, err
	}
	p := &peer{cmd: c, input: in, lines: make(chan []byte, 32), job: job}
	go func() {
		defer close(p.lines)
		s := bufio.NewScanner(out)
		s.Buffer(make([]byte, 65536), 32<<20)
		for s.Scan() {
			p.lines <- append([]byte(nil), s.Bytes()...)
		}
	}()
	_, err = p.call("initialize", obj{"protocolVersion": "2024-11-05", "clientInfo": obj{"name": "artemis-cua-adapter", "version": adapterVersion}, "capabilities": obj{}})
	if err != nil {
		p.stop()
		return nil, err
	}
	err = json.NewEncoder(in).Encode(obj{"jsonrpc": "2.0", "method": "notifications/initialized"})
	if err != nil {
		p.stop()
		return nil, err
	}
	return p, nil
}
func (p *peer) stop() {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.dead {
		return
	}
	p.dead = true
	_ = p.input.Close()
	closeJob(p.job)
	_ = p.cmd.Process.Kill()
	_ = p.cmd.Wait()
}
func (p *peer) call(method string, params obj) (obj, error) {
	p.mu.Lock()
	dead := p.dead
	p.mu.Unlock()
	if dead {
		return nil, errors.New("driver closed")
	}
	p.next++
	id := p.next
	if err := json.NewEncoder(p.input).Encode(obj{"jsonrpc": "2.0", "id": id, "method": method, "params": params}); err != nil {
		return nil, err
	}
	timer := time.NewTimer(45 * time.Second)
	defer timer.Stop()
	for {
		select {
		case raw, ok := <-p.lines:
			if !ok {
				return nil, errors.New("driver disconnected")
			}
			var r obj
			if json.Unmarshal(raw, &r) != nil {
				continue
			}
			if num(r["id"]) != int64(id) {
				continue
			}
			if r["error"] != nil {
				return nil, errors.New("driver protocol error")
			}
			return object(r["result"]), nil
		case <-timer.C:
			p.stop()
			return nil, errors.New("driver response timed out; action not replayed")
		}
	}
}
func (p *peer) tool(name string, args obj) (obj, error) {
	return p.call("tools/call", obj{"name": name, "arguments": args})
}

type adapter struct {
	peer    *peer
	tools   []any
	schemas map[string]obj
}

func newAdapter(p *peer) *adapter {
	a := &adapter{peer: p, schemas: map[string]obj{}}
	_ = json.Unmarshal(toolsJSON, &a.tools)
	for _, v := range a.tools {
		t := object(v)
		a.schemas[str(t["name"])] = object(t["inputSchema"])
	}
	return a
}
func (a *adapter) validate(name string, args obj) error {
	schema, ok := a.schemas[name]
	if !ok {
		return errors.New("tool is not exposed")
	}
	props := object(schema["properties"])
	for key := range args {
		if _, ok := props[key]; !ok {
			return fmt.Errorf("unsupported argument: %s", key)
		}
	}
	for _, key := range items(schema["required"]) {
		if _, ok := args[str(key)]; !ok {
			return fmt.Errorf("missing argument: %s", str(key))
		}
	}
	if name == "list_windows" {
		_, err := windowQuery(args)
		return err
	}
	if num(args["pid"]) <= 0 || num(args["window_id"]) <= 0 {
		return errors.New("a concrete pid and window_id are required")
	}
	if name == "save_screenshot" && !validFilename(str(args["filename"])) {
		return errors.New("filename must be a new plain PNG filename")
	}
	if name == "get_window_state" && (args["include_screenshot"] == false || args["include_accessibility_tree"] == false) {
		return errors.New("candidate observations require both screenshot and accessibility tree")
	}
	return nil
}
func validFilename(s string) bool {
	return s != "" && len(s) <= 160 && strings.EqualFold(filepath.Ext(s), ".png") && !strings.ContainsAny(s, "/\\:<>\"|?*\x00") && s != "." && s != ".." && !strings.HasSuffix(s, ".") && !strings.HasSuffix(s, " ")
}
func actionReceipt(r obj, write bool) obj {
	if !write {
		if r["isError"] == true {
			return receipt("not_started", "none", "observation_failed")
		}
		return receipt("completed", "none", "")
	}
	s := object(r["structuredContent"])
	// v0.36.0 class-level preflight refusal has no effect field. It occurs
	// before input injection. Do not classify actuator failures with a cause
	// as safe-to-repeat; retain their uncertain execution below.
	if s["code"] == "background_unavailable" && s["cause"] == nil && s["event_kind"] != nil && r["isError"] == true {
		return receipt("not_started", "none", "background_unavailable")
	}
	switch str(s["effect"]) {
	case "confirmed":
		return receipt("completed", "confirmed", "")
	case "partial":
		return receipt("partial", "confirmed", "partial_delivery")
	case "refused":
		return receipt("not_started", "none", "driver_refused")
	case "unverifiable", "suspected_noop":
		return receipt("completed", "possible", str(s["effect"]))
	default:
		return receipt("unknown", "possible", "unclassified_result")
	}
}
func imageSize(r obj) (int, int) {
	for _, v := range items(r["content"]) {
		c := object(v)
		if c["type"] == "image" {
			raw, err := base64.StdEncoding.DecodeString(str(c["data"]))
			if err == nil {
				cfg, _, err := image.DecodeConfig(bytes.NewReader(raw))
				if err == nil {
					return cfg.Width, cfg.Height
				}
			}
		}
	}
	return 0, 0
}
func source(r obj, args obj, started float64) obj {
	s := object(r["structuredContent"])
	w, h := imageSize(r)
	if r["isError"] == true || w <= 0 || h <= 0 || num(s["pid"]) != num(args["pid"]) || num(s["window_id"]) != num(args["window_id"]) {
		return nil
	}
	bounds := object(s["window_bounds"])
	rect := obj{"left": num(bounds["x"]), "top": num(bounds["y"]), "width": num(bounds["width"]), "height": num(bounds["height"])}
	screen, ok := verifyWindow(num(args["pid"]), num(args["window_id"]), rect)
	if !ok {
		return nil
	}
	return obj{"kind": "computer", "target_id": fmt.Sprint(num(args["window_id"])), "label": s["window_title"], "captured_at": started, "width": w, "height": h, "window": rect, "screen": screen}
}
func (a *adapter) observe(args obj) (obj, obj, error) {
	started := float64(time.Now().UnixMilli()) / 1000
	r, err := a.peer.tool("get_window_state", args)
	if err != nil {
		return nil, nil, err
	}
	return r, source(r, args, started), nil
}
func observationArgs(args obj) obj {
	return obj{"pid": args["pid"], "window_id": args["window_id"], "max_image_dimension": 1280, "max_elements": 250, "tree_format": "markdown"}
}
func (a *adapter) handle(name string, args obj) obj {
	if err := a.validate(name, args); err != nil {
		return failure(err.Error(), "not_started", "none", "invalid_input")
	}
	if name != "list_windows" {
		if !windowOwned(num(args["pid"]), num(args["window_id"])) {
			return failure("window is gone or belongs to a different process", "not_started", "none", "invalid_target")
		}
	}
	if name == "save_screenshot" {
		return a.save(args)
	}
	var r, src obj
	var err error
	if name == "get_window_state" {
		for k, v := range observationArgs(args) {
			if _, ok := args[k]; !ok {
				args[k] = v
			}
		}
		r, src, err = a.observe(args)
	} else if name == "list_windows" {
		query, _ := windowQuery(args) // Already validated above.
		r, err = a.peer.tool(name, query)
		if err == nil {
			r = windowListing(r, query)
		}
	} else {
		upstream := name
		if name == "click_at" {
			upstream = "click"
		}
		r, err = a.peer.tool(upstream, args)
	}
	if err != nil {
		a.peer.stop()
		if writeTool(name) {
			return failure(err.Error(), "unknown", "possible", "transport_lost")
		}
		return failure(err.Error(), "not_started", "none", "transport_lost")
	}
	rec := actionReceipt(r, writeTool(name))
	structured := object(r["structuredContent"])
	if writeTool(name) && rec["execution"] != "not_started" {
		// Observe once after delivery; NEVER repeat an uncertain write.
		observation, origin, e := a.observe(observationArgs(args))
		if e == nil && origin != nil {
			src = origin
			structured["observation"] = observation["structuredContent"]
			r["content"] = append(items(r["content"]), items(observation["content"])...)
		} else {
			structured["observation_error"] = "Post-action evidence unavailable; inspect before another action."
		}
	}
	if name == "verify_state" { // Read-only verification may include a screenshot; it is not a task verdict.
		src = source(r, args, float64(time.Now().UnixMilli())/1000)
	}
	if src != nil {
		rec["source"] = src
	} else if name == "get_window_state" && r["isError"] != true {
		return failure("capture target or geometry could not be validated", "not_started", "none", "invalid_capture")
	}
	if structured["effect"] == "refused" {
		r["isError"] = true
	}
	structured["artemis"] = rec
	r["structuredContent"] = structured
	return r
}
func (a *adapter) save(args obj) obj {
	root := os.Getenv("ARTEMIS_ARTIFACTS_DIR")
	if root == "" {
		return failure("task artifact directory unavailable", "not_started", "none", "missing_artifacts")
	}
	r, src, err := a.observe(observationArgs(args))
	if err != nil || src == nil {
		return failure("fresh window capture failed", "not_started", "none", "capture_failed")
	}
	var png []byte
	for _, v := range items(r["content"]) {
		b := object(v)
		if b["type"] == "image" && b["mimeType"] == "image/png" {
			png, _ = base64.StdEncoding.DecodeString(str(b["data"]))
			break
		}
	}
	if len(png) == 0 {
		return failure("expected PNG capture", "not_started", "none", "capture_format")
	}
	if err = os.MkdirAll(root, 0700); err != nil {
		return failure("artifact directory unavailable", "not_started", "none", "file_error")
	}
	path := filepath.Join(root, str(args["filename"]))
	f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return failure("file exists or cannot be created", "not_started", "none", "file_error")
	}
	n, writeErr := f.Write(png)
	closeErr := f.Close()
	actual, readErr := os.ReadFile(path)
	if writeErr != nil || closeErr != nil || readErr != nil || n != len(png) || !bytes.Equal(actual, png) {
		return failure("artifact write could not be verified", "partial", "possible", "file_error")
	}
	rec := receipt("completed", "confirmed", "")
	rec["source"] = src
	rec["attachments"] = []any{obj{"name": str(args["filename"]), "sha256": fmt.Sprintf("%x", sha256.Sum256(actual)), "media_type": "image/png"}}
	object(r["structuredContent"])["artemis"] = rec
	r["content"] = append(items(r["content"]), textBlock("Saved verified task attachment: "+str(args["filename"])))
	return r
}
func serve(a *adapter, in io.Reader, out io.Writer) {
	input := make(chan request)
	go func() {
		defer close(input)
		defer a.peer.stop()
		s := bufio.NewScanner(in)
		s.Buffer(make([]byte, 65536), 2<<20)
		for s.Scan() {
			var q request
			if json.Unmarshal(s.Bytes(), &q) == nil {
				if q.Method == "notifications/cancelled" {
					a.peer.stop()
					continue
				}
				input <- q
			}
		}
	}()
	enc := json.NewEncoder(out)
	for q := range input {
		if len(q.ID) == 0 {
			continue
		}
		var result any
		var rpcErr any
		switch q.Method {
		case "initialize":
			result = obj{"protocolVersion": "2024-11-05", "capabilities": obj{"tools": obj{}}, "serverInfo": obj{"name": "artemis-cua-driver", "version": adapterVersion}}
		case "ping":
			result = obj{}
		case "tools/list":
			result = obj{"tools": a.tools}
		case "tools/call":
			result = a.handle(str(q.Params["name"]), object(q.Params["arguments"]))
		default:
			rpcErr = obj{"code": -32601, "message": "method not found"}
		}
		response := obj{"jsonrpc": "2.0", "id": q.ID}
		if rpcErr != nil {
			response["error"] = rpcErr
		} else {
			response["result"] = result
		}
		if enc.Encode(response) != nil {
			return
		}
	}
}
func main() {
	exe, err := os.Executable()
	if err != nil {
		os.Exit(1)
	}
	p, err := startPeer(filepath.Join(filepath.Dir(exe), "cua-driver.exe"))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Cua startup failed:", err)
		os.Exit(1)
	}
	defer p.stop()
	serve(newAdapter(p), os.Stdin, os.Stdout)
}