//go:build e2e

// Package helpers provides shared test infrastructure for e2e tests.
// It contains the fake WebSocket server, stub dispatcher, mock channel,
// JSON decode helpers, and common test utilities used across all e2e test
// packages (api, protocol, dispatcher, workspace).
package helpers

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/wire"
)

// ---------------------------------------------------------------------------
// Fake WebSocket server
// ---------------------------------------------------------------------------

// FakeServer is a minimal AgentTeam-compatible WebSocket server for testing.
// It captures messages sent by the daemon (channel) and allows the test to
// inject messages toward the daemon.
type FakeServer struct {
	srv *httptest.Server

	mu   sync.Mutex
	conn *websocket.Conn

	// FromDaemon receives messages sent by the channel to the server.
	FromDaemon chan json.RawMessage
	// ToDaemon is used by tests to send messages to the channel.
	ToDaemon chan []byte
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(_ *http.Request) bool { return true },
}

const fakeWorkspaceID = "wsp_e2e"

// NewFakeServer starts an httptest.Server that speaks the AgentTeam WebSocket
// protocol at the workspace-scoped daemon connect path. Caller is responsible
// for closing the server.
func NewFakeServer() *FakeServer {
	fs := &FakeServer{
		FromDaemon: make(chan json.RawMessage, 100),
		ToDaemon:   make(chan []byte, 100),
	}

	mux := http.NewServeMux()
	handleConnect := func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		fs.mu.Lock()
		fs.conn = conn
		fs.mu.Unlock()

		go func() {
			for data := range fs.ToDaemon {
				if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
					return
				}
			}
		}()

		for {
			_, data, err := conn.ReadMessage()
			if err != nil {
				return
			}
			select {
			case fs.FromDaemon <- json.RawMessage(data):
			default:
			}
		}
	}

	mux.HandleFunc("/api/v1/workspaces/"+fakeWorkspaceID+"/daemon/connect", handleConnect)
	mux.HandleFunc("/api/v1/workspaces/"+fakeWorkspaceID+"/agents/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"messageId":"msg-fake-001"}`))
	})

	fs.srv = httptest.NewServer(mux)
	return fs
}

// URL returns the server's base URL (http://...).
func (fs *FakeServer) URL() string {
	return fs.srv.URL
}

// Send sends a JSON message from the server to the daemon.
func (fs *FakeServer) Send(v any) {
	data, err := json.Marshal(v)
	if err != nil {
		return
	}
	fs.ToDaemon <- data
}

// SendRaw sends raw bytes from the server to the daemon.
func (fs *FakeServer) SendRaw(data []byte) {
	fs.ToDaemon <- data
}

// Close shuts down the fake server.
func (fs *FakeServer) Close() {
	close(fs.ToDaemon)
	fs.mu.Lock()
	if fs.conn != nil {
		_ = fs.conn.Close()
	}
	fs.mu.Unlock()
	fs.srv.Close()
}

// ---------------------------------------------------------------------------
// Message reading helpers
// ---------------------------------------------------------------------------

// ReadMsg reads the next message from the daemon, blocking up to timeout.
func (fs *FakeServer) ReadMsg(timeout time.Duration) (json.RawMessage, error) {
	select {
	case msg := <-fs.FromDaemon:
		return msg, nil
	case <-time.After(timeout):
		return nil, fmt.Errorf("readMsg: timeout after %s", timeout)
	}
}

// ReadUntilType reads messages until one with the given "type" field is found.
func (fs *FakeServer) ReadUntilType(msgType string, timeout time.Duration) (json.RawMessage, error) {
	deadline := time.Now().Add(timeout)
	for {
		remaining := time.Until(deadline)
		if remaining <= 0 {
			return nil, fmt.Errorf("readUntilType(%q): timeout after %s", msgType, timeout)
		}
		msg, err := fs.ReadMsg(remaining)
		if err != nil {
			return nil, fmt.Errorf("readUntilType(%q): %w", msgType, err)
		}
		var env struct {
			Type string `json:"type"`
		}
		if json.Unmarshal(msg, &env) == nil && env.Type == msgType {
			return msg, nil
		}
	}
}

