package handlers

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"testing"
	"time"
	"unsafe"

	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	ctxpkg "github.com/sandbox0-ai/sandbox0/manager/procd/pkg/context"
	"github.com/sandbox0-ai/sandbox0/manager/procd/pkg/process"
	"go.uber.org/zap"
)

type fakeProcess struct {
	mu           sync.RWMutex
	outputCh     chan process.ProcessOutput
	onWrite      func([]byte)
	finished     bool
	exitCode     int
	state        process.ProcessState
	processType  process.ProcessType
	stdout       string
	stderr       string
	exitHandlers []process.ExitHandler
	handlerAdded chan struct{}
	writeErr     error
}

func (f *fakeProcess) ID() string { return "proc-test" }
func (f *fakeProcess) Type() process.ProcessType {
	if f.processType == "" {
		return process.ProcessTypeREPL
	}
	return f.processType
}
func (f *fakeProcess) PID() int        { return 1 }
func (f *fakeProcess) Start() error    { return nil }
func (f *fakeProcess) Stop() error     { return nil }
func (f *fakeProcess) Restart() error  { return nil }
func (f *fakeProcess) IsRunning() bool { return true }
func (f *fakeProcess) IsFinished() bool {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.finished
}
func (f *fakeProcess) State() process.ProcessState {
	f.mu.RLock()
	defer f.mu.RUnlock()
	if f.state == "" {
		return process.ProcessStateRunning
	}
	return f.state
}
func (f *fakeProcess) AddStartHandler(process.StartHandler) {}
func (f *fakeProcess) AddExitHandler(handler process.ExitHandler) {
	f.mu.Lock()
	f.exitHandlers = append(f.exitHandlers, handler)
	added := f.handlerAdded
	f.mu.Unlock()
	if added != nil {
		select {
		case <-added:
		default:
			close(added)
		}
	}
}
func (f *fakeProcess) Pause() error   { return nil }
func (f *fakeProcess) Resume() error  { return nil }
func (f *fakeProcess) IsPaused() bool { return false }
func (f *fakeProcess) WriteInput(data []byte) error {
	if f.writeErr != nil {
		return f.writeErr
	}
	if f.onWrite != nil {
		f.onWrite(data)
	}
	return nil
}
func (f *fakeProcess) ReadOutput() <-chan process.ProcessOutput { return f.outputCh }
func (f *fakeProcess) GetOutput() (stdout, stderr string) {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.stdout, f.stderr
}
func (f *fakeProcess) ResizePTY(process.PTYSize) error { return nil }
func (f *fakeProcess) SendSignal(syscall.Signal) error { return nil }
func (f *fakeProcess) ExitCode() (int, error) {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.exitCode, nil
}
func (f *fakeProcess) ResourceUsage() process.ResourceUsage { return process.ResourceUsage{} }

func (f *fakeProcess) triggerExit(event process.ExitEvent) {
	f.mu.Lock()
	f.finished = true
	f.exitCode = event.ExitCode
	f.state = event.State
	handlers := append([]process.ExitHandler(nil), f.exitHandlers...)
	f.mu.Unlock()
	for _, handler := range handlers {
		if handler != nil {
			handler(event)
		}
	}
}

func attachContext(manager *ctxpkg.Manager, ctx *ctxpkg.Context) {
	managerValue := reflect.ValueOf(manager).Elem()
	contextsField := managerValue.FieldByName("contexts")
	contexts := reflect.NewAt(contextsField.Type(), unsafe.Pointer(contextsField.UnsafeAddr())).Elem()
	if contexts.IsNil() {
		contexts.Set(reflect.MakeMap(contexts.Type()))
	}
	contexts.SetMapIndex(reflect.ValueOf(ctx.ID), reflect.ValueOf(ctx))
}

func newHandlerWithContext(proc process.Process, processType process.ProcessType) (*ContextHandler, *ctxpkg.Context) {
	manager := ctxpkg.NewManager()
	ctx := &ctxpkg.Context{
		ID:          "ctx-test",
		Type:        processType,
		Alias:       "python",
		MainProcess: proc,
	}
	attachContext(manager, ctx)
	return NewContextHandler(manager, zap.NewNop()), ctx
}

