// Copyright (c) 2026 Storm Dragon
// SPDX-License-Identifier: GPL-3.0-or-later

package tentacle

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"git.stormux.org/storm/cthulhu/internal/tentaclewire"
	sdk "git.stormux.org/storm/cthulhu/sdk/go/tentacle"
)

func TestSessionSDKChild(t *testing.T) {
	if len(os.Args) < 2 || os.Args[len(os.Args)-1] != "session-sdk-child" {
		return
	}
	client := sdk.NewSessionClient()
	commands := []sdk.Command{
		{ID: "large", Title: "Large", Handle: func(ctx context.Context, args json.RawMessage) (any, error) { return strings.Repeat("x", 48*1024), nil }},
		{ID: "request", Title: "Request", Handle: func(ctx context.Context, args json.RawMessage) (any, error) {
			var request struct {
				Method string          `json:"method"`
				Params json.RawMessage `json:"params"`
			}
			if err := json.Unmarshal(args, &request); err != nil {
				return nil, err
			}
			var value json.RawMessage
			err := client.Call(ctx, request.Method, request.Params, &value)
			var failure *sdk.SessionError
			if errors.As(err, &failure) {
				return map[string]any{"code": failure.Code}, nil
			}
			return map[string]any{"value": value}, err
		}},
		{ID: "event", Title: "Event", Handle: func(ctx context.Context, args json.RawMessage) (any, error) {
			select {
			case event := <-client.Events():
				return event, nil
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}},
		{ID: "wait", Title: "Wait", Handle: func(ctx context.Context, args json.RawMessage) (any, error) { <-ctx.Done(); return nil, ctx.Err() }},
	}
	if sdk.ServeSession(os.Stdin, os.Stdout, commands, nil, client) != nil {
		os.Exit(1)
	}
	os.Exit(0)
}

func sessionSDKProgram(t *testing.T, language string) []string {
	t.Helper()
	if language == "go" {
		executable, err := os.Executable()
		if err != nil {
			t.Fatal(err)
		}
		return []string{executable, "-test.run=^TestSessionSDKChild$", "--", "session-sdk-child"}
	}
	python, err := exec.LookPath("python3")
	if err != nil {
		t.Fatal("Python is required for session SDK conformance")
	}
	peer, err := filepath.Abs("testdata/session_peer.py")
	if err != nil {
		t.Fatal(err)
	}
	return []string{python, "-B", peer}
}

func startSessionSDK(t *testing.T, manager *Manager, language, id string, extraGrants ...string) *Session {
	t.Helper()
	grants := append([]string{"output-observe", "remote-present", "clipboard-exchange"}, extraGrants...)
	session, err := manager.Start(t.Context(), Spec{ID: id, Program: sessionSDKProgram(t, language), SessionEvents: true, RemoteGrants: grants})
	if err != nil {
		t.Fatal(err)
	}
	wait_signal(t, session.Ready())
	if err := session.Activate(); err != nil {
		state, reason := session.Status()
		t.Fatalf("activate: %v, %s %v", err, state, reason)
	}
	return session
}

func invokeSessionSDK(t *testing.T, session *Session, command string, params any) json.RawMessage {
	t.Helper()
	data, err := json.Marshal(params)
	if err != nil {
		t.Fatal(err)
	}
	pending, err := session.Invoke(t.Context(), command, data)
	if err != nil {
		t.Fatal(err)
	}
	result := wait_result(t, pending)
	if result.Err != nil {
		t.Fatal(result.Err)
	}
	return result.Value
}

func remoteCall(t *testing.T, session *Session, method string, params any) (json.RawMessage, int) {
	t.Helper()
	raw := invokeSessionSDK(t, session, "request", map[string]any{"method": method, "params": params})
	var result struct {
		Value json.RawMessage `json:"value"`
		Code  int             `json:"code"`
	}
	if err := json.Unmarshal(raw, &result); err != nil {
		t.Fatal(err)
	}
	return result.Value, result.Code
}

func TestRemoteWorkflowEpochCleanup(t *testing.T) {
	manager, err := NewManager(DefaultLimits(), nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := manager.Shutdown(ctx); err != nil {
			t.Error(err)
		}
	})
	session := startSessionSDK(t, manager, "go", "workflow")
	first, err := session.AuthorizeRemote("controller", []string{"remote-present"})
	if err != nil {
		t.Fatal(err)
	}
	second, err := session.AuthorizeRemote("controller", []string{"remote-present"})
	if err != nil {
		t.Fatal(err)
	}
	session.RevokeRemoteEpoch(first.Epoch)
	raw, code := remoteCall(t, session, "remote.open", tentaclewire.RemoteOpen{Action: second.Action, Epoch: second.Epoch, Role: second.Role, Grants: second.Grants})
	var lease tentaclewire.RemoteLease
	if code != 0 || json.Unmarshal(raw, &lease) != nil {
		t.Fatal("old workflow revoked its replacement")
	}
	delivery, err := session.RemoteContext(lease, "remote-present")
	if err != nil {
		t.Fatal(err)
	}
	session.RevokeRemoteEpoch(second.Epoch)
	if delivery.Err() == nil {
		t.Fatal("workflow close retained its live lease")
	}
	// Pending actions also release exclusive ownership on close, before their
	// timeout and without requiring a remote.open/remote.close pair.
	pending, err := session.AuthorizeRemote("controller", []string{"remote-present"})
	if err != nil {
		t.Fatal(err)
	}
	session.RevokeRemoteEpoch(pending.Epoch)
	other := startSessionSDK(t, manager, "go", "other-workflow")
	if _, err := other.AuthorizeRemote("controlled", []string{"output-observe"}); err != nil {
		t.Fatal("closed setup retained exclusive ownership", err)
	}
}

