package workbench

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func workbenchContext(tenant uint64, user string, role types.TenantRole) context.Context {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, tenant)
	ctx = context.WithValue(ctx, types.UserIDContextKey, user)
	return context.WithValue(ctx, types.TenantRoleContextKey, role)
}

func sqliteWorkbench(t *testing.T) *Service {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "workbench.db")+"?_busy_timeout=5000&_journal_mode=WAL"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = sqlDB.Close() })
	if err = db.AutoMigrate(&Resource{}, &Job{}, &Revision{}, &Event{}, &Consumption{}, &Balance{}, &BalanceHold{}, &Transaction{}, &Usage{}, &UsageQuota{}, &OperationLog{}, &PendingMemory{}); err != nil {
		t.Fatal(err)
	}
	return NewService(db)
}

func seedResource(t *testing.T, s *Service, tenant uint64, owner, kind string) *Resource {
	t.Helper()
	r := &Resource{ID: uuid.NewString(), TenantID: tenant, OwnerID: owner, Kind: kind, Name: "original", Status: "draft", Version: 1, Data: JSON{}}
	if err := s.db.Create(r).Error; err != nil {
		t.Fatal(err)
	}
	return r
}

func TestSaveRequiresCurrentLeaseAndScope(t *testing.T) {
	for _, state := range []string{"valid", "expired", "token", "status", "resource", "owner", "tenant"} {
		t.Run(state, func(t *testing.T) {
			s := sqliteWorkbench(t)
			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
			r := seedResource(t, s, 1, "alice", "training-maps")
			until := time.Now().Add(time.Minute)
			j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "lease", LeaseUntil: &until}
			switch state {
			case "expired":
				until = time.Now().Add(-time.Second)
			case "token":
				j.LeaseToken = "other"
			case "status":
				j.Status = "failed"
			case "resource":
				j.ResourceID = uuid.NewString()
			case "owner":
				j.OwnerID = "bob"
			case "tenant":
				j.TenantID = 2
			}
			if err := s.db.Create(&j).Error; err != nil {
				t.Fatal(err)
			}
			ctx = context.WithValue(ctx, leaseContextKey{}, jobLease{j.ID, "lease"})
			r.Name = "generated"
			err := s.save(ctx, r, 1)
			if state == "valid" {
				if err != nil {
					t.Fatal(err)
				}
			} else if !errors.Is(err, ErrConflict) {
				t.Fatalf("invalid lease saved: %v", err)
			}
			var stored Resource
			if err := s.db.First(&stored, "id = ?", r.ID).Error; err != nil {
				t.Fatal(err)
			}
			var revisions int64
			s.db.Model(&Revision{}).Count(&revisions)
			if state == "valid" && (stored.Version != 2 || revisions != 2) {
				t.Fatal("valid result must retain the original and create the next revision")
			}
			if state != "valid" && (stored.Version != 1 || stored.Name != "original" || revisions != 0) {
				t.Fatal("failed lease changed persisted state")
			}
		})
	}
}