func TestListContextsReturnsEmptyArray(t *testing.T) {
	handler := NewContextHandler(ctxpkg.NewManager(), zap.NewNop())
	req := httptest.NewRequest(http.MethodGet, "/contexts", nil)
	rec := httptest.NewRecorder()

	handler.List(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	var payload struct {
		Success bool `json:"success"`
		Data    struct {
			Contexts []ContextResponse `json:"contexts"`
		} `json:"data"`
	}
	if err := json.NewDecoder(rec.Body).Decode(&payload); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !payload.Success {
		t.Fatal("success = false, want true")
	}
	if payload.Data.Contexts == nil {
		t.Fatal("contexts slice is nil, want empty array")
	}
	if len(payload.Data.Contexts) != 0 {
		t.Fatalf("contexts = %d, want 0", len(payload.Data.Contexts))
	}
}

func TestCreateContextReturnsBadRequestForInvalidProcessConfig(t *testing.T) {
	tests := []struct {
		name string
		body string
	}{
		{name: "cmd missing command", body: `{"type":"cmd"}`},
		{name: "unsupported process type", body: `{"type":"wat"}`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			handler := NewContextHandler(ctxpkg.NewManager(), zap.NewNop())
			req := httptest.NewRequest(http.MethodPost, "/contexts", strings.NewReader(tt.body))
			rec := httptest.NewRecorder()

			handler.Create(rec, req)

			if rec.Code != http.StatusBadRequest {
				t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
			}
			var resp struct {
				Success bool `json:"success"`
				Error   struct {
					Code string `json:"code"`
				} `json:"error"`
			}
			if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
				t.Fatalf("json.Unmarshal() error = %v", err)
			}
			if resp.Success || resp.Error.Code != "invalid_request" {
				t.Fatalf("response = %+v, want invalid_request error", resp)
			}
		})
	}
}

func TestCreateContextWaitUntilDoneIncludesCMDTerminalStatus(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell command fixture requires POSIX")
	}
	handler := NewContextHandler(ctxpkg.NewManager(), zap.NewNop())
	req := httptest.NewRequest(http.MethodPost, "/contexts", strings.NewReader(`{
		"type": "cmd",
		"cmd": {"command": ["/bin/sh", "-c", "printf done; sleep 0.05; exit 7"]},
		"wait_until_done": true
	}`))
	rec := httptest.NewRecorder()

	handler.Create(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusCreated, rec.Body.String())
	}
	var resp struct {
		Success bool            `json:"success"`
		Data    ContextResponse `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if !resp.Success {
		t.Fatalf("success = false, body = %s", rec.Body.String())
	}
	if resp.Data.OutputRaw != "done" {
		t.Fatalf("output_raw = %q, want done", resp.Data.OutputRaw)
	}
	if resp.Data.ExitCode == nil || *resp.Data.ExitCode != 7 {
		t.Fatalf("exit_code = %#v, want 7", resp.Data.ExitCode)
	}
	if resp.Data.State != string(process.ProcessStateCrashed) {
		t.Fatalf("state = %q, want %q", resp.Data.State, process.ProcessStateCrashed)
	}
}

func TestCreateContextWaitUntilDoneIncludesCMDSplitOutput(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("shell command fixture requires POSIX")
	}
	handler := NewContextHandler(ctxpkg.NewManager(), zap.NewNop())
	req := httptest.NewRequest(http.MethodPost, "/contexts", strings.NewReader(`{
		"type": "cmd",
		"cmd": {"command": ["/bin/sh", "-c", "printf out; printf err >&2; exit 7"]},
		"wait_until_done": true
	}`))
	rec := httptest.NewRecorder()

	handler.Create(rec, req)

	if rec.Code != http.StatusCreated {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusCreated, rec.Body.String())
	}
	var resp struct {
		Success bool            `json:"success"`
		Data    ContextResponse `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if !resp.Success {
		t.Fatalf("success = false, body = %s", rec.Body.String())
	}
	if resp.Data.Stdout == nil || *resp.Data.Stdout != "out" {
		t.Fatalf("stdout = %#v, want out", resp.Data.Stdout)
	}
	if resp.Data.Stderr == nil || *resp.Data.Stderr != "err" {
		t.Fatalf("stderr = %#v, want err", resp.Data.Stderr)
	}
	if !strings.Contains(resp.Data.OutputRaw, "out") || !strings.Contains(resp.Data.OutputRaw, "err") {
		t.Fatalf("output_raw = %q, want merged output containing stdout and stderr", resp.Data.OutputRaw)
	}
	if resp.Data.ExitCode == nil || *resp.Data.ExitCode != 7 {
		t.Fatalf("exit_code = %#v, want 7", resp.Data.ExitCode)
	}
	if resp.Data.State != string(process.ProcessStateCrashed) {
		t.Fatalf("state = %q, want %q", resp.Data.State, process.ProcessStateCrashed)
	}
}

