package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/application/repository"
	werrors "github.com/Tencent/WeKnora/internal/errors"
	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type classificationTestModel struct {
	*autoTagChatModel
	call func(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error)
}

func (m classificationTestModel) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
	return m.call(ctx, messages, options)
}

func classificationFixture(t *testing.T) (*KnowledgeClassificationService, context.Context, *autoTagChatModel) {
	t.Helper()
	var db *gorm.DB
	var e error
	if dsn := os.Getenv("CLASSIFICATION_TEST_POSTGRES_DSN"); dsn != "" {
		admin, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
		require.NoError(t, err)
		schema := "classification_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
		require.NoError(t, admin.Exec("CREATE SCHEMA "+schema).Error)
		t.Cleanup(func() {
			require.NoError(t, admin.Exec("DROP SCHEMA "+schema+" CASCADE").Error)
			raw, err := admin.DB()
			require.NoError(t, err)
			raw.Close()
		})
		db, e = gorm.Open(postgres.Open(dsn+" search_path="+schema), &gorm.Config{})
	} else {
		db, e = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	}
	require.NoError(t, e)
	raw, e := db.DB()
	require.NoError(t, e)
	if db.Dialector.Name() == "sqlite" {
		raw.SetMaxOpenConns(1)
	} else {
		raw.SetMaxOpenConns(12)
	}
	t.Cleanup(func() { raw.Close() })
	require.NoError(t, db.AutoMigrate(&types.User{}, &types.Tenant{}, &types.TenantMember{}, &types.KnowledgeBase{}, &types.Knowledge{}, &types.Chunk{}, &types.KnowledgeTag{}, &types.KnowledgeTagRelation{}, &types.AuditLog{}, &KnowledgeClassificationJob{}, &KnowledgeClassificationItem{}))
	require.NoError(t, db.Exec("CREATE UNIQUE INDEX idx_classification_active ON knowledge_classification_jobs(tenant_id,owner_id,knowledge_base_id) WHERE status='running'").Error)
	require.NoError(t, db.Create(&types.Tenant{ID: 1, Name: "test"}).Error)
	require.NoError(t, db.Create(&types.User{ID: "alice", TenantID: 1, Email: "alice@example.test", Username: "Alice", IsActive: true}).Error)
	require.NoError(t, db.Create(&types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleContributor, Status: types.TenantMemberStatusActive}).Error)
	kb := &types.KnowledgeBase{ID: "kb", TenantID: 1, CreatorID: "alice", Type: "document", SummaryModelID: "model"}
	require.NoError(t, db.Create(kb).Error)
	require.NoError(t, db.Create(&types.Knowledge{ID: "doc", TenantID: 1, KnowledgeBaseID: "kb", Title: "product document", ParseStatus: types.ParseStatusCompleted}).Error)
	require.NoError(t, db.Create(&types.Chunk{ID: "chunk", TenantID: 1, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: "Product education: " + strings.Repeat("x", 100000), ChunkType: types.ChunkTypeText, IndexStatus: "ready"}).Error)
	for _, tag := range []*types.KnowledgeTag{{ID: "new", TenantID: 1, KnowledgeBaseID: "kb", Name: "产品 / 使用说明 / 完整分类路径"}, {ID: "old", TenantID: 1, KnowledgeBaseID: "kb", Name: "现有分类"}} {
		require.NoError(t, db.Create(tag).Error)
	}
	model := &autoTagChatModel{response: `{"matches":[{"index":1,"confidence":0.95,"reason":"资料介绍了产品的使用方法，适合归入使用说明。"}]}`}
	s := &KnowledgeClassificationService{db: db, models: &stubModelService{chatModel: model}, bases: versionCopyKBService{bases: map[string]*types.KnowledgeBase{"kb": kb}}}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "alice")
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleContributor)
	return s, ctx, model
}
func TestKnowledgeClassificationRequiresExplicitApplyAndPreservesTags(t *testing.T) {
	s, ctx, model := classificationFixture(t)
	repo := repository.NewKnowledgeRepository(s.db)
	require.NoError(t, repo.SetKnowledgeTags(ctx, "doc", []string{"old"}))
	job, e := s.Create(ctx, "kb", "", []string{"doc"})
	require.NoError(t, e)
	_, e = s.Create(ctx, "kb", "", []string{"doc"})
	require.ErrorContains(t, e, "已有分类任务")
	require.NoError(t, s.RunNext(context.Background()))
	jobs, e := s.List(ctx, "kb")
	require.NoError(t, e)
	require.Len(t, jobs, 1)
	require.Equal(t, "completed", jobs[0].Status)
	require.Equal(t, 1, jobs[0].Processed)
	item := jobs[0].Items[0]
	require.Equal(t, "suggested", item.Status)
	var suggestions []ClassificationSuggestion
	require.NoError(t, json.Unmarshal(item.Suggestions, &suggestions))
	require.Equal(t, "产品 / 使用说明 / 完整分类路径", suggestions[0].Name)
	tags, e := repo.GetKnowledgeTags(ctx, []string{"doc"})
	require.NoError(t, e)
	require.Len(t, tags["doc"], 1)
	require.Less(t, len(model.messages[1].Content), 5000, "must not read whole large chunk into model")
	require.NoError(t, s.Apply(ctx, "kb", job.ID, []string{item.ID}))
	require.NoError(t, s.Apply(ctx, "kb", job.ID, []string{item.ID}))
	tags, e = repo.GetKnowledgeTags(ctx, []string{"doc"})
	require.NoError(t, e)
	require.Len(t, tags["doc"], 2)
	var count int64
	require.NoError(t, s.db.Model(&types.AuditLog{}).Where("action=?", "knowledge.classification_applied").Count(&count).Error)
	require.Positive(t, count)
}
func TestKnowledgeClassificationRejectsStaleContentCategoryAndOtherUser(t *testing.T) {
	for _, change := range []string{"content", "category", "user", "permission"} {
		t.Run(change, func(t *testing.T) {
			s, ctx, _ := classificationFixture(t)
			job, e := s.Create(ctx, "kb", "model", []string{"doc"})
			require.NoError(t, e)
			require.NoError(t, s.RunNext(context.Background()))
			jobs, e := s.List(ctx, "kb")
			require.NoError(t, e)
			switch change {
			case "content":
				require.NoError(t, s.db.Model(&types.Chunk{}).Where("id='chunk'").Updates(map[string]any{"content": "new text", "content_revision": 1}).Error)
			case "category":
				require.NoError(t, s.db.Model(&types.KnowledgeTag{}).Where("id='new'").Update("name", "changed").Error)
			case "user":
				ctx = context.WithValue(ctx, types.UserIDContextKey, "bob")
				ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin)
			case "permission":
				ctx = accesspolicy.WithSnapshot(ctx, &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.read": true}})
			}
			require.Error(t, s.Apply(ctx, "kb", job.ID, []string{jobs[0].Items[0].ID}))
			tags, e := repository.NewKnowledgeRepository(s.db).GetKnowledgeTags(ctx, []string{"doc"})
			require.NoError(t, e)
			require.Empty(t, tags["doc"])
		})
	}
}
func TestKnowledgeClassificationWorkerRevocationFailureAndRetry(t *testing.T) {
	s, ctx, model := classificationFixture(t)
	job, e := s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, e)
	require.NoError(t, s.db.Model(&types.User{}).Where("id='alice'").Update("is_active", false).Error)
	require.NoError(t, s.RunNext(context.Background()))
	jobs, e := s.List(ctx, "kb")
	require.NoError(t, e)
	require.Equal(t, "completed_with_errors", jobs[0].Status)
	require.Empty(t, model.messages)
	require.NoError(t, s.db.Model(&types.User{}).Where("id='alice'").Update("is_active", true).Error)
	require.NoError(t, s.Retry(ctx, "kb", job.ID))
	require.NoError(t, s.RunNext(context.Background()))
	jobs, e = s.List(ctx, "kb")
	require.NoError(t, e)
	require.Equal(t, "completed", jobs[0].Status)
	require.Equal(t, 1, jobs[0].Processed)
	require.Zero(t, jobs[0].Failed)
}
func TestKnowledgeClassificationRejectsInventedTagsAndResumesExpiredLease(t *testing.T) {
	s, ctx, model := classificationFixture(t)
	model.response = `{"matches":[{"index":999,"confidence":1,"reason":"invented"},{"index":1,"confidence":0.1,"reason":"uncertain"}]}`
	job, e := s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, e)
	require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("job_id=?", job.ID).Updates(map[string]any{"status": "running", "lease_token": "old-worker", "lease_until": time.Now().Add(-time.Minute)}).Error)
	require.NoError(t, s.RunNext(context.Background()))
	jobs, e := s.List(ctx, "kb")
	require.NoError(t, e)
	require.Equal(t, "no_match", jobs[0].Items[0].Status)
	require.Error(t, s.Apply(ctx, "kb", job.ID, []string{jobs[0].Items[0].ID}))
}