func openRemote(t *testing.T, session *Session) tentaclewire.RemoteLease {
	t.Helper()
	action, err := session.AuthorizeRemote("controlled", []string{"output-observe"})
	if err != nil {
		t.Fatal(err)
	}
	raw := invokeSessionSDK(t, session, "event", nil)
	var event tentaclewire.SessionEvent
	if json.Unmarshal(raw, &event) != nil || event.Action != action.Action || event.Kind != "remote.action" {
		t.Fatal("lost local action", string(raw))
	}
	value, code := remoteCall(t, session, "remote.open", tentaclewire.RemoteOpen{Action: event.Action, Epoch: event.Epoch, Role: event.Role, Grants: event.Grants})
	var lease tentaclewire.RemoteLease
	if code != 0 || json.Unmarshal(value, &lease) != nil || lease.Lease == "" || lease.Epoch != event.Epoch {
		t.Fatal("open", code, string(value))
	}
	return lease
}

func TestSessionSDKLifecycleConformance(t *testing.T) {
	for _, language := range []string{"go", "python"} {
		t.Run(language, func(t *testing.T) {
			limits := DefaultLimits()
			limits.RequestTimeout = 5 * time.Second
			manager := test_manager(t, limits, nil)
			session := startSessionSDK(t, manager, language, "session-sdk")
			if _, code := remoteCall(t, session, "remote.open", tentaclewire.RemoteOpen{Action: strings.Repeat("a", 32), Epoch: 1, Role: "controlled", Grants: []string{"output-observe"}}); code != -32001 {
				t.Fatal("child minted consent", code)
			}
			if _, err := session.AuthorizeRemote("controller", []string{"keyboard-capture"}); !errors.Is(err, ErrPermission) {
				t.Fatal("ungranted approval", err)
			}
			if _, err := session.AuthorizeRemote("controller", []string{"output-observe"}); !errors.Is(err, ErrPermission) {
				t.Fatal("wrong role", err)
			}
			lease := openRemote(t, session)
			delivery, err := session.RemoteContext(lease, "output-observe")
			if err != nil {
				t.Fatal(err)
			}
			if _, err := session.RemoteContext(lease, "clipboard-exchange"); !errors.Is(err, ErrPermission) {
				t.Fatal("install grant became session grant")
			}
			if _, code := remoteCall(t, session, "remote.key", map[string]any{}); code != -32602 {
				t.Fatal("malformed typed key operation accepted", code)
			}
			if _, code := remoteCall(t, session, "remote.braille", map[string]any{}); code != -32602 {
				t.Fatal("malformed typed Braille operation accepted", code)
			}
			if _, code := remoteCall(t, session, "remote.mouse", map[string]any{}); code != -32601 {
				t.Fatal("unimplemented desktop operation accepted", code)
			}
			// One blocked handler cannot block the endpoint receive loop or lifecycle calls.
			waiting, err := session.Invoke(t.Context(), "wait", nil)
			if err != nil {
				t.Fatal(err)
			}
			heartbeat := tentaclewire.RemoteOperation{Lease: lease.Lease, Epoch: lease.Epoch, Sequence: 1}
			if _, code := remoteCall(t, session, "remote.heartbeat", heartbeat); code != 0 {
				t.Fatal("heartbeat", code)
			}
			if _, code := remoteCall(t, session, "remote.heartbeat", heartbeat); code != -32001 {
				t.Fatal("duplicate operation accepted")
			}
			heartbeat.Sequence++
			if _, code := remoteCall(t, session, "remote.close", heartbeat); code != 0 {
				t.Fatal("close", code)
			}
			if delivery.Err() == nil {
				t.Fatal("buffered delivery survived close")
			}
			var event tentaclewire.SessionEvent
			raw := invokeSessionSDK(t, session, "event", nil)
			if json.Unmarshal(raw, &event) != nil || event.Kind != "remote.revoked" || event.Reason != "closed" {
				t.Fatal("missing revocation", string(raw))
			}
			if _, code := remoteCall(t, session, "remote.close", heartbeat); code != 0 {
				t.Fatal("close not idempotent")
			}
			newLease := openRemote(t, session)
			if newLease.Epoch <= lease.Epoch {
				t.Fatal("epoch reused")
			}
			if _, code := remoteCall(t, session, "remote.heartbeat", heartbeat); code != -32001 {
				t.Fatal("stale heartbeat accepted")
			}
			if _, code := remoteCall(t, session, "remote.close", heartbeat); code != 0 {
				t.Fatal("stale close failed")
			}
			current, err := session.RemoteContext(newLease, "output-observe")
			if err != nil {
				t.Fatal("stale close revoked new lease")
			}
			session.Disable()
			if current.Err() == nil {
				t.Fatal("disable left lease live")
			}
			if result := wait_result(t, waiting); result.Err == nil {
				t.Fatal("blocked handler survived disable")
			}
			wait_signal(t, session.Done())
		})
	}
}