func TestManualSaveRejectsActiveJobsAndUnscopedRows(t *testing.T) {
	s := sqliteWorkbench(t)
	s.agents = generationAgentStub{}
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	r := seedResource(t, s, 1, "alice", "training-maps")
	r.Data = asJSON(MapData{Topic: "topic", AgentID: "author", DaysCount: 1})
	if err := s.db.Save(r).Error; err != nil {
		t.Fatal(err)
	}
	if _, err := s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err != nil {
		t.Fatal(err)
	}
	r.Name = "changed"
	if err := s.save(ctx, r, 1); !errors.Is(err, ErrConflict) {
		t.Fatalf("queued job must block edit: %v", err)
	}
	for _, c := range []context.Context{workbenchContext(1, "bob", types.TenantRoleContributor), workbenchContext(2, "alice", types.TenantRoleAdmin)} {
		if err := s.save(c, r, 1); err == nil {
			t.Fatal("cross-owner or cross-tenant write succeeded")
		}
		if _, err := s.submit(c, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err == nil {
			t.Fatal("cross-owner or cross-tenant submit succeeded")
		}
	}
}

func TestEmptyClaimDoesNotLogRecordNotFound(t *testing.T) {
	s := sqliteWorkbench(t)
	var output bytes.Buffer
	s.db = s.db.Session(&gorm.Session{Logger: logger.New(log.New(&output, "", 0), logger.Config{LogLevel: logger.Warn})})
	j, err := s.claim(context.Background())
	if err != nil || j != nil || strings.Contains(output.String(), "record not found") {
		t.Fatalf("empty queue must be quiet: job=%v err=%v log=%s", j, err, output.String())
	}
}

func TestClaimRequiresAffectedRow(t *testing.T) {
	s := sqliteWorkbench(t)
	j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: uuid.NewString(), Kind: "training-maps", Status: "queued"}
	if err := s.db.Create(&j).Error; err != nil {
		t.Fatal(err)
	}
	if err := s.db.Exec("CREATE TRIGGER refuse_claim BEFORE UPDATE ON workbench_jobs BEGIN SELECT RAISE(IGNORE); END;").Error; err != nil {
		t.Fatal(err)
	}
	claimed, err := s.claim(context.Background())
	if claimed != nil || !errors.Is(err, ErrConflict) {
		t.Fatalf("unclaimed job must not run: job=%v err=%v", claimed, err)
	}
}

func TestExpiredLeaseCannotPublishProgress(t *testing.T) {
	s := sqliteWorkbench(t)
	until := time.Now().Add(-time.Second)
	j := Job{ID: uuid.NewString(), Status: "running", LeaseToken: "lease", LeaseUntil: &until}
	if err := s.db.Create(&j).Error; err != nil {
		t.Fatal(err)
	}
	ctx := context.WithValue(context.Background(), leaseContextKey{}, jobLease{j.ID, j.LeaseToken})
	if err := s.progress(ctx, j.ID, 90, "done"); !errors.Is(err, ErrConflict) {
		t.Fatalf("expired progress update allowed: %v", err)
	}
}

func TestWorkerRejectsSuspendedMembership(t *testing.T) {
	s := sqliteWorkbench(t)
	if err := s.db.AutoMigrate(&types.TenantMember{}); err != nil {
		t.Fatal(err)
	}
	member := types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleAdmin, Status: types.TenantMemberStatusSuspended}
	if err := s.db.Create(&member).Error; err != nil {
		t.Fatal(err)
	}
	r := seedResource(t, s, 1, "alice", "training-maps")
	until := time.Now().Add(time.Minute)
	j := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "valid-lease", LeaseUntil: &until, Request: asJSON(ActionRequest{Version: 1})}
	if err := s.db.Create(&j).Error; err != nil {
		t.Fatal(err)
	}
	s.run(context.Background(), &j)
	var stored Job
	if err := s.db.First(&stored, "id = ?", j.ID).Error; err != nil {
		t.Fatal(err)
	}
	if stored.Status != "failed" || stored.Error != ErrForbidden.Error() {
		t.Fatalf("suspended user executed job: %+v", stored)
	}
}

func TestResourceJobListIsScopedAndDoesNotExposeLease(t *testing.T) {
	s := sqliteWorkbench(t)
	r := seedResource(t, s, 1, "alice", "training-maps")
	jobs := []Job{
		{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "private-token", Request: asJSON(ActionRequest{Advice: "private-advice"})},
		{ID: uuid.NewString(), TenantID: 1, OwnerID: "bob", ResourceID: r.ID, Kind: r.Kind, Status: "queued"},
		{ID: uuid.NewString(), TenantID: 2, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "queued"},
	}
	if err := s.db.Create(&jobs).Error; err != nil {
		t.Fatal(err)
	}
	for _, role := range []types.TenantRole{types.TenantRoleContributor, types.TenantRoleAdmin} {
		w := httptest.NewRecorder()
		c, _ := gin.CreateTestContext(w)
		c.Params = gin.Params{{Key: "kind", Value: r.Kind}, {Key: "id", Value: r.ID}}
		c.Request = httptest.NewRequest(http.MethodGet, "/", nil).WithContext(workbenchContext(1, "alice", role))
		NewHandler(s).Jobs(c)
		var result struct {
			Total int64 `json:"total"`
			Data  []Job `json:"data"`
		}
		if err := json.Unmarshal(w.Body.Bytes(), &result); err != nil {
			t.Fatal(err)
		}
		want := int64(1)
		if role == types.TenantRoleAdmin {
			want = 2
		}
		if w.Code != 200 || result.Total != want || int64(len(result.Data)) != want || strings.Contains(w.Body.String(), "private-token") || strings.Contains(w.Body.String(), "private-advice") {
			t.Fatalf("job list leaked scope or internal fields: %d %s", w.Code, w.Body.String())
		}
	}
}