func TestExecReturnsCMDTerminalStatus(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 1)
	var proc *fakeProcess
	proc = &fakeProcess{
		outputCh:    outputCh,
		processType: process.ProcessTypeCMD,
		onWrite: func([]byte) {
			outputCh <- process.ProcessOutput{Source: process.OutputSourceStdout, Data: []byte("failed\n")}
			proc.triggerExit(process.ExitEvent{ExitCode: 9, State: process.ProcessStateCrashed})
			close(outputCh)
		},
	}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeCMD)
	req := httptest.NewRequest(http.MethodPost, "/contexts/"+ctx.ID+"/exec", strings.NewReader(`{"data":"run"}`))
	req = mux.SetURLVars(req, map[string]string{"id": ctx.ID})
	rec := httptest.NewRecorder()

	handler.Exec(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	var resp struct {
		Success bool                `json:"success"`
		Data    ContextExecResponse `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if !resp.Success {
		t.Fatalf("success = false, body = %s", rec.Body.String())
	}
	if resp.Data.OutputRaw != "failed\n" {
		t.Fatalf("output_raw = %q, want failed newline", resp.Data.OutputRaw)
	}
	if resp.Data.ExitCode == nil || *resp.Data.ExitCode != 9 {
		t.Fatalf("exit_code = %#v, want 9", resp.Data.ExitCode)
	}
	if resp.Data.State != string(process.ProcessStateCrashed) {
		t.Fatalf("state = %q, want %q", resp.Data.State, process.ProcessStateCrashed)
	}
}

func TestExecReturnsCMDSplitOutput(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 2)
	var proc *fakeProcess
	proc = &fakeProcess{
		outputCh:    outputCh,
		processType: process.ProcessTypeCMD,
		stdout:      "out",
		stderr:      "err",
		onWrite: func([]byte) {
			outputCh <- process.ProcessOutput{Source: process.OutputSourceStdout, Data: []byte("out")}
			outputCh <- process.ProcessOutput{Source: process.OutputSourceStderr, Data: []byte("err")}
			proc.triggerExit(process.ExitEvent{ExitCode: 7, State: process.ProcessStateCrashed})
			close(outputCh)
		},
	}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeCMD)
	req := httptest.NewRequest(http.MethodPost, "/contexts/"+ctx.ID+"/exec", strings.NewReader(`{"data":"run"}`))
	req = mux.SetURLVars(req, map[string]string{"id": ctx.ID})
	rec := httptest.NewRecorder()

	handler.Exec(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
	}
	var resp struct {
		Success bool                `json:"success"`
		Data    ContextExecResponse `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if !resp.Success {
		t.Fatalf("success = false, body = %s", rec.Body.String())
	}
	if resp.Data.Stdout == nil || *resp.Data.Stdout != "out" {
		t.Fatalf("stdout = %#v, want out", resp.Data.Stdout)
	}
	if resp.Data.Stderr == nil || *resp.Data.Stderr != "err" {
		t.Fatalf("stderr = %#v, want err", resp.Data.Stderr)
	}
	if resp.Data.OutputRaw != "outerr" {
		t.Fatalf("output_raw = %q, want outerr", resp.Data.OutputRaw)
	}
	if resp.Data.ExitCode == nil || *resp.Data.ExitCode != 7 {
		t.Fatalf("exit_code = %#v, want 7", resp.Data.ExitCode)
	}
}

func TestExecInputSync_PromptBeforeOutputReturnsEmpty(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 2)
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func([]byte) {
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePTY, Data: []byte("hello world")}
			close(outputCh)
		},
	}

	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)
	output, execErr, timedOut := handler.execInputSync(ctx, "print('hello world')\n", context.Background())

	if timedOut {
		t.Fatal("execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("execInputSync() error = %v", execErr)
	}
	if output != "_S0_> hello world" {
		t.Errorf("output = %q, want %q when prompt arrives first", output, "_S0_> hello world")
	}
}