func TestKnowledgeClassificationOldLeaseCannotOverwriteRecoveredResult(t *testing.T) {
	s, ctx, _ := classificationFixture(t)
	job, err := s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, err)
	old := job.Items[0]
	old.LeaseToken = "expired-worker"
	require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("id=?", old.ID).Updates(map[string]any{"status": "running", "lease_token": old.LeaseToken, "lease_until": time.Now().Add(-time.Minute)}).Error)
	require.NoError(t, s.RunNext(context.Background()))
	before, err := s.List(ctx, "kb")
	require.NoError(t, err)
	require.NoError(t, s.finish(context.Background(), *job, old, "failed", "late failure", "", nil))
	after, err := s.List(ctx, "kb")
	require.NoError(t, err)
	require.Equal(t, before, after)
}

func TestKnowledgeClassificationConcurrentCompletionCounts(t *testing.T) {
	s, ctx, _ := classificationFixture(t)
	job, err := s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, err)
	const total = 8
	items := make([]KnowledgeClassificationItem, total)
	for i := range items {
		items[i] = KnowledgeClassificationItem{ID: fmt.Sprintf("parallel-%d", i), JobID: job.ID, KnowledgeID: fmt.Sprintf("parallel-doc-%d", i), Status: "running", LeaseToken: fmt.Sprintf("lease-%d", i), Suggestions: types.JSON(`[]`)}
	}
	require.NoError(t, s.db.Where("job_id=?", job.ID).Delete(&KnowledgeClassificationItem{}).Error)
	require.NoError(t, s.db.Create(&items).Error)
	require.NoError(t, s.db.Model(job).Update("total", total).Error)
	var wg sync.WaitGroup
	failures := make(chan error, total)
	for _, item := range items {
		wg.Add(1)
		go func(item KnowledgeClassificationItem) {
			defer wg.Done()
			failures <- s.finish(context.Background(), *job, item, "no_match", "", "", nil)
		}(item)
	}
	wg.Wait()
	close(failures)
	for err := range failures {
		require.NoError(t, err)
	}
	require.NoError(t, s.db.First(job, "id=?", job.ID).Error)
	require.Equal(t, total, job.Processed)
	require.Equal(t, "completed", job.Status)
}