func deleteWorkbenchResource(s *Service, ctx context.Context, r *Resource, detach bool) *httptest.ResponseRecorder {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "kind", Value: r.Kind}, {Key: "id", Value: r.ID}}
	c.Request = httptest.NewRequest(http.MethodDelete, "/?detach="+map[bool]string{true: "true", false: "false"}[detach], nil).WithContext(ctx)
	NewHandler(s).Delete(c)
	return w
}

func TestDeleteBlocksQueuedJobAndDetachesWithRevision(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	r := seedResource(t, s, 1, "alice", "training-maps")
	r.Status = "confirmed"
	r.Data = asJSON(MapData{Days: []Day{{Day: 1, Title: "day"}}})
	if err := s.db.Save(r).Error; err != nil {
		t.Fatal(err)
	}
	flow := seedResource(t, s, 1, "alice", "flows")
	flow.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "author"}}})
	if err := s.db.Save(flow).Error; err != nil {
		t.Fatal(err)
	}
	s.agents = generationAgentStub{}
	task := seedResource(t, s, 1, "alice", "content-tasks")
	task.Data = asJSON(TaskData{MapID: r.ID, FlowID: flow.ID, Day: 1, AgentID: "author", Review: "old review", Composed: "old draft"})
	if err := s.db.Model(task).Update("data", task.Data).Error; err != nil {
		t.Fatal(err)
	}
	if _, err := s.submit(ctx, task.Kind, task.ID, "decompose", ActionRequest{Version: 1}); err != nil {
		t.Fatal(err)
	}
	if w := deleteWorkbenchResource(s, ctx, r, true); w.Code != 409 {
		t.Fatalf("active dependent must block delete: %d %s", w.Code, w.Body.String())
	}
	if err := s.db.Model(&Job{}).Where("resource_id = ?", task.ID).Update("status", "failed").Error; err != nil {
		t.Fatal(err)
	}
	if w := deleteWorkbenchResource(s, ctx, r, false); w.Code != 409 {
		t.Fatal("delete requires explicit detach")
	}
	if w := deleteWorkbenchResource(s, ctx, r, true); w.Code != 200 {
		t.Fatalf("detach delete: %d %s", w.Code, w.Body.String())
	}
	stored, err := s.get(ctx, task.Kind, task.ID)
	if err != nil {
		t.Fatal(err)
	}
	d, _ := decode[TaskData](stored.Data)
	var revisions int64
	s.db.Model(&Revision{}).Where("resource_id = ?", task.ID).Count(&revisions)
	if stored.Version != 2 || stored.Status != "needs_configuration" || d.MapID != "" || d.Review != "old review" || d.Composed != "old draft" || revisions != 1 {
		t.Fatal("detachment must preserve historical output and record a revision")
	}
}