// MustReadUntilType is like ReadUntilType but calls t.Fatalf on failure.
func (fs *FakeServer) MustReadUntilType(msgType string, timeout time.Duration) json.RawMessage {
	msg, err := fs.ReadUntilType(msgType, timeout)
	if err != nil {
		panic(fmt.Sprintf("MustReadUntilType(%q): %v", msgType, err))
	}
	return msg
}

// DrainReady reads and discards messages until the "ready" message is found.
func (fs *FakeServer) DrainReady() {
	msg, err := fs.ReadUntilType("ready", 5*time.Second)
	if err != nil {
		panic(fmt.Sprintf("DrainReady: %v", err))
	}
	_ = msg
}

// DrainUntil reads and discards messages until one of the given type is found.
func (fs *FakeServer) DrainUntil(msgType string) json.RawMessage {
	msg, err := fs.ReadUntilType(msgType, 5*time.Second)
	if err != nil {
		panic(fmt.Sprintf("DrainUntil(%q): %v", msgType, err))
	}
	return msg
}

// ---------------------------------------------------------------------------
// JSON decode helpers
// ---------------------------------------------------------------------------

// DecodeType extracts the "type" field from a JSON message.
func DecodeType(raw json.RawMessage) string {
	var env struct {
		Type string `json:"type"`
	}
	_ = json.Unmarshal(raw, &env)
	return env.Type
}

// DecodeField extracts a string field from a JSON message.
func DecodeField(raw json.RawMessage, field string) string {
	var m map[string]json.RawMessage
	if err := json.Unmarshal(raw, &m); err != nil {
		return ""
	}
	v, ok := m[field]
	if !ok {
		return ""
	}
	var s string
	if json.Unmarshal(v, &s) == nil {
		return s
	}
	return string(v)
}

// DecodeIntField extracts a numeric field from a JSON message.
func DecodeIntField(raw json.RawMessage, field string) int {
	var m map[string]json.RawMessage
	if err := json.Unmarshal(raw, &m); err != nil {
		return 0
	}
	v, ok := m[field]
	if !ok {
		return 0
	}
	var n float64
	if json.Unmarshal(v, &n) == nil {
		return int(n)
	}
	return 0
}

// DecodeBoolField extracts a boolean field from a JSON message.
func DecodeBoolField(raw json.RawMessage, field string) bool {
	var m map[string]json.RawMessage
	if err := json.Unmarshal(raw, &m); err != nil {
		return false
	}
	v, ok := m[field]
	if !ok {
		return false
	}
	var b bool
	if json.Unmarshal(v, &b) == nil {
		return b
	}
	return false
}

// ---------------------------------------------------------------------------
// Mock channel
// ---------------------------------------------------------------------------

// MockChannel is a test channel that captures outbound messages.
type MockChannel struct {
	IDVal      string
	EnabledVal bool
	SentMsgs   []types.OutboundMessage
	SentMeta   []types.Metadata
	DebugOnVal bool
}

func (m *MockChannel) ID() string                    { return m.IDVal }
func (m *MockChannel) Enabled() bool                 { return m.EnabledVal }
func (m *MockChannel) Start(_ context.Context) error { return nil }
func (m *MockChannel) Stop(_ context.Context) error  { return nil }
func (m *MockChannel) Send(_ context.Context, meta types.Metadata, msg types.OutboundMessage) error {
	m.SentMsgs = append(m.SentMsgs, msg)
	m.SentMeta = append(m.SentMeta, meta)
	return nil
}
func (m *MockChannel) MCPHandler() types.MCPHandler { return nil }
func (m *MockChannel) ToggleDebug() bool            { m.DebugOnVal = !m.DebugOnVal; return m.DebugOnVal }

// LastReply returns the content of the last sent message.
func (m *MockChannel) LastReply() string {
	if len(m.SentMsgs) == 0 {
		return ""
	}
	return m.SentMsgs[len(m.SentMsgs)-1].Content
}

// ---------------------------------------------------------------------------
// Stub dispatcher
// ---------------------------------------------------------------------------

