package access

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
	"github.com/tidbcloud/AgentTeam/internal/store/testfixture"
)

type ownershipScriptRow struct {
	values []any
	err    error
}

func (row ownershipScriptRow) Scan(destinations ...any) error {
	if row.err != nil {
		return row.err
	}
	if len(destinations) != len(row.values) {
		return fmt.Errorf("scan destination count %d does not match values %d", len(destinations), len(row.values))
	}
	for index, value := range row.values {
		destination := reflect.ValueOf(destinations[index])
		if destination.Kind() != reflect.Pointer || destination.IsNil() {
			return fmt.Errorf("scan destination %d is not a pointer", index)
		}
		destination.Elem().Set(reflect.ValueOf(value))
	}
	return nil
}

type ownershipRecordingTx struct {
	rows    []ownershipScriptRow
	queries []string
	args    [][]any
}

func (tx *ownershipRecordingTx) QueryRow(_ context.Context, query string, args ...any) coredb.Row {
	tx.queries = append(tx.queries, query)
	tx.args = append(tx.args, append([]any(nil), args...))
	if len(tx.rows) == 0 {
		return ownershipScriptRow{err: fmt.Errorf("unexpected ownership query")}
	}
	row := tx.rows[0]
	tx.rows = tx.rows[1:]
	return row
}

func (*ownershipRecordingTx) Query(context.Context, string, ...any) (coredb.Rows, error) {
	return nil, fmt.Errorf("unexpected Query call")
}

func (*ownershipRecordingTx) Exec(context.Context, string, ...any) (coredb.Result, error) {
	return nil, fmt.Errorf("unexpected Exec call")
}

func (*ownershipRecordingTx) Commit(context.Context) error   { return nil }
func (*ownershipRecordingTx) Rollback(context.Context) error { return nil }

func validOwnershipRows() []ownershipScriptRow {
	return []ownershipScriptRow{
		{values: []any{"mem_manager", "wsp_ownership", "human", "usr_manager", "active"}},
		{values: []any{"mac_agent", "wsp_ownership", "usr_manager", "active"}},
		{values: []any{"agt_actor", "wsp_ownership", "mac_agent", "mem_agent", "inactive"}},
		{values: []any{"mem_agent", "wsp_ownership", "agent", "active"}},
	}
}

func expectedOwnershipPrincipal() AuthenticatedAgentPrincipal {
	return AuthenticatedAgentPrincipal{
		WorkspaceID: "wsp_ownership",
		AgentID:     "agt_actor",
		MemberID:    "mem_agent",
		MachineID:   "mac_agent",
	}
}

func expectedOwnershipManager() HumanManager {
	return HumanManager{UserID: "usr_manager", MemberID: "mem_manager"}
}

func TestLockAndRevalidateAgentPrincipalManagerUsesLockOrder(t *testing.T) {
	tx := &ownershipRecordingTx{rows: validOwnershipRows()}

	err := LockAndRevalidateAgentPrincipalManager(
		context.Background(), tx, expectedOwnershipPrincipal(), expectedOwnershipManager(),
	)
	require.NoError(t, err)
	require.Len(t, tx.queries, 4)
	require.Contains(t, tx.queries[0], "FROM members")
	require.Contains(t, tx.queries[1], "FROM machines")
	require.Contains(t, tx.queries[2], "FROM agents")
	require.Contains(t, tx.queries[3], "FROM members")
	for _, query := range tx.queries {
		require.Contains(t, query, "FOR UPDATE")
		require.NotContains(t, query, "JOIN")
	}
	require.Equal(t, [][]any{
		{"wsp_ownership", "mem_manager"},
		{"wsp_ownership", "mac_agent"},
		{"wsp_ownership", "agt_actor"},
		{"wsp_ownership", "mem_agent"},
	}, tx.args)
}

