package knowledgeqa

import (
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"

	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func TestStorePersistsOpaqueRootHandleAuthority(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:kqa_root_handles?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&RootHandleRecord{}))
	store := NewRuntime(db)
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	handle := protocol.RootHandle{
		ID: "kqr_1", OrganizationID: "org_1", WorkspaceID: "wsp_1",
		RequesterID: "usr_1", TurnID: "kqt_1", KnowledgeBaseID: "kb_1",
		SourceID: "src_1", RootKind: "primary",
		SourceName: "tidb", DefaultRef: "master", ResolvedRef: "release-7.5",
		ResolvedRefKind: "branch",
		RepoRole:        "primary", Revision: "sha-fixed",
		RepositoryRef: "internal-repo-ref", FreshnessStatus: "current",
		ExpiresAt: now.Add(time.Minute), CreatedAt: now,
	}
	require.NoError(t, store.PutRootHandles(context.Background(), []protocol.RootHandle{handle}))
	retryHandle := handle
	retryHandle.ID = "kqr_retry"
	require.NoError(t, store.PutRootHandles(context.Background(), []protocol.RootHandle{retryHandle}))
	principal := protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
	}
	loaded, err := store.LoadRootHandles(context.Background(), principal, []string{handle.ID})
	require.NoError(t, err)
	require.Equal(t, []protocol.RootHandle{handle}, loaded)
	queryCount := 0
	require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:count-root-handle-batch", func(*gorm.DB) {
		queryCount++
	}))
	loadedBatch, err := store.LoadRootHandles(context.Background(), principal, []string{handle.ID, retryHandle.ID})
	require.NoError(t, err)
	require.Len(t, loadedBatch, 2)
	require.Equal(t, 1, queryCount, "batch size must not change root-handle query count")
	queryCount = 0
	foreignPrincipal := principal
	foreignPrincipal.OrganizationID = "org_other"
	loadedBatch, err = store.LoadRootHandles(context.Background(), foreignPrincipal, []string{handle.ID, retryHandle.ID})
	require.NoError(t, err)
	require.Empty(t, loadedBatch)
	require.Equal(t, 1, queryCount)

	require.NoError(t, store.DeleteRootHandles(context.Background(), protocol.Principal{
		OrganizationID: "org_other", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
	}, []string{handle.ID}))
	loadedBatch, err = store.LoadRootHandles(context.Background(), principal, []string{handle.ID, retryHandle.ID})
	require.NoError(t, err)
	require.Len(t, loadedBatch, 2)

	require.NoError(t, store.DeleteRootHandles(context.Background(), protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
	}, []string{handle.ID}))
	loadedBatch, err = store.LoadRootHandles(context.Background(), principal, []string{handle.ID, retryHandle.ID})
	require.NoError(t, err)
	require.Equal(t, []protocol.RootHandle{retryHandle}, loadedBatch,
		"cleanup from an older attempt must preserve retry handles for the same turn")
}

func TestStoreExpiresHandlesInBoundedBatches(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:kqa_root_expiry?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&RootHandleRecord{}))
	store := NewRuntime(db)
	now := time.Date(2026, 7, 25, 10, 0, 0, 0, time.UTC)
	for index, expiry := range []time.Time{now.Add(-time.Minute), now.Add(-time.Second), now.Add(time.Minute)} {
		handle := protocol.RootHandle{
			ID:             []string{"kqr_1", "kqr_2", "kqr_3"}[index],
			OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1",
			TurnID: "kqt_1", KnowledgeBaseID: "kb_1", SourceID: "src_1",
			RootKind: "primary", RepoRole: "primary",
			Revision: "sha", RepositoryRef: "repo", ExpiresAt: expiry, CreatedAt: now.Add(-time.Hour),
		}
		require.NoError(t, store.PutRootHandles(context.Background(), []protocol.RootHandle{handle}))
	}
	deleted, err := store.DeleteExpiredRootHandles(context.Background(), now, 1)
	require.NoError(t, err)
	require.EqualValues(t, 1, deleted)
	deleted, err = store.DeleteExpiredRootHandles(context.Background(), now, 10)
	require.NoError(t, err)
	require.EqualValues(t, 1, deleted)
	loaded, err := store.LoadRootHandles(context.Background(), protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
	}, []string{"kqr_3"})
	require.NoError(t, err)
	require.Len(t, loaded, 1)
}