func TestSessionLeaseExpiryAndExclusiveOwner(t *testing.T) {
	limits := DefaultLimits()
	limits.RequestTimeout = 5 * time.Second
	manager := test_manager(t, limits, nil)
	first := startSessionSDK(t, manager, "go", "first")
	second := startSessionSDK(t, manager, "python", "second")
	lease := openRemote(t, first)
	delivery, err := first.RemoteContext(lease, "output-observe")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := second.AuthorizeRemote("controlled", []string{"output-observe"}); !errors.Is(err, ErrBusy) {
		t.Fatal("two owners", err)
	}
	select {
	case <-delivery.Done():
	case <-time.After(RemoteHeartbeatTimeout + time.Second):
		t.Fatal("heartbeat did not expire")
	}
	if _, code := remoteCall(t, first, "remote.heartbeat", tentaclewire.RemoteOperation{Lease: lease.Lease, Epoch: lease.Epoch, Sequence: 1}); code != -32001 {
		t.Fatal("expired lease revived")
	}
	openRemote(t, second)
}

func TestSessionSDKCallTimeout(t *testing.T) {
	for _, blockedWriter := range []bool{false, true} {
		for _, language := range []string{"go", "python"} {
			t.Run(language, func(t *testing.T) {
				program := sessionSDKProgram(t, language)
				ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
				defer cancel()
				command := exec.CommandContext(ctx, program[0], program[1:]...)
				input, err := command.StdinPipe()
				if err != nil {
					t.Fatal(err)
				}
				output, err := command.StdoutPipe()
				if err != nil {
					t.Fatal(err)
				}
				if err := command.Start(); err != nil {
					t.Fatal(err)
				}
				defer input.Close()
				writer := json.NewEncoder(input)
				reader := bufio.NewReader(output)
				writer.Encode(map[string]any{"jsonrpc": "2.0", "id": "1", "method": "initialize", "params": Initialization{Protocol: Version{1, 1}, Session: "test", APIs: []string{"commands", "session-events"}, Capabilities: []string{}}})
				if _, err := read_frame(reader); err != nil {
					t.Fatal(err)
				}
				writer.Encode(map[string]any{"jsonrpc": "2.0", "method": "initialized", "params": map[string]string{"session": "test"}})
				requestID := "2"
				if blockedWriter {
					writer.Encode(map[string]any{"jsonrpc": "2.0", "id": "2", "method": "command.invoke", "params": map[string]any{"command": "large", "arguments": nil}})
					requestID = "3"
				}
				writer.Encode(map[string]any{"jsonrpc": "2.0", "id": requestID, "method": "command.invoke", "params": map[string]any{"command": "request", "arguments": map[string]any{"method": "remote.heartbeat", "params": map[string]any{}}}})
				if !blockedWriter {
					frame, err := read_frame(reader)
					if err != nil || !strings.Contains(string(frame), `"id":"c:1"`) {
						t.Fatal("missing child request", string(frame), err)
					}
				}
				// Keep stdin open and never reply: SDK must wake a blocked receive loop.
				_ = command.Wait()
				if ctx.Err() != nil {
					t.Fatal("SDK stayed blocked after request timeout")
				}
			})
		}
	}

}