func TestLockAndRevalidateAgentPrincipalManagerChecksEveryLockedIdentity(t *testing.T) {
	tests := []struct {
		name     string
		rowIndex int
		value    int
		changed  string
	}{
		{name: "manager id", rowIndex: 0, value: 0, changed: "mem_other"},
		{name: "manager workspace", rowIndex: 0, value: 1, changed: "wsp_other"},
		{name: "manager type", rowIndex: 0, value: 2, changed: "agent"},
		{name: "manager user", rowIndex: 0, value: 3, changed: "usr_other"},
		{name: "manager inactive", rowIndex: 0, value: 4, changed: "inactive"},
		{name: "machine id", rowIndex: 1, value: 0, changed: "mac_other"},
		{name: "machine workspace", rowIndex: 1, value: 1, changed: "wsp_other"},
		{name: "machine creator", rowIndex: 1, value: 2, changed: "usr_other"},
		{name: "machine inactive", rowIndex: 1, value: 3, changed: "inactive"},
		{name: "Agent id", rowIndex: 2, value: 0, changed: "agt_other"},
		{name: "Agent workspace", rowIndex: 2, value: 1, changed: "wsp_other"},
		{name: "Agent machine", rowIndex: 2, value: 2, changed: "mac_other"},
		{name: "Agent member", rowIndex: 2, value: 3, changed: "mem_other"},
		{name: "Agent archived", rowIndex: 2, value: 4, changed: "archived"},
		{name: "Agent member id", rowIndex: 3, value: 0, changed: "mem_other"},
		{name: "Agent member workspace", rowIndex: 3, value: 1, changed: "wsp_other"},
		{name: "Agent member type", rowIndex: 3, value: 2, changed: "human"},
		{name: "Agent member inactive", rowIndex: 3, value: 3, changed: "inactive"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			rows := validOwnershipRows()
			rows[test.rowIndex].values[test.value] = test.changed
			tx := &ownershipRecordingTx{rows: rows}

			err := LockAndRevalidateAgentPrincipalManager(
				context.Background(), tx, expectedOwnershipPrincipal(), expectedOwnershipManager(),
			)
			require.ErrorIs(t, err, storeerr.ErrForbidden)
			require.Len(t, tx.queries, test.rowIndex+1)
		})
	}
}

type ownershipFixture struct {
	workspaceID       string
	managerUserID     string
	managerMemberID   string
	otherUserID       string
	agentID           string
	agentMemberID     string
	otherAgentMember  string
	machineID         string
	sameManagerMachID string
	otherManagerMach  string
}

func seedOwnershipFixture(t *testing.T, ctx context.Context, db coredb.QueryExecer, suffix string) ownershipFixture {
	t.Helper()
	now := time.Date(2026, 8, 4, 10, 0, 0, 0, time.UTC)
	fixture := ownershipFixture{
		workspaceID:       "wsp_own_" + suffix,
		managerUserID:     "usr_own_mgr_" + suffix,
		managerMemberID:   "mem_own_mgr_" + suffix,
		otherUserID:       "usr_own_other_" + suffix,
		agentID:           "agt_own_" + suffix,
		agentMemberID:     "mem_own_agent_" + suffix,
		otherAgentMember:  "mem_own_agent_other_" + suffix,
		machineID:         "mac_own_" + suffix,
		sameManagerMachID: "mac_own_same_" + suffix,
		otherManagerMach:  "mac_own_other_" + suffix,
	}
	mustExec(t, ctx, db, `
INSERT INTO users (id,email,name,display_name,avatar_url,is_admin,password_hash,email_verified_at,created_at,updated_at)
VALUES
($1,$1,$1,$1,'',FALSE,'hash',$3,$3,$3),
($2,$2,$2,$2,'',FALSE,'hash',$3,$3,$3)
`, fixture.managerUserID, fixture.otherUserID, now)
	mustExec(t, ctx, db, `
INSERT INTO workspaces (id,organization_id,slug,name,latest_message_seq,created_at,updated_at)
VALUES ($1,$2,$1,$1,0,$3,$3)
`, fixture.workspaceID, testfixture.SeedWorkspaceOrganization(t, ctx, db, fixture.workspaceID, now), now)
	mustExec(t, ctx, db, `
INSERT INTO members (id,workspace_id,type,display_name,user_id,workspace_role,status,created_at,updated_at)
VALUES
($1,$3,'human',$1,$4,'member','active',$6,$6),
($2,$3,'agent',$2,NULL,NULL,'active',$6,$6),
($5,$3,'agent',$5,NULL,NULL,'active',$6,$6)
`, fixture.managerMemberID, fixture.agentMemberID, fixture.workspaceID, fixture.managerUserID, fixture.otherAgentMember, now)
	mustExec(t, ctx, db, `
INSERT INTO machines (id,workspace_id,created_by_user_id,name,api_key,status,created_at,updated_at)
VALUES
($1,$4,$5,$1,$1,'active',$7,$7),
($2,$4,$5,$2,$2,'active',$7,$7),
($3,$4,$6,$3,$3,'active',$7,$7)
`, fixture.machineID, fixture.sameManagerMachID, fixture.otherManagerMach, fixture.workspaceID, fixture.managerUserID, fixture.otherUserID, now)
	mustExec(t, ctx, db, `
INSERT INTO agents (
  id,workspace_id,machine_id,member_id,name,runtime,model,description,auth_token_hash,
  status,activity_detail,created_at,updated_at
) VALUES ($1,$2,$3,$4,$1,'codex','gpt-5','','hash','inactive','',$5,$5)
`, fixture.agentID, fixture.workspaceID, fixture.machineID, fixture.agentMemberID, now)
	return fixture
}