// StubDispatcher provides a fully-featured test double for e2e tests.
// It records all Dispatch calls and provides no-op implementations
// for methods that the agentteam handlers invoke.
type StubDispatcher struct {
	mu             sync.Mutex
	DispatchedMsgs []*types.Message
	StartedAgents  []string
	StoppedAgents  []string
	SleptAgents    []string
	RestoredAgents []string
}

func (s *StubDispatcher) Session(string) (types.SessionSnapshot, bool) {
	return types.SessionSnapshot{}, false
}
func (s *StubDispatcher) Sessions() []types.SessionSnapshot { return nil }
func (s *StubDispatcher) History(string, int, int64, int64) ([]types.HistoryMessage, error) {
	return nil, nil
}
func (s *StubDispatcher) PrepareVisibleReply(string, ...string) string     { return "" }
func (s *StubDispatcher) ConfirmReactionReply(string, string, string) bool { return false }
func (s *StubDispatcher) ReconcileDeferredDeliveries(string, []string) ([]string, error) {
	return nil, nil
}
func (s *StubDispatcher) AppendDeferredDelivery(string, int, json.RawMessage) error { return nil }
func (s *StubDispatcher) HasDeferredDeliveries(string) bool                         { return false }

func (s *StubDispatcher) Dispatch(_ context.Context, _ types.ServerConnection, msg *types.Message) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.DispatchedMsgs = append(s.DispatchedMsgs, msg)
}

func (s *StubDispatcher) DispatchForDeliveryAck(ctx context.Context, ch types.ServerConnection, msg *types.Message) error {
	s.Dispatch(ctx, ch, msg)
	return nil
}

func (s *StubDispatcher) StartAgentForAgentTeam(_, serverAgentID string, _ types.AgentStartConfig, _ *types.Message) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.StartedAgents = append(s.StartedAgents, serverAgentID)
	return nil
}

func (s *StubDispatcher) StopByServerAgentID(serverAgentID string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.StoppedAgents = append(s.StoppedAgents, serverAgentID)
	return true
}

func (s *StubDispatcher) InterruptByServerAgentID(serverAgentID string) bool { return false }

func (s *StubDispatcher) InterruptActiveRunByServerAgentID(serverAgentID string) bool {
	return false
}

func (s *StubDispatcher) SleepByServerAgentID(serverAgentID string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.SleptAgents = append(s.SleptAgents, serverAgentID)
	return true
}

func (s *StubDispatcher) RestoreAgentRecoveryPoint(_, serverAgentID, _, _, _, _, _, _ string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.RestoredAgents = append(s.RestoredAgents, serverAgentID)
	return nil
}

func (s *StubDispatcher) GCRecoveryPoints(string, []types.RecoveryArtifactRequest) error {
	return nil
}
func (s *StubDispatcher) UpdateAgentIdentity(string, string, string) {}
func (s *StubDispatcher) SubscribeTrajectory(string) <-chan types.TrajectoryLogItem {
	ch := make(chan types.TrajectoryLogItem)
	close(ch)
	return ch
}
func (s *StubDispatcher) UnsubscribeTrajectory(string, <-chan types.TrajectoryLogItem) {}
func (s *StubDispatcher) ServerConnections() types.ServerConnectionAccessor            { return nil }
func (s *StubDispatcher) TrajectoryHistory(string, int, int64) []types.TrajectoryLogItem {
	return nil
}

// GetDispatchedMsgs returns a snapshot of all dispatched messages.
func (s *StubDispatcher) GetDispatchedMsgs() []*types.Message {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := make([]*types.Message, len(s.DispatchedMsgs))
	copy(cp, s.DispatchedMsgs)
	return cp
}

// GetStartedAgents returns a snapshot of all started agent IDs.
func (s *StubDispatcher) GetStartedAgents() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := make([]string, len(s.StartedAgents))
	copy(cp, s.StartedAgents)
	return cp
}

// ---------------------------------------------------------------------------
// Channel factory
// ---------------------------------------------------------------------------