func TestSessionQueueOverflowRevokesLease(t *testing.T) {
	manager := test_manager(t, DefaultLimits(), nil)
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	output, revoke := context.WithCancel(ctx)
	session := &Session{manager: manager, spec: Spec{SessionEvents: true, RemoteGrants: []string{"output-observe"}}, ctx: ctx, cancel: cancel, outputContext: output, revokeOutput: revoke, state: Active, outbound: make(chan outgoing, 1), revoked: make(chan struct{})}
	action, err := session.AuthorizeRemote("controlled", []string{"output-observe"})
	if err != nil {
		t.Fatal(err)
	}
	params, _ := json.Marshal(tentaclewire.RemoteOpen{Action: action.Action, Epoch: action.Epoch, Role: action.Role, Grants: action.Grants})
	session.mutex.Lock()
	value, failure := session.remote_call_locked("remote.open", params)
	session.mutex.Unlock()
	if failure != nil {
		t.Fatal(failure)
	}
	delivery, err := session.RemoteContext(value.(tentaclewire.RemoteLease), "output-observe")
	if err != nil {
		t.Fatal(err)
	}
	session.RevokeRemote() // Revocation cannot be queued behind the unread action.
	if delivery.Err() == nil {
		t.Fatal("queue blocked immediate revocation")
	}
	if state, _ := session.Status(); state != Failed {
		t.Fatal("overflow did not retire session", state)
	}
}

func TestSessionHostRejectsMalformedChildRequests(t *testing.T) {
	for _, frame := range []string{
		`{"jsonrpc":"2.0","id":"1","method":"remote.close","params":{}}`,
		`{"jsonrpc":"2.0","id":"c:01","method":"remote.close","params":{}}`,
		`{"jsonrpc":"2.0","id":"c:0","method":"remote.close","params":{}}`,
		`{"jsonrpc":"2.0","id":null,"method":"remote.close","params":{}}`,
		`{"jsonrpc":"2.0","id":"c:1","method":"remote.close","method":"remote.open","params":{}}`,
		`{"jsonrpc":"2.0","id":"c:1","method":"remote.close","params":{},"result":{}}`,
		`{"jsonrpc":"2.0","id":"c:1","method":"remote.close"}`,
	} {
		t.Run(frame, func(t *testing.T) {
			session := &Session{spec: Spec{SessionEvents: true}, ctx: t.Context(), state: Active, outbound: make(chan outgoing, 32)}
			if err := session.accept_child_request([]byte(frame)); err == nil {
				t.Fatal("malformed request accepted")
			}
		})
	}
	session := &Session{spec: Spec{SessionEvents: true}, ctx: t.Context(), state: Active, outbound: make(chan outgoing, 32)}
	frame := []byte(`{"jsonrpc":"2.0","id":"c:1","method":"remote.unknown","params":{}}`)
	if err := session.accept_child_request(frame); err != nil {
		t.Fatal(err)
	}
	reply, err := decode_response((<-session.outbound).frame)
	if err != nil || reply.Error == nil || reply.Error.Code != -32601 {
		t.Fatal("unknown method response", reply, err)
	}
	if err := session.accept_child_request(frame); err == nil {
		t.Fatal("duplicate child ID accepted")
	}
	session.spec.SessionEvents = false
	if err := session.accept_child_request([]byte(`{"jsonrpc":"2.0","id":"c:2","method":"remote.close","params":{}}`)); err == nil {
		t.Fatal("1.0 accepted child request")
	}
}