func TestStoreExpiresHandlesWithinWorkspaceScopeBeforeLimit(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:kqa_root_scoped_expiry?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&RootHandleRecord{}))
	store := NewRuntime(db)
	now := time.Date(2026, 7, 25, 11, 0, 0, 0, time.UTC)
	for _, handle := range []protocol.RootHandle{
		{ID: "kqr_foreign", OrganizationID: "org_1", WorkspaceID: "wsp_foreign", RequesterID: "usr_1", TurnID: "kqt_1", KnowledgeBaseID: "kb_1", SourceID: "src_1", RootKind: "primary", RepoRole: "primary", Revision: "sha", RepositoryRef: "repo", ExpiresAt: now.Add(-time.Minute), CreatedAt: now.Add(-time.Hour)},
		{ID: "kqr_local", OrganizationID: "org_1", WorkspaceID: "wsp_local", RequesterID: "usr_1", TurnID: "kqt_2", KnowledgeBaseID: "kb_2", SourceID: "src_2", RootKind: "primary", RepoRole: "primary", Revision: "sha", RepositoryRef: "repo", ExpiresAt: now.Add(-time.Second), CreatedAt: now.Add(-time.Hour)},
	} {
		require.NoError(t, store.PutRootHandles(context.Background(), []protocol.RootHandle{handle}))
	}

	deleted, err := store.DeleteExpiredRootHandles(context.Background(), now, 1, []string{"wsp_local"})
	require.NoError(t, err)
	require.EqualValues(t, 1, deleted)

	var foreignCount int64
	require.NoError(t, db.Model(&RootHandleRecord{}).Where("id = ?", "kqr_foreign").Count(&foreignCount).Error)
	require.EqualValues(t, 1, foreignCount)
}

