package app

import (
	"context"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/coredb/testkit"
	"github.com/tidbcloud/AgentTeam/internal/store/agent"
)

func TestUpsertAgentRoleMaterializationConcurrentSameRoleVersionKeepsSingleProjectionSet(t *testing.T) {
	ctx := context.Background()
	rootDSN := appTiDB.RootDSN(t)
	dsn := testkit.StartInitializedTiDBWithRootDSN(t, rootDSN)

	sqlDB, err := testkit.OpenTiDBContext(ctx, dsn)
	require.NoError(t, err)
	t.Cleanup(func() { _ = sqlDB.Close() })
	sqlDB.SetMaxOpenConns(8)
	sqlDB.SetMaxIdleConns(8)

	db := coredb.WrapSQLDB(sqlDB)
	now := time.Unix(100, 0).UTC()
	fixture := seedAppAdminFixture(t, db, "agent_role_materialization_concurrent", now)
	seedMachineWorkspaceAgent(t, db, fixture.WorkspaceID, fixture.UserID, "mac_role_materialization_a", "agt_role_materialization_a", "mem_role_materialization_a", "Role Agent A", "active", now)
	seedMachineWorkspaceAgent(t, db, fixture.WorkspaceID, fixture.UserID, "mac_role_materialization_b", "agt_role_materialization_b", "mem_role_materialization_b", "Role Agent B", "active", now)

	barrierStore := New(&manifestProjectionBarrierDB{
		DB:      db,
		barrier: newManifestProjectionBarrier(2),
	}, Options{})

	inputForAgent := func(agentID string) agent.AgentRoleMaterializationInput {
		return agent.AgentRoleMaterializationInput{
			UserID:      fixture.UserID,
			WorkspaceID: fixture.WorkspaceID,
			AgentID:     agentID,
			Source: agent.AgentRoleMaterializationSource{
				BindingKind: "my_role",
				RoleID:      "role_backend_copy",
			},
			RoleID:         "role_backend_copy",
			RoleVersionID:  "rlv_backend_shared",
			RealizationID:  "rrl_backend_shared",
			DefinitionPath: "memory/system/identity.md",
			ManifestItems: []agent.AgentRoleMaterializationItem{
				{Type: "role_definition_asset", Name: "IDENTITY.md", Content: "# Backend"},
				{Type: "portable_knowledge", Name: "Runbook", Content: "# Runbook"},
				{Type: "skill", Name: "Deploy", Content: "# Skill"},
			},
		}
	}

	errs := make(chan error, 2)
	start := make(chan struct{})
	var wg sync.WaitGroup
	for _, agentID := range []string{"agt_role_materialization_a", "agt_role_materialization_b"} {
		wg.Add(1)
		go func(agentID string) {
			defer wg.Done()
			<-start
			errs <- barrierStore.UpsertAgentRoleMaterialization(ctx, inputForAgent(agentID))
		}(agentID)
	}

	close(start)
	wg.Wait()
	close(errs)

	for err := range errs {
		require.NoError(t, err)
	}

	var rowCount int
	require.NoError(t, db.QueryRow(ctx, `
SELECT COUNT(*)
FROM role_package_manifest_items
WHERE role_version_id = $1
`, "rlv_backend_shared").Scan(&rowCount))
	require.Equal(t, 3, rowCount)

	materialization, err := agent.New(db).LoadAgentWorkspaceMaterialization(ctx, "agt_role_materialization_a")
	require.NoError(t, err)
	require.NotNil(t, materialization)
	require.Equal(t, []agent.WorkspaceSeedFile{
		{Path: "memory/system/identity.md", Content: "# Backend"},
		{Path: ".agents/skills/Deploy/SKILL.md", Content: "# Skill"},
		{Path: ".agents/knowledge/Runbook.md", Content: "# Runbook"},
	}, materialization.SeedFiles)
}

type manifestProjectionBarrierDB struct {
	coredb.DB
	barrier *manifestProjectionBarrier
}

func (db *manifestProjectionBarrierDB) BeginTx(ctx context.Context, options coredb.TxOptions) (coredb.Tx, error) {
	tx, err := db.DB.BeginTx(ctx, options)
	if err != nil {
		return nil, err
	}
	return &manifestProjectionBarrierTx{Tx: tx, barrier: db.barrier}, nil
}

type manifestProjectionBarrierTx struct {
	coredb.Tx
	barrier        *manifestProjectionBarrier
	manifestSynced bool
}

func (tx *manifestProjectionBarrierTx) Exec(ctx context.Context, query string, args ...any) (coredb.Result, error) {
	result, err := tx.Tx.Exec(ctx, query, args...)
	if err != nil {
		return result, err
	}
	if !tx.manifestSynced && strings.Contains(query, "role_package_manifest_items") {
		tx.manifestSynced = true
		tx.barrier.Wait()
	}
	return result, nil
}

type manifestProjectionBarrier struct {
	want    int
	mu      sync.Mutex
	arrived int
	ready   chan struct{}
}

func newManifestProjectionBarrier(want int) *manifestProjectionBarrier {
	return &manifestProjectionBarrier{
		want:  want,
		ready: make(chan struct{}),
	}
}

func (barrier *manifestProjectionBarrier) Wait() {
	barrier.mu.Lock()
	barrier.arrived++
	ready := barrier.ready
	if barrier.arrived == barrier.want {
		close(barrier.ready)
	}
	barrier.mu.Unlock()
	<-ready
}