func TestExecInputSync_OutputBeforePromptReturnsOutput(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 2)
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func([]byte) {
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePTY, Data: []byte("hello world")}
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			close(outputCh)
		},
	}

	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)
	output, execErr, timedOut := handler.execInputSync(ctx, "print('hello world')\n", context.Background())

	if timedOut {
		t.Fatal("execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("execInputSync() error = %v", execErr)
	}
	if output != "_S0_> hello world" {
		t.Errorf("output = %q, want %q", output, "_S0_> hello world")
	}
}

func TestExecInputSync_REPLPromptFormatting_FirstAndMultipleRuns(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 8)
	writeCount := 0
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func([]byte) {
			writeCount++
			switch writeCount {
			case 1:
				outputCh <- process.ProcessOutput{
					Source: process.OutputSourcePTY,
					Data:   []byte("print('hello')\nhello\n_S0_> "),
				}
				outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			case 2:
				outputCh <- process.ProcessOutput{
					Source: process.OutputSourcePTY,
					Data:   []byte("print('world')\nworld\n_S0_> "),
				}
				outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			default:
				t.Fatalf("unexpected write count: %d", writeCount)
			}
		},
	}

	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)

	firstOutput, execErr, timedOut := handler.execInputSync(ctx, "print('hello')", context.Background())
	if timedOut {
		t.Fatal("first execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("first execInputSync() error = %v", execErr)
	}
	wantFirst := "_S0_> print('hello')\nhello\n"
	if firstOutput != wantFirst {
		t.Fatalf("first output = %q, want %q", firstOutput, wantFirst)
	}

	secondOutput, execErr, timedOut := handler.execInputSync(ctx, "print('world')", context.Background())
	if timedOut {
		t.Fatal("second execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("second execInputSync() error = %v", execErr)
	}
	wantSecond := "_S0_> print('world')\nworld\n"
	if secondOutput != wantSecond {
		t.Fatalf("second output = %q, want %q", secondOutput, wantSecond)
	}
}

func TestExecInputSync_AppendsNewlineForREPL(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 1)
	var captured string
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func(data []byte) {
			captured = string(data)
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			close(outputCh)
		},
	}

	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)
	_, execErr, timedOut := handler.execInputSync(ctx, "print('hello world')", context.Background())

	if timedOut {
		t.Fatal("execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("execInputSync() error = %v", execErr)
	}
	if captured != "print('hello world')\n" {
		t.Errorf("input = %q, want %q", captured, "print('hello world')\n")
	}
}

func TestExecInputSync_DoesNotAppendNewlineForCMD(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 1)
	var captured string
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func(data []byte) {
			captured = string(data)
			outputCh <- process.ProcessOutput{Source: process.OutputSourcePrompt}
			close(outputCh)
		},
	}

	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeCMD)
	_, execErr, timedOut := handler.execInputSync(ctx, "echo hello", context.Background())

	if timedOut {
		t.Fatal("execInputSync() timed out")
	}
	if execErr != nil {
		t.Fatalf("execInputSync() error = %v", execErr)
	}
	if captured != "echo hello" {
		t.Errorf("input = %q, want %q", captured, "echo hello")
	}
}

func TestNormalizeStringMap_ReturnsEmptyMapForNil(t *testing.T) {
	got := normalizeStringMap(nil)
	if got == nil {
		t.Fatal("normalizeStringMap(nil) returned nil")
	}
	if len(got) != 0 {
		t.Fatalf("normalizeStringMap(nil) length = %d, want 0", len(got))
	}
}