func TestStorePersistsMonotonicOwnerBoundAccountingCheckpoints(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:kqa_accounting?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&AccountingCheckpointRecord{}))
	store := NewRuntime(db)
	now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
	principal := protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
	}
	checkpoint := protocol.AccountingCheckpoint{
		Principal: principal, AttemptID: "kqa_1", Generation: 1,
		State: protocol.AccountingCheckpointRunning,
		Usage: map[string]any{"input_tokens": float64(10)}, SettlementRefs: []string{"usage_1"},
		ExpiresAt: now.Add(time.Hour), UpdatedAt: now,
	}
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), checkpoint))
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), checkpoint))
	conflictingReplay := checkpoint
	conflictingReplay.Usage = map[string]any{"input_tokens": float64(11)}
	require.ErrorContains(t, store.PutAccountingCheckpoint(context.Background(), conflictingReplay), "generation conflicts")
	rebound := checkpoint
	rebound.Principal.RequesterID = "usr_other"
	require.ErrorContains(t, store.PutAccountingCheckpoint(context.Background(), rebound), "authority conflicts")

	stale := checkpoint
	stale.Generation = 0
	stale.Usage = map[string]any{}
	stale.SettlementRefs = nil
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), stale))

	next := checkpoint
	next.Generation = 2
	next.Usage = map[string]any{"input_tokens": float64(18), "output_tokens": float64(4)}
	next.SettlementRefs = []string{"usage_1", "usage_2"}
	next.UpdatedAt = now.Add(time.Second)
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), next))

	regressed := next
	regressed.Generation = 3
	regressed.Usage = map[string]any{"input_tokens": float64(17)}
	require.ErrorContains(t, store.PutAccountingCheckpoint(context.Background(), regressed), "cannot regress")
	reordered := next
	reordered.Generation = 3
	reordered.SettlementRefs = []string{"usage_2", "usage_1"}
	require.ErrorContains(t, store.PutAccountingCheckpoint(context.Background(), reordered), "cannot be reordered")

	terminal := next
	terminal.State = protocol.AccountingCheckpointTerminal
	terminal.UpdatedAt = now.Add(2 * time.Second)
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), terminal))
	loaded, found, err := store.LoadAccountingCheckpoint(context.Background(), principal, "kqa_1", now)
	require.NoError(t, err)
	require.True(t, found)
	require.Equal(t, 2, loaded.Generation)
	require.Equal(t, protocol.AccountingCheckpointTerminal, loaded.State)
	require.Equal(t, []string{"usage_1", "usage_2"}, loaded.SettlementRefs)
	require.Equal(t, float64(18), loaded.Usage["input_tokens"])
	afterTerminal := terminal
	afterTerminal.Generation = 3
	afterTerminal.Usage = map[string]any{"input_tokens": float64(19), "output_tokens": float64(4)}
	afterTerminal.SettlementRefs = []string{"usage_1", "usage_2", "usage_3"}
	require.ErrorContains(t, store.PutAccountingCheckpoint(context.Background(), afterTerminal), "terminal accounting checkpoint cannot advance")
	loaded, found, err = store.LoadAccountingCheckpoint(context.Background(), principal, "kqa_1", now)
	require.NoError(t, err)
	require.True(t, found)
	require.Equal(t, 2, loaded.Generation)
	require.Equal(t, []string{"usage_1", "usage_2"}, loaded.SettlementRefs)

	otherOwner := principal
	otherOwner.RequesterID = "usr_other"
	_, found, err = store.LoadAccountingCheckpoint(context.Background(), otherOwner, "kqa_1", now)
	require.NoError(t, err)
	require.False(t, found)

	deleted, err := store.DeleteExpiredAccountingCheckpoints(context.Background(), now.Add(2*time.Hour), 1)
	require.NoError(t, err)
	require.EqualValues(t, 1, deleted)
	_, found, err = store.LoadAccountingCheckpoint(context.Background(), principal, "kqa_1", now.Add(2*time.Hour))
	require.NoError(t, err)
	require.False(t, found)
}

func TestStorePersistsAccountingCheckpointPastLegacy64Boundaries(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:kqa_accounting_legacy_limits?mode=memory&cache=shared"), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(&AccountingCheckpointRecord{}))
	store := NewRuntime(db)
	now := time.Date(2026, 8, 2, 5, 0, 0, 0, time.UTC)
	principal := protocol.Principal{
		OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_65",
	}
	refs := make([]string, 65)
	for index := range refs {
		prefix := fmt.Sprintf("usage_%03d_", index)
		refs[index] = prefix + strings.Repeat("x", 255-len(prefix))
	}
	encodedRefs, err := json.Marshal(refs)
	require.NoError(t, err)
	require.Greater(t, len(encodedRefs), 16<<10, "regression must cross the former lineage envelope")

	checkpoint := protocol.AccountingCheckpoint{
		Principal: principal, AttemptID: "kqa_65", Generation: 65,
		State: protocol.AccountingCheckpointRunning,
		Usage: map[string]any{"input_tokens": float64(65)}, SettlementRefs: refs,
		ExpiresAt: now.Add(time.Hour), UpdatedAt: now,
	}
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), checkpoint))

	nextRefPrefix := "usage_065_"
	next := checkpoint
	next.Generation = 66
	next.Usage = map[string]any{"input_tokens": float64(66)}
	next.SettlementRefs = append(append([]string(nil), refs...), nextRefPrefix+strings.Repeat("x", 255-len(nextRefPrefix)))
	next.UpdatedAt = now.Add(time.Second)
	require.NoError(t, store.PutAccountingCheckpoint(context.Background(), next))

	loaded, found, err := store.LoadAccountingCheckpoint(context.Background(), principal, checkpoint.AttemptID, now)
	require.NoError(t, err)
	require.True(t, found)
	require.Equal(t, 66, loaded.Generation)
	require.Equal(t, next.SettlementRefs, loaded.SettlementRefs)
	require.Equal(t, float64(66), loaded.Usage["input_tokens"])
}

