package session

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"

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

type trajectoryHandlerTestStore struct {
	sessions        []types.SessionSnapshot
	historyAgentID  string
	historyLimit    int
	historyBefore   int64
	history         []types.TrajectoryLogItem
	latestIDs       []string
	latest          map[string]types.TrajectoryLogItem
	allLive         chan types.TrajectoryLogItem
	unsubscribedAll bool
}

func (s *trajectoryHandlerTestStore) Session(string) (types.SessionSnapshot, bool) {
	return types.SessionSnapshot{}, true
}

func (s *trajectoryHandlerTestStore) Sessions() []types.SessionSnapshot { return s.sessions }

func (s *trajectoryHandlerTestStore) SubscribeTrajectory(string) <-chan types.TrajectoryLogItem {
	ch := make(chan types.TrajectoryLogItem)
	close(ch)
	return ch
}

func (s *trajectoryHandlerTestStore) SubscribeLiveTrajectory(string) <-chan types.TrajectoryLogItem {
	ch := make(chan types.TrajectoryLogItem)
	close(ch)
	return ch
}

func (s *trajectoryHandlerTestStore) SubscribeAllLiveTrajectory() <-chan types.TrajectoryLogItem {
	return s.allLive
}

func (s *trajectoryHandlerTestStore) UnsubscribeTrajectory(string, <-chan types.TrajectoryLogItem) {}

func (s *trajectoryHandlerTestStore) UnsubscribeAllTrajectory(<-chan types.TrajectoryLogItem) {
	s.unsubscribedAll = true
}

func (s *trajectoryHandlerTestStore) TrajectoryHistory(agentID string, limit int, before int64) []types.TrajectoryLogItem {
	s.historyAgentID = agentID
	s.historyLimit = limit
	s.historyBefore = before
	return s.history
}

func (s *trajectoryHandlerTestStore) LatestTrajectoryByAgent(agentIDs []string) map[string]types.TrajectoryLogItem {
	s.latestIDs = append([]string(nil), agentIDs...)
	return s.latest
}

func TestTrajectoryLatestHandlerUsesActiveSessionAgentIDs(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := &trajectoryHandlerTestStore{
		sessions: []types.SessionSnapshot{
			{AgentConfigID: "config-a", AgentID: "runtime-a"},
			{AgentConfigID: "config-a", AgentID: "runtime-a"},
			{AgentConfigID: "config-b"},
		},
		latest: map[string]types.TrajectoryLogItem{
			"config-a": {
				AgentID:   "config-a",
				Timestamp: 1234,
				Entry:     types.TrajectoryEntry{Kind: "status", Activity: "working", Detail: "Writing code"},
			},
		},
	}
	router := gin.New()
	router.GET("/trajectory/latest", NewTrajectoryLatestHandler(store))

	recorder := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/trajectory/latest", nil)
	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Equal(t, []string{"config-a", "runtime-a", "config-b"}, store.latestIDs)
	var body map[string]types.TrajectoryLogItem
	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body))
	require.Equal(t, "Writing code", body["config-a"].Entry.Detail)
}

func TestTrajectoryHistoryHandlerReturnsHistoryWithoutActiveSession(t *testing.T) {
	gin.SetMode(gin.TestMode)
	store := &trajectoryHandlerTestStore{
		history: []types.TrajectoryLogItem{
			{
				AgentID:   "agent-a",
				Timestamp: 1234,
				Entry:     types.TrajectoryEntry{Kind: "text", Text: "restored"},
			},
		},
	}
	router := gin.New()
	router.GET("/sessions/:agentId/trajectory/history", NewTrajectoryHistoryHandler(store))

	recorder := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/sessions/agent-a/trajectory/history?limit=50&before=2000", nil)
	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Equal(t, "agent-a", store.historyAgentID)
	require.Equal(t, 50, store.historyLimit)
	require.Equal(t, int64(2000), store.historyBefore)
	var body []types.TrajectoryLogItem
	require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &body))
	require.Len(t, body, 1)
	require.Equal(t, "restored", body[0].Entry.Text)
}

func TestTrajectoryLiveHandlerStreamsAllAgentEvents(t *testing.T) {
	gin.SetMode(gin.TestMode)
	ch := make(chan types.TrajectoryLogItem, 1)
	ch <- types.TrajectoryLogItem{
		AgentID:   "agent-a",
		Timestamp: 1234,
		Entry:     types.TrajectoryEntry{Kind: "tool_start", ToolName: "message_history_read"},
	}
	close(ch)
	store := &trajectoryHandlerTestStore{allLive: ch}
	router := gin.New()
	router.GET("/trajectory/live", NewTrajectoryLiveHandler(store))

	recorder := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/trajectory/live", nil)
	router.ServeHTTP(recorder, req)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.True(t, store.unsubscribedAll)
	body := recorder.Body.String()
	require.Contains(t, body, "event: trajectory")
	require.Contains(t, body, `"agentId":"agent-a"`)
	require.Contains(t, body, `"toolName":"message_history_read"`)
	require.Contains(t, body, "event: done")
}