package router

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/api"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
	"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)

func TestHandlerRoutesByResolvedWorkspaceGroup(t *testing.T) {
	t.Parallel()

	upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		require.Equal(t, "ws_1", request.Header.Get(api.HeaderRouteWorkspaceID))
		require.Equal(t, "ws_1", request.Header.Get(api.HeaderWorkspaceID))
		require.Equal(t, "workspace-large-a", request.Header.Get(HeaderResolvedWorkspaceGroupID))
		writer.Header().Set("X-Upstream-Path", request.URL.Path)
		_, _ = writer.Write([]byte("ok"))
	}))
	t.Cleanup(upstream.Close)

	handler, err := NewHandler(fakeResolver{
		target: workspace.WorkspaceRouteTarget{
			WorkspaceID:      "ws_1",
			WorkspaceGroupID: "workspace-large-a",
			ServiceRef:       "workspace-large-a",
		},
	}, map[string]string{"workspace-large-a": upstream.URL})
	require.NoError(t, err)

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodPost, "/api/channels", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "ws_1")
	request.Header.Set(api.HeaderWorkspaceID, "ws_1")
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusOK, recorder.Code)
	require.Equal(t, "ok", recorder.Body.String())
	require.Equal(t, "/api/channels", recorder.Header().Get("X-Upstream-Path"))
}

func TestHandlerRejectsPlacementWorkspaceMismatch(t *testing.T) {
	t.Parallel()

	handler, err := NewHandler(fakeResolver{
		target: workspace.WorkspaceRouteTarget{
			WorkspaceID:      "ws_other",
			WorkspaceGroupID: "workspace-default",
		},
	}, map[string]string{"workspace-default": "http://127.0.0.1:1"})
	require.NoError(t, err)

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/ws_1/kb", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "ws_1")
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusConflict, recorder.Code)
	require.Contains(t, recorder.Body.String(), "placement identity does not match")
}

func TestHandlerRejectsMissingWorkspaceID(t *testing.T) {
	t.Parallel()

	handler, err := NewHandler(fakeResolver{}, map[string]string{"workspace-default": "http://127.0.0.1:1"})
	require.NoError(t, err)

	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/api/channels", nil))

	require.Equal(t, http.StatusBadRequest, recorder.Code)
	require.Contains(t, recorder.Body.String(), "X-AgentTeam-Route-Workspace-Id is required")
}

func TestHandlerMapsResolverErrors(t *testing.T) {
	t.Parallel()

	handler, err := NewHandler(fakeResolver{err: storeerr.NotFound("workspace placement not found")}, map[string]string{"workspace-default": "http://127.0.0.1:1"})
	require.NoError(t, err)

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodGet, "/api/channels", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "ws_missing")
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusNotFound, recorder.Code)
	require.Contains(t, recorder.Body.String(), "workspace placement not found")
}

func TestHandlerRejectsUnconfiguredWorkspaceGroup(t *testing.T) {
	t.Parallel()

	handler, err := NewHandler(fakeResolver{
		target: workspace.WorkspaceRouteTarget{
			WorkspaceID:      "ws_1",
			WorkspaceGroupID: "workspace-large-a",
		},
	}, map[string]string{"workspace-default": "http://127.0.0.1:1"})
	require.NoError(t, err)

	recorder := httptest.NewRecorder()
	request := httptest.NewRequest(http.MethodPost, "/api/channels", nil)
	request.Header.Set(api.HeaderRouteWorkspaceID, "ws_1")
	handler.ServeHTTP(recorder, request)

	require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
	require.Contains(t, recorder.Body.String(), "workspace group workspace-large-a has no configured upstream")
}

func TestNewHandlerValidatesUpstreams(t *testing.T) {
	t.Parallel()

	_, err := NewHandler(fakeResolver{}, map[string]string{"workspace-default": "://bad"})
	require.Error(t, err)
	require.Contains(t, err.Error(), "parse workspace router upstream")

	_, err = NewHandler(fakeResolver{}, map[string]string{"workspace-default": "localhost:3001"})
	require.EqualError(t, err, "workspace router upstream for group workspace-default must include scheme and host")

	_, err = NewHandler(nil, map[string]string{})
	require.EqualError(t, err, "workspace router resolver is required")
}

type fakeResolver struct {
	target workspace.WorkspaceRouteTarget
	err    error
}

func (resolver fakeResolver) ResolveWorkspaceRouteTarget(_ context.Context, workspaceID string) (workspace.WorkspaceRouteTarget, error) {
	if resolver.err != nil {
		return workspace.WorkspaceRouteTarget{}, resolver.err
	}
	if resolver.target.WorkspaceID == "" {
		return workspace.WorkspaceRouteTarget{}, fmt.Errorf("unexpected workspace id %s", workspaceID)
	}
	return resolver.target, nil
}