package exec

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/wire/apiclient"
)

func TestRuntimeRoutesReturnBearerBoundMergedCatalogAndInvokeExtension(t *testing.T) {
	var received extensionInvokeRequest
	var receivedBody []byte
	upstream := httptest.NewServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		require.Equal(
			t,
			"/api/v1/workspaces/ws-1/agents/remote-agent/runtime/ext/gh/actions/list",
			request.URL.EscapedPath(),
		)
		var err error
		receivedBody, err = io.ReadAll(request.Body)
		require.NoError(t, err)
		require.NoError(t, json.Unmarshal(receivedBody, &received))
		writer.Header().Set("Content-Type", "application/json")
		_, _ = writer.Write([]byte(
			`{"domain":"gh","action":"list","data":{"items":[1]}}`,
		))
	}))
	defer upstream.Close()

	config := testRuntimeAuthorizationConfig()
	config.RuntimeCatalog = testRuntimeCatalog(t)
	config.ProcessGeneration = "process-1"
	config.Server = &apiclient.Client{
		ServerURL: upstream.URL, AgentID: "remote-agent",
		AuthToken: "agent-token", RouteWorkspaceID: "ws-1",
		HTTPClient: upstream.Client(),
	}
	store := newExecAuthorizationTestStore(t)
	lease, err := store.Create(config)
	require.NoError(t, err)
	router := gin.New()
	RegisterRoutes(
		router.Group("/api/internal/exec"),
		RouteDependencies{Authorizations: store},
	)

	catalog := runtimeRequest(
		router,
		lease.Token,
		http.MethodGet,
		"/api/internal/exec/servers/server-1/agents/agent-1/catalog",
		nil,
	)
	require.Equal(t, http.StatusOK, catalog.Code)
	var document runtimeCatalogDocument
	require.NoError(t, json.Unmarshal(catalog.Body.Bytes(), &document))
	view, err := document.view()
	require.NoError(t, err)
	builtin, ok := view.Lookup("workspace", "list")
	require.True(t, ok)
	require.Equal(t, RuntimeActionBuiltin, builtin.Source)
	extension, ok := view.Lookup("gh", "list")
	require.True(t, ok)
	require.Equal(t, RuntimeActionExtension, extension.Source)
	// Discovery contains action content only, never the HTTP cache validator or
	// the deleted Catalog identity fields.
	require.NotContains(t, catalog.Body.String(), `"ext-v1-auth"`)
	require.NotContains(t, catalog.Body.String(), "extensionCatalogVersion")
	require.NotContains(t, catalog.Body.String(), "extensionCatalogDigest")

	invoke := runtimeRequest(
		router,
		lease.Token,
		http.MethodPost,
		"/api/internal/exec/servers/server-1/agents/agent-1/gh/list",
		strings.NewReader(`{"limit":5}`),
	)
	require.Equal(t, http.StatusOK, invoke.Code)
	require.JSONEq(
		t,
		`{"domain":"gh","action":"list","data":{"items":[1]}}`,
		invoke.Body.String(),
	)
	require.NotContains(t, string(receivedBody), "definitionDigest")
	require.NotContains(t, string(receivedBody), "catalogVersion")
	require.JSONEq(t, `{"limit":5}`, string(received.Input))
}

func TestDynamicBuiltinRoutePreservesTypedValidation(t *testing.T) {
	store := newExecAuthorizationTestStore(t)
	lease, err := store.Create(testRuntimeAuthorizationConfig())
	require.NoError(t, err)
	router := gin.New()
	RegisterRoutes(
		router.Group("/api/internal/exec"),
		RouteDependencies{Authorizations: store},
	)

	for name, body := range map[string]string{
		"unknown field": `{"unexpected":true}`,
		"trailing JSON": `{} {}`,
		"oversize":      `{"query":"` + strings.Repeat("x", maxActionBodyBytes) + `"}`,
	} {
		t.Run(name, func(t *testing.T) {
			response := runtimeRequest(
				router,
				lease.Token,
				http.MethodPost,
				"/api/internal/exec/servers/server-1/agents/agent-1/workspace/list",
				strings.NewReader(body),
			)
			require.Equal(t, http.StatusBadRequest, response.Code)
			require.Contains(t, response.Body.String(), `"code":"validation"`)
		})
	}
}

func TestRuntimeRouteAuthenticatesBeforeReadingBody(t *testing.T) {
	router := gin.New()
	RegisterRoutes(
		router.Group("/api/internal/exec"),
		RouteDependencies{Authorizations: newExecAuthorizationTestStore(t)},
	)
	body := &countingExecReader{
		reader: bytes.NewBufferString(strings.Repeat("x", maxActionBodyBytes+1)),
	}
	request := httptest.NewRequest(
		http.MethodPost,
		"/api/internal/exec/servers/server-1/agents/agent-1/workspace/list",
		body,
	)
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	require.Equal(t, http.StatusForbidden, response.Code)
	require.Zero(t, body.reads)
}