func TestManualUpdateCannotInjectStateAndReassignmentInvalidatesFollowing(t *testing.T) {
	original := TaskData{Segments: []Segment{{ID: "s1", AgentID: "a", Status: "accepted", Content: "first"}, {ID: "s2", AgentID: "a", Status: "accepted", Content: "second"}, {ID: "s3", AgentID: "a", Status: "generated", Content: "third"}}, Review: "review", Composed: "final"}
	r := &Resource{Kind: "content-tasks", Data: asJSON(original), Status: "completed"}
	injected, _ := decode[TaskData](asJSON(original))
	injected.Segments[1].Content = "injected"
	if err := applyManualUpdate(r, resourceRequest{Name: "task", Data: asJSON(injected)}); !errors.Is(err, ErrInvalid) {
		t.Fatalf("injected content accepted: %v", err)
	}
	changed, _ := decode[TaskData](asJSON(original))
	changed.Segments[1].AgentID = "b"
	if err := applyManualUpdate(r, resourceRequest{Name: "task", Data: asJSON(changed)}); err != nil {
		t.Fatal(err)
	}
	actual, _ := decode[TaskData](r.Data)
	if actual.Segments[0].Content != "first" || actual.Segments[0].Status != "accepted" || actual.Segments[1].AgentID != "b" || actual.Segments[1].Content != "second" || actual.Segments[1].Status != "needs_review" || actual.Segments[2].Status != "needs_review" || actual.Review != "" || actual.Composed != "" {
		t.Fatalf("reassignment must invalidate its dependent output: %+v", actual)
	}
	md := MapData{Topic: "before", DaysCount: 1, Days: []Day{{Day: 1, Confirmed: true}}}
	r = &Resource{Kind: "training-maps", Data: asJSON(md), Status: "confirmed"}
	md.Days[0].Title = "injected"
	if err := applyManualUpdate(r, resourceRequest{Data: asJSON(md)}); !errors.Is(err, ErrInvalid) {
		t.Fatal("direct map content mutation accepted")
	}
	md, _ = decode[MapData](r.Data)
	md.Topic = "after"
	if err := applyManualUpdate(r, resourceRequest{Name: "map", Data: asJSON(md)}); err != nil {
		t.Fatal(err)
	}
	md, _ = decode[MapData](r.Data)
	if len(md.Days) != 1 || md.Days[0].Confirmed || r.Status != "needs_review" {
		t.Fatal("configuration change must retain output and require its review")
	}
}