func (fixture ownershipFixture) principal() AuthenticatedAgentPrincipal {
	return AuthenticatedAgentPrincipal{
		WorkspaceID: fixture.workspaceID,
		AgentID:     fixture.agentID,
		MemberID:    fixture.agentMemberID,
		MachineID:   fixture.machineID,
	}
}

func (fixture ownershipFixture) manager() HumanManager {
	return HumanManager{UserID: fixture.managerUserID, MemberID: fixture.managerMemberID}
}

func TestLockAndRevalidateAgentPrincipalManagerAgainstCoreFacts(t *testing.T) {
	ctx := context.Background()
	db := startAccessTestPool(t, ctx)
	tests := []struct {
		name   string
		mutate func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture)
	}{
		{name: "current principal and manager"},
		{name: "manager member mismatch", mutate: func(_ *testing.T, _ context.Context, _ coredb.QueryExecer, fixture *ownershipFixture) {
			fixture.managerMemberID = "mem_missing"
		}},
		{name: "manager inactive", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE members SET status='inactive' WHERE workspace_id=$1 AND id=$2`, fixture.workspaceID, fixture.managerMemberID)
		}},
		{name: "machine mismatch", mutate: func(_ *testing.T, _ context.Context, _ coredb.QueryExecer, fixture *ownershipFixture) {
			fixture.machineID = "mac_missing"
		}},
		{name: "machine inactive", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE machines SET status='inactive' WHERE workspace_id=$1 AND id=$2`, fixture.workspaceID, fixture.machineID)
		}},
		{name: "machine creator drift", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE machines SET created_by_user_id=$1 WHERE workspace_id=$2 AND id=$3`, fixture.otherUserID, fixture.workspaceID, fixture.machineID)
		}},
		{name: "Agent mismatch", mutate: func(_ *testing.T, _ context.Context, _ coredb.QueryExecer, fixture *ownershipFixture) {
			fixture.agentID = "agt_missing"
		}},
		{name: "Agent archived", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE agents SET status='archived' WHERE workspace_id=$1 AND id=$2`, fixture.workspaceID, fixture.agentID)
		}},
		{name: "Agent member rebind", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE agents SET member_id=$1 WHERE workspace_id=$2 AND id=$3`, fixture.otherAgentMember, fixture.workspaceID, fixture.agentID)
		}},
		{name: "Agent machine rebind to same manager", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE agents SET machine_id=$1 WHERE workspace_id=$2 AND id=$3`, fixture.sameManagerMachID, fixture.workspaceID, fixture.agentID)
		}},
		{name: "Agent machine rebind to different manager", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE agents SET machine_id=$1 WHERE workspace_id=$2 AND id=$3`, fixture.otherManagerMach, fixture.workspaceID, fixture.agentID)
		}},
		{name: "Agent member mismatch", mutate: func(_ *testing.T, _ context.Context, _ coredb.QueryExecer, fixture *ownershipFixture) {
			fixture.agentMemberID = "mem_missing"
		}},
		{name: "Agent member inactive", mutate: func(t *testing.T, ctx context.Context, db coredb.QueryExecer, fixture *ownershipFixture) {
			mustExec(t, ctx, db, `UPDATE members SET status='inactive' WHERE workspace_id=$1 AND id=$2`, fixture.workspaceID, fixture.agentMemberID)
		}},
	}

	for index, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			fixture := seedOwnershipFixture(t, ctx, db, fmt.Sprintf("%02d_%s", index, strings.ReplaceAll(test.name, " ", "_")))
			if test.mutate != nil {
				test.mutate(t, ctx, db, &fixture)
			}
			tx, err := db.(coredb.TxBeginner).BeginTx(ctx, coredb.TxOptions{})
			require.NoError(t, err)
			t.Cleanup(func() { _ = tx.Rollback(context.Background()) })

			err = LockAndRevalidateAgentPrincipalManager(ctx, tx, fixture.principal(), fixture.manager())
			if test.mutate == nil {
				require.NoError(t, err)
			} else {
				require.ErrorIs(t, err, storeerr.ErrForbidden)
			}
		})
	}
}