func TestKnowledgeClassificationMasksUntrustedErrors(t *testing.T) {
	const secret = "sk-provider-private-value response-body password=secret"
	for _, source := range []error{errors.New(secret), werrors.NewBadRequestError(secret).WithDetails(secret), werrors.NewForbiddenError(secret)} {
		t.Run(fmt.Sprintf("%T-%s", source, classificationErrorCode(source)), func(t *testing.T) {
			s, ctx, model := classificationFixture(t)
			s.models = &stubModelService{chatModel: classificationTestModel{autoTagChatModel: model, call: func(_ context.Context, _ []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
				require.True(t, json.Valid(options.Format), "JSON response mode required")
				return nil, source
			}}}
			_, err := s.Create(ctx, "kb", "model", []string{"doc"})
			require.NoError(t, err)
			require.NoError(t, s.RunNext(context.Background()))
			jobs, err := s.List(ctx, "kb")
			require.NoError(t, err)
			require.Equal(t, "failed", jobs[0].Items[0].Status)
			body, err := json.Marshal(jobs)
			require.NoError(t, err)
			require.NotContains(t, string(body), secret)
			require.NotContains(t, KnowledgeClassificationPublicError(source).Error(), secret)
			require.Nil(t, KnowledgeClassificationPublicError(source).Details)
		})
	}
	require.Equal(t, "已知的安全说明", KnowledgeClassificationPublicError(classificationConflict("已知的安全说明")).Message)
}

func TestKnowledgeClassificationRevokedDuringModelCall(t *testing.T) {
	for _, revoke := range []string{"membership", "permission", "hidden"} {
		t.Run(revoke, func(t *testing.T) {
			s, ctx, model := classificationFixture(t)
			revoked := false
			s.snapshot = func(context.Context) (*accesspolicy.Snapshot, error) {
				snapshot := &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"knowledge.write": !revoked || revoke != "permission"}}
				if revoked && revoke == "hidden" {
					snapshot.Hidden = map[string]bool{"knowledge:kb": true}
				}
				return snapshot, nil
			}
			s.models = &stubModelService{chatModel: classificationTestModel{autoTagChatModel: model, call: func(_ context.Context, _ []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) {
				revoked = true
				if revoke == "membership" {
					require.NoError(t, s.db.Model(&types.TenantMember{}).Where("user_id='alice'").Update("status", "inactive").Error)
				}
				return &types.ChatResponse{Content: model.response}, nil
			}}}
			_, err := s.Create(ctx, "kb", "model", []string{"doc"})
			require.NoError(t, err)
			require.NoError(t, s.RunNext(context.Background()))
			jobs, err := s.List(ctx, "kb")
			require.NoError(t, err)
			require.Equal(t, "failed", jobs[0].Items[0].Status)
			require.JSONEq(t, `[]`, string(jobs[0].Items[0].Suggestions))
		})
	}
}

func TestKnowledgeClassificationRetryBlockedByAnotherActiveJob(t *testing.T) {
	s, ctx, _ := classificationFixture(t)
	job, err := s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, err)
	require.NoError(t, s.db.Model(&KnowledgeClassificationItem{}).Where("job_id=?", job.ID).Update("status", "failed").Error)
	require.NoError(t, s.db.Model(job).Updates(map[string]any{"status": "completed_with_errors", "processed": 1, "failed": 1}).Error)
	_, err = s.Create(ctx, "kb", "model", []string{"doc"})
	require.NoError(t, err)
	require.ErrorContains(t, s.Retry(ctx, "kb", job.ID), "已有分类任务")
	var item KnowledgeClassificationItem
	require.NoError(t, s.db.First(&item, "job_id=?", job.ID).Error)
	require.Equal(t, "failed", item.Status)
}