func TestGetContext_EncodesNilEnvVarsAsEmptyObject(t *testing.T) {
	proc := &fakeProcess{outputCh: make(chan process.ProcessOutput)}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)
	ctx.EnvVars = nil

	req := httptest.NewRequest(http.MethodGet, "/contexts/"+ctx.ID, nil)
	req = mux.SetURLVars(req, map[string]string{"id": ctx.ID})
	rec := httptest.NewRecorder()

	handler.Get(rec, req)

	if rec.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
	}

	var resp struct {
		Success bool `json:"success"`
		Data    struct {
			EnvVars map[string]string `json:"env_vars"`
		} `json:"data"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if !resp.Success {
		t.Fatalf("success = false, body = %s", rec.Body.String())
	}
	if resp.Data.EnvVars == nil {
		t.Fatal("env_vars decoded as nil, want empty object")
	}
	if len(resp.Data.EnvVars) != 0 {
		t.Fatalf("env_vars length = %d, want 0", len(resp.Data.EnvVars))
	}
	if strings.Contains(rec.Body.String(), "\"env_vars\":null") {
		t.Fatalf("response body contains null env_vars: %s", rec.Body.String())
	}
	if !strings.Contains(rec.Body.String(), "\"env_vars\":{}") {
		t.Fatalf("response body does not contain empty object env_vars: %s", rec.Body.String())
	}
}

func TestWriteInputRejectsMissingData(t *testing.T) {
	proc := &fakeProcess{outputCh: make(chan process.ProcessOutput)}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)

	req := httptest.NewRequest(http.MethodPost, "/contexts/"+ctx.ID+"/input", strings.NewReader(`{}`))
	req = mux.SetURLVars(req, map[string]string{"id": ctx.ID})
	rec := httptest.NewRecorder()

	handler.WriteInput(rec, req)

	if rec.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
	}
	var resp struct {
		Success bool `json:"success"`
		Error   struct {
			Code string `json:"code"`
		} `json:"error"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if resp.Success || resp.Error.Code != "invalid_request" {
		t.Fatalf("response = %+v, want invalid_request error", resp)
	}
}

