package kb_test

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

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/AgentTeam/internal/service/kb"
)

func TestKBHandlerSetsOneAgentKnowledgeBaseAccess(t *testing.T) {
	mux, svc := newKBTestMuxAndServiceWithWiki(t, nil)
	ctx := context.Background()
	base, err := svc.CreateKnowledgeBase(ctx, "wsp_1", "owner", kb.KBMutation{Name: "Runbooks"})
	require.NoError(t, err)

	req := httptest.NewRequest(http.MethodPut, "/api/v1/kb/agents/agt_reader/knowledge-bases/"+base.ID, bytes.NewBufferString(`{"knowledgeBaseId":"`+base.ID+`","accessLevel":"write"}`))
	req.Header.Set("Authorization", "Bearer "+jwt(map[string]any{"sub": "owner", "owned_agent_ids": []string{"agt_reader"}}))
	req.Header.Set("X-Workspace-Id", "wsp_1")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	require.Contains(t, rec.Body.String(), `"accessLevel":"write"`)
}

func TestKBHandlerBulkSetsAgentPermissions(t *testing.T) {
	mux, svc := newKBTestMuxAndServiceWithWiki(t, nil)
	ctx := context.Background()
	keep, err := svc.CreateKnowledgeBase(ctx, "wsp_1", "owner", kb.KBMutation{Name: "Keep"})
	require.NoError(t, err)
	remove, err := svc.CreateKnowledgeBase(ctx, "wsp_1", "owner", kb.KBMutation{Name: "Remove"})
	require.NoError(t, err)
	_, err = svc.GrantAgentKnowledgeBaseAccess(ctx, "wsp_1", "owner", "agt_reader", remove.ID, "read")
	require.NoError(t, err)

	req := httptest.NewRequest(http.MethodPut, "/api/v1/kb/agents/agt_reader/permissions", bytes.NewBufferString(`{"kbIds":["`+keep.ID+`"]}`))
	req.Header.Set("Authorization", "Bearer "+jwt(map[string]any{"sub": "owner"}))
	req.Header.Set("X-Workspace-Id", "wsp_1")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	require.Equal(t, http.StatusOK, rec.Code)
	require.Contains(t, rec.Body.String(), `"agentId":"agt_reader"`)
	require.Contains(t, rec.Body.String(), keep.ID)
	require.Contains(t, rec.Body.String(), remove.ID)

	items, err := svc.ListAgentAccessibleKnowledgeBases(ctx, "wsp_1", "owner", "agt_reader", 1, 20)
	require.NoError(t, err)
	require.Len(t, items, 1)
	require.Equal(t, keep.ID, items[0].ID)
}

func TestKBHandlerBulkSetsAgentPermissionAccessLevels(t *testing.T) {
	mux, svc := newKBTestMuxAndServiceWithWiki(t, nil)
	ctx := context.Background()
	base, err := svc.CreateKnowledgeBase(ctx, "wsp_1", "owner", kb.KBMutation{Name: "Runbooks"})
	require.NoError(t, err)

	req := httptest.NewRequest(http.MethodPut, "/api/v1/kb/agents/agt_reader/permissions", bytes.NewBufferString(`{"knowledgeBases":[{"knowledgeBaseId":"`+base.ID+`","accessLevel":"write"}]}`))
	req.Header.Set("Authorization", "Bearer "+jwt(map[string]any{"sub": "owner"}))
	req.Header.Set("X-Workspace-Id", "wsp_1")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)
	require.Equal(t, http.StatusOK, rec.Code)

	items, err := svc.ListAgentAccessibleKnowledgeBases(ctx, "wsp_1", "owner", "agt_reader", 1, 20)
	require.NoError(t, err)
	require.Len(t, items, 1)
	require.Equal(t, base.ID, items[0].ID)
	require.Equal(t, "write", items[0].AccessLevel)
}

func TestKBHandlerAgentAccessBatchAcceptsTemplateLocale(t *testing.T) {
	mux, _ := newKBTestMuxAndServiceWithWiki(t, nil)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/kb/agents/access:batch", bytes.NewBufferString(`{"agentIds":["agt_reader"],"templateLocale":"zh-CN"}`))
	req.Header.Set("Authorization", "Bearer "+jwt(map[string]any{"sub": "owner"}))
	req.Header.Set("X-Workspace-Id", "wsp_1")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	require.Equal(t, http.StatusOK, rec.Code)
	var body struct {
		Agents map[string][]map[string]any `json:"agents"`
	}
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
	items, ok := body.Agents["agt_reader"]
	require.True(t, ok)
	require.Empty(t, items)
}

func TestKBHandlerBulkAgentPermissionsRequireOwnedAgent(t *testing.T) {
	mux, svc := newKBTestMuxAndServiceWithWiki(t, nil)
	ctx := context.Background()
	base, err := svc.CreateKnowledgeBase(ctx, "wsp_1", "owner", kb.KBMutation{Name: "Runbooks"})
	require.NoError(t, err)

	req := httptest.NewRequest(http.MethodPut, "/api/v1/kb/agents/agt_unowned/permissions", bytes.NewBufferString(`{"kbIds":["`+base.ID+`"]}`))
	req.Header.Set("Authorization", "Bearer "+jwt(map[string]any{"sub": "owner"}))
	req.Header.Set("X-Workspace-Id", "wsp_1")
	rec := httptest.NewRecorder()
	mux.ServeHTTP(rec, req)

	require.Equal(t, http.StatusForbidden, rec.Code)
	require.Contains(t, rec.Body.String(), "agent must be owned")
}