func TestAccountingCheckpointDocumentSizeBoundaries(t *testing.T) {
	require.NoError(t, validateAccountingDocumentSizes(
		make([]byte, protocol.MaxAccountingUsageJSONBytes),
		make([]byte, protocol.MaxAccountingSettlementLineageJSONBytes),
		protocol.MaxAccountingSettlementReferences,
	))

	tests := []struct {
		name           string
		usageBytes     int
		lineageBytes   int
		referenceCount int
		wantError      string
	}{
		{
			name: "usage", usageBytes: protocol.MaxAccountingUsageJSONBytes + 1,
			lineageBytes: 2, referenceCount: 0, wantError: "usage is invalid or too large",
		},
		{
			name: "encoded lineage", usageBytes: 2,
			lineageBytes:   protocol.MaxAccountingSettlementLineageJSONBytes + 1,
			referenceCount: 1, wantError: "storage boundary",
		},
		{
			name: "reference count", usageBytes: 2, lineageBytes: 2,
			referenceCount: protocol.MaxAccountingSettlementReferences + 1,
			wantError:      "storage boundary",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			err := validateAccountingDocumentSizes(
				make([]byte, test.usageBytes), make([]byte, test.lineageBytes), test.referenceCount,
			)
			require.ErrorContains(t, err, test.wantError)
		})
	}
}

func TestAccountingCheckpointRejectsOversizedSettlementReference(t *testing.T) {
	_, err := normalizeAccountingSettlementRefs([]string{
		strings.Repeat("x", protocol.MaxAccountingSettlementReferenceBytes+1),
	})
	require.ErrorContains(t, err, "storage boundary")
}

func TestAccountingCheckpointRecordRejectsSettlementReferenceCountOutsideStorageEnvelope(t *testing.T) {
	refs := make([]string, protocol.MaxAccountingSettlementReferences+1)
	for index := range refs {
		refs[index] = fmt.Sprintf("usage_%d", index)
	}
	now := time.Date(2026, 8, 3, 13, 0, 0, 0, time.UTC)
	_, err := accountingRecordFromCheckpoint(protocol.AccountingCheckpoint{
		Principal: protocol.Principal{
			OrganizationID: "org_1", WorkspaceID: "wsp_1", RequesterID: "usr_1", TurnID: "kqt_1",
		},
		AttemptID: "kqa_1", Generation: 1, State: protocol.AccountingCheckpointRunning,
		SettlementRefs: refs, ExpiresAt: now.Add(time.Hour), UpdatedAt: now,
	})
	require.ErrorContains(t, err, "storage boundary")
}

func TestAccountingCheckpointLineageUsesLongTextOnMySQL(t *testing.T) {
	db, err := gorm.Open(mysql.New(mysql.Config{
		DSN:                       "gorm:gorm@tcp(localhost:9910)/gorm?charset=utf8mb4&parseTime=True&loc=Local",
		SkipInitializeWithVersion: true,
	}), &gorm.Config{DisableAutomaticPing: true})
	require.NoError(t, err)

	statement := &gorm.Statement{DB: db}
	require.NoError(t, statement.Parse(&AccountingCheckpointRecord{}))
	field := statement.Schema.LookUpField("SettlementRefsJSON")
	require.NotNil(t, field)
	require.Equal(t, "LONGTEXT NOT NULL", db.Migrator().FullDataTypeOf(field).SQL)
}