func TestWriteInputReturnsGoneForFinishedProcess(t *testing.T) {
	proc := &fakeProcess{
		outputCh: make(chan process.ProcessOutput),
		writeErr: process.ErrProcessFinished,
	}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)

	req := httptest.NewRequest(http.MethodPost, "/contexts/"+ctx.ID+"/input", strings.NewReader(`{"data":"hello"}`))
	req = mux.SetURLVars(req, map[string]string{"id": ctx.ID})
	rec := httptest.NewRecorder()

	handler.WriteInput(rec, req)

	if rec.Code != http.StatusGone {
		t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusGone, rec.Body.String())
	}
	var resp struct {
		Success bool `json:"success"`
		Error   struct {
			Code string `json:"code"`
		} `json:"error"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("json.Unmarshal() error = %v", err)
	}
	if resp.Success || resp.Error.Code != "process_finished" {
		t.Fatalf("response = %+v, want process_finished error", resp)
	}
}

func TestNewWSOutputMessage(t *testing.T) {
	msg := newWSOutputMessage(process.OutputSourcePTY, "hello")
	if msg.Type != "output" {
		t.Fatalf("type = %q, want %q", msg.Type, "output")
	}
	if msg.Source != "pty" {
		t.Fatalf("source = %q, want %q", msg.Source, "pty")
	}
	if msg.Data != "hello" {
		t.Fatalf("data = %q, want %q", msg.Data, "hello")
	}
}

func TestNewWSDoneMessage(t *testing.T) {
	msg := newWSRequestDoneMessage("req-1")
	if msg.Type != "done" {
		t.Fatalf("type = %q, want %q", msg.Type, "done")
	}
	if msg.RequestID != "req-1" {
		t.Fatalf("request_id = %q, want %q", msg.RequestID, "req-1")
	}
	if msg.ExitCode != nil {
		t.Fatal("exit_code should be nil for request done")
	}
	if msg.State != "" {
		t.Fatal("state should be empty for request done")
	}
}

func TestNewWSProcessDoneMessage(t *testing.T) {
	msg := newWSProcessDoneMessage(0, process.ProcessStateStopped)
	if msg.Type != "done" {
		t.Fatalf("type = %q, want %q", msg.Type, "done")
	}
	if msg.RequestID != "" {
		t.Fatalf("request_id = %q, want empty", msg.RequestID)
	}
	if msg.ExitCode == nil || *msg.ExitCode != 0 {
		t.Fatalf("exit_code = %#v, want 0", msg.ExitCode)
	}
	if msg.State != string(process.ProcessStateStopped) {
		t.Fatalf("state = %q, want %q", msg.State, process.ProcessStateStopped)
	}
}

func TestParseSignalRejectsOutOfRangeNumericSignal(t *testing.T) {
	if _, err := parseSignal("999"); err == nil {
		t.Fatal("parseSignal(999) error = nil, want unsupported signal")
	}
	if _, err := parseSignal("0"); err == nil {
		t.Fatal("parseSignal(0) error = nil, want unsupported signal")
	}
	sig, err := parseSignal("64")
	if err != nil {
		t.Fatalf("parseSignal(64) error = %v", err)
	}
	if sig != syscall.Signal(64) {
		t.Fatalf("parseSignal(64) = %v, want 64", sig)
	}
}

func TestWebSocketSendsProcessDoneMessage(t *testing.T) {
	outputCh := make(chan process.ProcessOutput, 2)
	proc := &fakeProcess{
		outputCh:     outputCh,
		processType:  process.ProcessTypeCMD,
		handlerAdded: make(chan struct{}),
	}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeCMD)

	router := mux.NewRouter()
	router.HandleFunc("/contexts/{id}/ws", handler.WebSocket)
	server := httptest.NewServer(router)
	defer server.Close()

	wsURL, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf("url.Parse() error = %v", err)
	}
	wsURL.Scheme = "ws"
	wsURL.Path = "/contexts/" + ctx.ID + "/ws"

	conn, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
	if err != nil {
		t.Fatalf("Dial() error = %v", err)
	}
	defer conn.Close()

	select {
	case <-proc.handlerAdded:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for exit handler registration")
	}

	outputCh <- process.ProcessOutput{Source: process.OutputSourceStdout, Data: []byte("done\n")}
	proc.triggerExit(process.ExitEvent{ExitCode: 7, State: process.ProcessStateCrashed})
	close(outputCh)

	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	var messages []map[string]any
	for {
		var msg map[string]any
		if err := conn.ReadJSON(&msg); err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure) {
				break
			}
			t.Fatalf("ReadJSON() error = %v", err)
		}
		messages = append(messages, msg)
	}

	if len(messages) != 2 {
		t.Fatalf("messages length = %d, want 2", len(messages))
	}
	if got := messages[0]["type"]; got != "output" {
		t.Fatalf("first message type = %v, want output", got)
	}
	if got := messages[1]["type"]; got != "done" {
		t.Fatalf("second message type = %v, want done", got)
	}
	if got := messages[1]["exit_code"]; got != float64(7) {
		t.Fatalf("exit_code = %v, want 7", got)
	}
	if got := messages[1]["state"]; got != string(process.ProcessStateCrashed) {
		t.Fatalf("state = %v, want %q", got, process.ProcessStateCrashed)
	}
}

func TestWebSocketInputPassesREPLBytesWithoutAppendingNewline(t *testing.T) {
	outputCh := make(chan process.ProcessOutput)
	written := make(chan string, 1)
	proc := &fakeProcess{
		outputCh: outputCh,
		onWrite: func(data []byte) {
			written <- string(data)
		},
	}
	handler, ctx := newHandlerWithContext(proc, process.ProcessTypeREPL)

	router := mux.NewRouter()
	router.HandleFunc("/contexts/{id}/ws", handler.WebSocket)
	server := httptest.NewServer(router)
	defer server.Close()

	wsURL, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf("url.Parse() error = %v", err)
	}
	wsURL.Scheme = "ws"
	wsURL.Path = "/contexts/" + ctx.ID + "/ws"

	conn, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
	if err != nil {
		t.Fatalf("Dial() error = %v", err)
	}
	defer conn.Close()

	if err := conn.WriteJSON(wsControlMessage{Type: "input", Data: "e"}); err != nil {
		t.Fatalf("WriteJSON() error = %v", err)
	}

	select {
	case got := <-written:
		if got != "e" {
			t.Fatalf("input = %q, want %q", got, "e")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for WebSocket input")
	}
}