// Uses only an isolated schema. Set WORKBENCH_TEST_POSTGRES_DSN to opt in.
func TestPostgresResourceRaces(t *testing.T) {
	dsn := os.Getenv("WORKBENCH_TEST_POSTGRES_DSN")
	if dsn == "" {
		t.Skip("WORKBENCH_TEST_POSTGRES_DSN is not set")
	}
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	sqlDB, _ := db.DB()
	defer sqlDB.Close()
	schema := "wb_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
	if err = db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
		t.Fatal(err)
	}
	defer db.Exec("DROP SCHEMA " + schema + " CASCADE")
	// Every pooled connection uses the isolated schema, including parallel workers.
	testDSN := dsn + " search_path=" + schema
	if strings.Contains(dsn, "://") {
		separator := "?"
		if strings.Contains(dsn, "?") {
			separator = "&"
		}
		testDSN = dsn + separator + "search_path=" + schema
	}
	testDB, err := gorm.Open(postgres.Open(testDSN), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal(err)
	}
	testSQL, _ := testDB.DB()
	defer testSQL.Close()
	if err = testDB.AutoMigrate(&Resource{}, &Job{}, &Revision{}, &Event{}, &Consumption{}); err != nil {
		t.Fatal(err)
	}
	s := NewService(testDB)
	s.agents = generationAgentStub{}
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	seedGeneratableMap := func() *Resource {
		r := seedResource(t, s, 1, "alice", "training-maps")
		r.Data = asJSON(MapData{Topic: "topic", AgentID: "author", DaysCount: 1})
		if err := s.db.Save(r).Error; err != nil {
			t.Fatal(err)
		}
		return r
	}
	t.Run("lease expires while waiting for lock", func(t *testing.T) {
		r := seedResource(t, s, 1, "alice", "training-maps")
		until := time.Now().Add(150 * time.Millisecond)
		job := Job{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", ResourceID: r.ID, Kind: r.Kind, Status: "running", LeaseToken: "locked-lease", LeaseUntil: &until}
		if err := s.db.Create(&job).Error; err != nil {
			t.Fatal(err)
		}
		locked := s.db.Begin()
		if locked.Error != nil {
			t.Fatal(locked.Error)
		}
		defer locked.Rollback()
		if err := locked.Exec("SELECT id FROM workbench_jobs WHERE id = ? FOR UPDATE", job.ID).Error; err != nil {
			t.Fatal(err)
		}
		leaseCtx := context.WithValue(ctx, leaseContextKey{}, jobLease{job.ID, job.LeaseToken})
		finished := make(chan error, 1)
		go func() { candidate := *r; candidate.Name = "too late"; finished <- s.save(leaseCtx, &candidate, 1) }()
		time.Sleep(200 * time.Millisecond)
		if err := locked.Commit().Error; err != nil {
			t.Fatal(err)
		}
		if err := <-finished; !errors.Is(err, ErrConflict) {
			t.Fatalf("lease expired under lock was allowed to save: %v", err)
		}
	})
	for i := 0; i < 12; i++ {
		r := seedGeneratableMap()
		start := make(chan struct{})
		var wg sync.WaitGroup
		var editErr, submitErr error
		wg.Add(2)
		go func() {
			defer wg.Done()
			<-start
			candidate := *r
			candidate.Name = "edited"
			editErr = s.save(ctx, &candidate, 1)
		}()
		go func() {
			defer wg.Done()
			<-start
			_, submitErr = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1})
		}()
		close(start)
		wg.Wait()
		if (editErr == nil) == (submitErr == nil) {
			t.Fatalf("exactly one concurrent write must succeed: edit=%v submit=%v", editErr, submitErr)
		}
	}
	t.Run("binding versus parent delete", func(t *testing.T) {
		for i := 0; i < 8; i++ {
			m := seedResource(t, s, 1, "alice", "training-maps")
			if err := s.db.Model(m).Update("status", "confirmed").Error; err != nil {
				t.Fatal(err)
			}
			flow := seedResource(t, s, 1, "alice", "flows")
			task := &Resource{ID: uuid.NewString(), TenantID: 1, OwnerID: "alice", Kind: "content-tasks", Name: "bound task", Version: 1, Data: asJSON(TaskData{MapID: m.ID, FlowID: flow.ID})}
			start := make(chan struct{})
			var wg sync.WaitGroup
			var createErr error
			var deleted *httptest.ResponseRecorder
			wg.Add(2)
			go func() { defer wg.Done(); <-start; createErr = s.create(ctx, task) }()
			go func() { defer wg.Done(); <-start; deleted = deleteWorkbenchResource(s, ctx, m, true) }()
			close(start)
			wg.Wait()
			if deleted.Code != 200 {
				t.Fatalf("parent deletion failed: %d %s", deleted.Code, deleted.Body.String())
			}
			if createErr == nil {
				stored, err := s.get(ctx, task.Kind, task.ID)
				if err != nil {
					t.Fatal(err)
				}
				data, _ := decode[TaskData](stored.Data)
				if data.MapID != "" {
					t.Fatal("concurrent create left dangling reference")
				}
			}
		}
	})
	for i := 0; i < 12; i++ {
		r := seedGeneratableMap()
		start := make(chan struct{})
		var wg sync.WaitGroup
		var deleted *httptest.ResponseRecorder
		var submitErr error
		wg.Add(2)
		go func() { defer wg.Done(); <-start; deleted = deleteWorkbenchResource(s, ctx, r, true) }()
		go func() {
			defer wg.Done()
			<-start
			_, submitErr = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1})
		}()
		close(start)
		wg.Wait()
		if (deleted.Code == 200) == (submitErr == nil) {
			t.Fatalf("delete and submit must not both succeed: delete=%d submit=%v", deleted.Code, submitErr)
		}
	}
	r := seedGeneratableMap()
	if _, err = s.submit(ctx, r.Kind, r.ID, "generate", ActionRequest{Version: 1}); err != nil {
		t.Fatal(err)
	}
	var wg sync.WaitGroup
	var mu sync.Mutex
	seen := map[string]bool{}
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			job, err := s.claim(ctx)
			if err != nil {
				t.Errorf("claim: %v", err)
				return
			}
			if job != nil {
				mu.Lock()
				if seen[job.ID] {
					t.Error("same job claimed twice")
				}
				seen[job.ID] = true
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
}