func TestRuntimeExtensionRouteRevokeCancelsUpstreamAndFencesPublication(t *testing.T) {
	started := make(chan struct{})
	upstreamCanceled := make(chan struct{})
	config := testRuntimeAuthorizationConfig()
	config.RuntimeCatalog = testRuntimeCatalog(t)
	config.ProcessGeneration = "process-1"
	config.Server = &apiclient.Client{
		ServerURL: "http://agentteam.invalid", AgentID: "remote-agent",
		AuthToken: "agent-token", RouteWorkspaceID: "ws-1",
		HTTPClient: &http.Client{Transport: roundTripFunc(func(
			request *http.Request,
		) (*http.Response, error) {
			close(started)
			<-request.Context().Done()
			close(upstreamCanceled)
			return nil, request.Context().Err()
		})},
	}
	store := newExecAuthorizationTestStore(t)
	lease, err := store.Create(config)
	require.NoError(t, err)
	router := gin.New()
	RegisterRoutes(
		router.Group("/api/internal/exec"),
		RouteDependencies{Authorizations: store},
	)
	done := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		done <- runtimeRequest(
			router,
			lease.Token,
			http.MethodPost,
			"/api/internal/exec/servers/server-1/agents/agent-1/gh/list",
			strings.NewReader(`{}`),
		)
	}()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("Extension upstream did not start")
	}
	drained := store.Revoke(lease.AuthorizationID)
	select {
	case <-upstreamCanceled:
	case <-time.After(time.Second):
		t.Fatal("revocation did not cancel Extension upstream")
	}
	response := <-done
	require.Equal(t, http.StatusForbidden, response.Code)
	require.NotContains(t, response.Body.String(), `"items"`)
	select {
	case <-drained:
	case <-time.After(time.Second):
		t.Fatal("revoked authorization did not drain")
	}
}

func TestRuntimeRouteRevokeInterruptsBodyReadAndAllowsDrain(t *testing.T) {
	store := newExecAuthorizationTestStore(t)
	lease, err := store.Create(testRuntimeAuthorizationConfig())
	require.NoError(t, err)
	router := gin.New()
	RegisterRoutes(
		router.Group("/api/internal/exec"),
		RouteDependencies{Authorizations: store},
	)
	body := newBlockingExecBody()
	request := httptest.NewRequest(
		http.MethodPost,
		"/api/internal/exec/servers/server-1/agents/agent-1/workspace/list",
		body,
	)
	request.Header.Set("Authorization", "Bearer "+lease.Token)
	response := httptest.NewRecorder()
	done := make(chan struct{})
	go func() {
		router.ServeHTTP(runtimeRouteRecorder{ResponseRecorder: response}, request)
		close(done)
	}()
	select {
	case <-body.started:
	case <-time.After(time.Second):
		t.Fatal("route did not start reading body")
	}
	drained := store.Revoke(lease.AuthorizationID)
	select {
	case <-body.closed:
	case <-time.After(time.Second):
		t.Fatal("revocation did not close body")
	}
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("route remained blocked after revocation")
	}
	require.Equal(t, http.StatusForbidden, response.Code)
	select {
	case <-drained:
	case <-time.After(time.Second):
		t.Fatal("authorization did not drain")
	}
}

type countingExecReader struct {
	reader *bytes.Buffer
	reads  int
}

type blockingExecBody struct {
	started chan struct{}
	closed  chan struct{}
}

func newBlockingExecBody() *blockingExecBody {
	return &blockingExecBody{started: make(chan struct{}), closed: make(chan struct{})}
}

func (body *blockingExecBody) Read([]byte) (int, error) {
	select {
	case <-body.started:
	default:
		close(body.started)
	}
	<-body.closed
	return 0, io.ErrClosedPipe
}

func (body *blockingExecBody) Close() error {
	select {
	case <-body.closed:
	default:
		close(body.closed)
	}
	return nil
}

func (reader *countingExecReader) Read(buffer []byte) (int, error) {
	reader.reads++
	return reader.reader.Read(buffer)
}

type roundTripFunc func(*http.Request) (*http.Response, error)

func (function roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return function(request)
}

func runtimeRequest(
	router http.Handler,
	tokenValue,
	method,
	path string,
	body io.Reader,
) *httptest.ResponseRecorder {
	request := httptest.NewRequest(method, path, body)
	request.Header.Set("Authorization", "Bearer "+tokenValue)
	response := httptest.NewRecorder()
	router.ServeHTTP(runtimeRouteRecorder{ResponseRecorder: response}, request)
	return response
}

type runtimeRouteRecorder struct {
	*httptest.ResponseRecorder
}

func (runtimeRouteRecorder) SetReadDeadline(time.Time) error { return nil }