func TestSessionSDKRejectsMalformedHostMessages(t *testing.T) {
	initialize := `{"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocol":{"major":1,"minor":1},"session":"test","apis":["commands","session-events"],"capabilities":[]}}` + "\n"
	activate := `{"jsonrpc":"2.0","method":"initialized","params":{"session":"test"}}` + "\n"
	event := `{"jsonrpc":"2.0","method":"session.event","params":{"sequence":1,"kind":"remote.action","action":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa","epoch":1,"role":"controlled","grants":["output-observe"]}}` + "\n"
	cases := map[string]string{
		"null capabilities":       strings.Replace(initialize, `"capabilities":[]`, `"capabilities":null`, 1),
		"null event field":        initialize + activate + strings.Replace(event, `"epoch":1`, `"epoch":1,"lease":null`, 1),
		"old host":                strings.Replace(initialize, `"minor":1`, `"minor":0`, 1),
		"missing event API":       strings.Replace(initialize, `"commands","session-events"`, `"commands"`, 1),
		"legacy capabilities":     strings.Replace(initialize, `"capabilities":[]`, `"capabilities":["commands"]`, 1),
		"unrequested response":    initialize + activate + `{"jsonrpc":"2.0","id":"c:1","result":null}` + "\n",
		"event before activation": initialize + event,
		"event with request ID":   initialize + activate + strings.Replace(event, `"method":`, `"id":"2","method":`, 1),
		"duplicate event":         initialize + activate + event + event,
		"skipped sequence":        initialize + activate + strings.Replace(event, `"sequence":1`, `"sequence":2`, 1),
		"event role mismatch":     initialize + activate + strings.Replace(event, `"controlled"`, `"controller"`, 1),
		"unknown event":           initialize + activate + strings.Replace(event, `"remote.action"`, `"remote.execute"`, 1),
	}
	overflow := initialize + activate
	for sequence := 1; sequence <= 33; sequence++ {
		overflow += strings.Replace(event, `"sequence":1`, `"sequence":`+strconv.Itoa(sequence), 1)
	}
	cases["event overflow"] = overflow
	outputEvent := `{"jsonrpc":"2.0","method":"session.event","params":{"sequence":1,"kind":"remote.output","lease":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","epoch":1,"speech":{"kind":"cancel"}}}` + "\n"
	cases["output without lease"] = initialize + activate + strings.Replace(outputEvent, `"lease":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",`, "", 1)
	cases["output with role"] = initialize + activate + strings.Replace(outputEvent, `"epoch":1`, `"epoch":1,"role":"controller"`, 1)
	cases["null speech"] = initialize + activate + strings.Replace(outputEvent, `{"kind":"cancel"}`, `null`, 1)
	cases["malformed speech"] = initialize + activate + strings.Replace(outputEvent, `{"kind":"cancel"}`, `{"kind":"pause","paused":null}`, 1)
	cases["speech on action"] = initialize + activate + strings.Replace(event, `"epoch":1`, `"epoch":1,"speech":{"kind":"cancel"}`, 1)
	for _, language := range []string{"go", "python"} {
		for name, input := range cases {
			t.Run(language+"/"+name, func(t *testing.T) {
				program := sessionSDKProgram(t, language)
				ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
				defer cancel()
				command := exec.CommandContext(ctx, program[0], program[1:]...)
				command.Stdin = strings.NewReader(input)
				if err := command.Run(); err == nil {
					t.Fatal("invalid host message accepted")
				}
				if ctx.Err() != nil {
					t.Fatal("SDK hung on invalid message")
				}
			})
		}
	}
}

func TestSessionAggregateQueueBound(t *testing.T) {
	session := &Session{spec: Spec{SessionEvents: true}, outbound: make(chan outgoing, 32)}
	frame := make([]byte, MaxFrameBytes)
	for range 4 {
		if !session.enqueue_locked(outgoing{frame: frame}) {
			t.Fatal("queue rejected boundary")
		}
	}
	if session.enqueue_locked(outgoing{frame: []byte{1}}) {
		t.Fatal("aggregate byte ceiling exceeded")
	}
}

func TestPendingRemoteRevocationReachesBothSDKs(t *testing.T) {
	for _, language := range []string{"go", "python"} {
		for _, reason := range []string{"local", "expired", "replaced"} {
			t.Run(language+"/"+reason, func(t *testing.T) {
				session := startSessionSDK(t, test_manager(t, DefaultLimits(), nil), language, "pending-revoke")
				action, err := session.AuthorizeRemote("controller", []string{"remote-present"})
				if err != nil {
					t.Fatal(err)
				}
				invokeSessionSDK(t, session, "event", nil)
				session.mutex.Lock()
				if reason == "expired" {
					session.refresh_remote_locked(time.Millisecond)
				} else {
					session.revoke_remote_locked(reason, true)
				}
				session.mutex.Unlock()
				var event tentaclewire.SessionEvent
				raw := invokeSessionSDK(t, session, "event", nil)
				if json.Unmarshal(raw, &event) != nil || event.Kind != "remote.revoked" || event.Action != action.Action || event.Lease != "" || event.Epoch != action.Epoch || event.Reason != reason {
					t.Fatal("pending action revocation was not delivered", string(raw))
				}
				if _, code := remoteCall(t, session, "remote.open", tentaclewire.RemoteOpen{Action: action.Action, Epoch: action.Epoch, Role: action.Role, Grants: action.Grants}); code != -32001 {
					t.Fatal("retired action opened a lease", code)
				}
			})
		}
	}
}