package machine

import (
	"archive/zip"
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"testing"

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

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/diagnosticbundle"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
)

func TestDiagnosticDownloadScopesArchiveToCurrentDesktopAgents(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		http.NotFound(writer, request)
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	configDir := t.TempDir()
	profileManager := profile.New(configDir)
	server := profile.AgentTeamServer{
		ID:          "server-current",
		Enabled:     true,
		ServerURL:   upstream.URL,
		WorkspaceID: "wsp_1",
		MachineID:   "mch_local",
		APIKey:      "machine-secret",
	}
	require.NoError(t, profileManager.Servers().Add(&server))
	require.NoError(t, profileManager.Agents().Add(agent.Config{
		ID:                  "agent-a",
		Runtime:             "opencode",
		RuntimeDebugEnabled: true,
		Source:              agent.NewAgentTeamSource(server.ID),
	}))
	require.NoError(t, profileManager.Agents().Add(agent.Config{
		ID:      "agent-other",
		Runtime: "opencode",
		Source:  agent.NewAgentTeamSource("server-other"),
	}))
	writeMachineDiagnostic(t, configDir, "agent-a", "current-agent")
	writeMachineDiagnostic(t, configDir, "agent-other", "foreign-agent")

	sessions := uisession.NewStore(uisession.Options{})
	sessionID := token.CreateSession()
	expiresAt, ok := token.SessionExpiresAt(sessionID)
	require.True(t, ok)
	_, err := sessions.Select(sessionID, expiresAt, target.ID)
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiagnosticDownloadRoute(router, Dependencies{
		Auth:     authManager,
		Agents:   profileManager.Agents(),
		Servers:  profileManager.Servers(),
		Sessions: sessions,
		DiagnosticPaths: diagnosticbundle.Paths{
			ConfigDir: configDir,
		},
	})

	path := "/api/internal/machines/mch_local/diagnostics/archive?workspaceId=wsp_1&accessToken=" +
		url.QueryEscape(sessionID)
	response := performDiagnosticRequest(router, path, "")
	require.Equal(t, http.StatusOK, response.Code)
	require.Equal(t, "application/zip", response.Header().Get("Content-Type"))
	require.Equal(t, "no-store", response.Header().Get("Cache-Control"))
	require.Contains(t, response.Header().Get("Content-Disposition"), "agentteam-diagnostics-")

	entries := diagnosticArchiveEntries(t, response.Body.Bytes())
	require.Contains(t, entries, "agents/agent-a/opencode-runtime-debug.jsonl")
	require.Contains(t, string(entries["agents/agent-a/opencode-runtime-debug.jsonl"]), "current-agent")
	require.NotContains(t, entries, "agents/agent-other/opencode-runtime-debug.jsonl")
	require.NotContains(t, string(response.Body.Bytes()), "foreign-agent")

	wrongMachine := performDiagnosticRequest(
		router,
		"/api/internal/machines/mch_other/diagnostics/archive?workspaceId=wsp_1&accessToken="+url.QueryEscape(sessionID),
		"",
	)
	require.Equal(t, http.StatusConflict, wrongMachine.Code)

	wrongWorkspace := performDiagnosticRequest(
		router,
		"/api/internal/machines/mch_local/diagnostics/archive?workspaceId=wsp_other&accessToken="+url.QueryEscape(sessionID),
		"",
	)
	require.Equal(t, http.StatusConflict, wrongWorkspace.Code)

	mixedCredentials := performDiagnosticRequest(router, path, sessionID)
	require.Equal(t, http.StatusBadRequest, mixedCredentials.Code)

	require.True(t, acquireDiagnosticDownload())
	busy := performDiagnosticRequest(router, path, "")
	releaseDiagnosticDownload()
	require.Equal(t, http.StatusTooManyRequests, busy.Code)
	require.Equal(t, "5", busy.Header().Get("Retry-After"))
}

func TestDiagnosticDownloadAcceptsAuthorizationHeader(t *testing.T) {
	setMachineTestEncryptionKeys(t)
	upstream := newMachineTestServer(t, func(writer http.ResponseWriter, request *http.Request) {
		http.NotFound(writer, request)
	})
	authManager, target := loginMachineTestUser(t, upstream.URL)
	configDir := t.TempDir()
	profileManager := profile.New(configDir)
	server := profile.AgentTeamServer{
		ID:          "server-current",
		Enabled:     true,
		ServerURL:   upstream.URL,
		WorkspaceID: "wsp_1",
		MachineID:   "mch_local",
		APIKey:      "machine-secret",
	}
	require.NoError(t, profileManager.Servers().Add(&server))
	sessions := uisession.NewStore(uisession.Options{})
	sessionID := token.CreateSession()
	expiresAt, ok := token.SessionExpiresAt(sessionID)
	require.True(t, ok)
	_, err := sessions.Select(sessionID, expiresAt, target.ID)
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiagnosticDownloadRoute(router, Dependencies{
		Auth:            authManager,
		Agents:          profileManager.Agents(),
		Servers:         profileManager.Servers(),
		Sessions:        sessions,
		DiagnosticPaths: diagnosticbundle.Paths{ConfigDir: configDir},
	})
	response := performDiagnosticRequest(
		router,
		"/api/internal/machines/mch_local/diagnostics/archive?workspaceId=wsp_1",
		sessionID,
	)
	require.Equal(t, http.StatusOK, response.Code)
	require.Contains(t, diagnosticArchiveEntries(t, response.Body.Bytes()), "manifest.json")
}

func writeMachineDiagnostic(t *testing.T, configDir, agentID, value string) {
	t.Helper()
	dir := filepath.Join(configDir, "agents", agentID, "diagnostics")
	require.NoError(t, os.MkdirAll(dir, 0o755))
	require.NoError(t, os.WriteFile(
		filepath.Join(dir, "opencode-runtime-debug.jsonl"),
		[]byte(value+"\n"),
		0o600,
	))
}

func performDiagnosticRequest(handler http.Handler, path, sessionID string) *httptest.ResponseRecorder {
	request := httptest.NewRequest(http.MethodGet, path, nil)
	if sessionID != "" {
		request.Header.Set("Authorization", "Bearer "+sessionID)
	}
	recorder := httptest.NewRecorder()
	handler.ServeHTTP(recorder, request)
	return recorder
}

func diagnosticArchiveEntries(t *testing.T, data []byte) map[string][]byte {
	t.Helper()
	reader, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	require.NoError(t, err)
	entries := make(map[string][]byte, len(reader.File))
	for _, file := range reader.File {
		source, err := file.Open()
		require.NoError(t, err)
		content, err := io.ReadAll(source)
		require.NoError(t, err)
		require.NoError(t, source.Close())
		entries[file.Name] = content
	}
	return entries
}