// StartChannel creates a real wire.Protocol connected to the given FakeServer.
// Returns the channel and the stub dispatcher.
func StartChannel(fs *FakeServer) (types.ServerConnection, *StubDispatcher) {
	disp := &StubDispatcher{}

	cfgYAML := fmt.Sprintf("server_url: %s\napi_key: test-api-key\nagent_id: test-agent\nworkspace_id: %s\n", fs.URL(), fakeWorkspaceID)

	chCfg := profile.AgentTeamServer{
		ID:      "e2e-agentteam",
		Enabled: true,
		Config:  types.RawMessage(cfgYAML),
	}

	ch, err := wire.New(chCfg, disp)
	if err != nil {
		panic(fmt.Sprintf("wire.New: %v", err))
	}

	ctx, cancel := context.WithCancel(context.Background())
	if err := ch.Start(ctx); err != nil {
		cancel()
		panic(fmt.Sprintf("channel.Start: %v", err))
	}

	_ = cancel

	return ch, disp
}

// StopChannel stops a channel created by StartChannel.
func StopChannel(ch types.ServerConnection) {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	_ = ch.Stop(ctx)
}

// ---------------------------------------------------------------------------
// Build and environment helpers
// ---------------------------------------------------------------------------

// BuildAgentTeam builds the agentteam binary and returns its path.
func BuildAgentTeam() (string, error) {
	bin := filepath.Join(os.TempDir(), "agentteam")
	root, err := RepoRoot()
	if err != nil {
		return "", err
	}
	cmd := exec.Command("go", "build", "-o", bin, "./cmd/agentteam")
	cmd.Dir = root
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("build agentteam: %v\n%s", err, out)
	}
	return bin, nil
}

// RepoRoot walks up from the current directory to find go.mod.
func RepoRoot() (string, error) {
	dir, err := os.Getwd()
	if err != nil {
		return "", err
	}
	for {
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			return dir, nil
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return "", fmt.Errorf("cannot find go.mod")
		}
		dir = parent
	}
}

// FailpointEnv returns environment variables with GO_FAILPOINTS set to inject
// the given config directory via the agentteam-config-dir failpoint.
func FailpointEnv(cfgDir string) []string {
	return append(os.Environ(),
		"GO_FAILPOINTS="+fmt.Sprintf(`agentteam-config-dir=return("%s")`, cfgDir),
		"HOME="+os.TempDir(),
	)
}

// ---------------------------------------------------------------------------
// Message construction helpers
// ---------------------------------------------------------------------------

// MakeMessage creates a types.Message for testing.
func MakeMessage(channelID string, meta types.Metadata, prompt string) *types.Message {
	return &types.Message{
		ID:      "msg-1",
		AgentID: "",
		Input: types.MessageInput{
			Prompt:  prompt,
			Channel: channelID,
			Meta:    meta,
		},
	}
}

// MakeAgentTeamMeta creates a AgentTeamMeta for testing.
func MakeAgentTeamMeta(agentID, target string) *types.AgentTeamMeta {
	return &types.AgentTeamMeta{
		ID:           "msg-agentteam-1",
		AuthorTarget: "user-1",
		Target:       target,
		AgentID:      agentID,
		AuthorProfile: types.AuthorProfile{
			ActorType: "human",
			TargetID:  "user-1",
		},
		Timestamp: fmt.Sprintf("%d", time.Now().UnixMilli()),
	}
}

// ---------------------------------------------------------------------------
// Archive helpers
// ---------------------------------------------------------------------------

// CreateTestZip creates a zip archive at path with the given files map.
func CreateTestZip(path string, files map[string]string) error {
	var buf bytes.Buffer
	w := zip.NewWriter(&buf)
	for name, content := range files {
		fw, err := w.Create(name)
		if err != nil {
			return err
		}
		fw.Write([]byte(content))
	}
	w.Close()
	return os.WriteFile(path, buf.Bytes(), 0644)
}

// CreateTestTarGz creates a tar.gz archive at path with the given files map.
func CreateTestTarGz(path string, files map[string]string) error {
	var buf bytes.Buffer
	gw := gzip.NewWriter(&buf)
	tw := tar.NewWriter(gw)
	for name, content := range files {
		tw.WriteHeader(&tar.Header{
			Name: name,
			Mode: 0644,
			Size: int64(len(content)),
		})
		tw.Write([]byte(content))
	}
	tw.Close()
	gw.Close()
	return os.WriteFile(path, buf.Bytes(), 0644)
}