package taskrun_test

import (
	"context"
	"encoding/json"
	"errors"
	"net/url"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"chat-agent-system/internal/app/uow"
	"chat-agent-system/internal/artifactvault"
	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/ctxmgr"
	"chat-agent-system/internal/db"
	"chat-agent-system/internal/docgov"
	"chat-agent-system/internal/stream"
	"chat-agent-system/internal/taskrun"
	"chat-agent-system/internal/workspace"

	"github.com/jackc/pgx/v5/pgxpool"
)

func init() {
	if _, err := os.Stat("deploy/postgres/goose"); err == nil {
		return
	}
	for _, cand := range []string{"../..", "../../.."} {
		if fi, err := os.Stat(cand + "/deploy/postgres/goose"); err == nil && fi.IsDir() {
			_ = os.Chdir(cand)
			return
		}
	}
}

var (
	isoCleanup func()
)

func TestMain(m *testing.M) {
	code := m.Run()
	if isoCleanup != nil {
		isoCleanup()
	}
	os.Exit(code)
}

func rewriteDatabaseName(dsn, name string) string {
	if u, err := url.Parse(dsn); err == nil && u.Scheme != "" {
		u.Path = "/" + name
		return u.String()
	}
	return dsn
}

var isoOnce sync.Once
var isoPoolA, isoPoolB *pgxpool.Pool
var isoErr error

func openIsolatedM3(t *testing.T) (*pgxpool.Pool, *pgxpool.Pool) {
	t.Helper()
	if testing.Short() {
		t.Skip("U8 PG races are not part of -short")
	}
	dsn := db.SkipIfNoTestDatabaseURL(t)
	db.RefuseSharedEnvDSN(t, dsn)

	isoOnce.Do(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
		defer cancel()
		adminDSN := rewriteDatabaseName(dsn, "postgres")
		admin, err := pgxpool.New(ctx, adminDSN)
		if err != nil {
			isoErr = err
			return
		}
		const isolated = "agent_rel_m3_u8"
		_, _ = admin.Exec(ctx, `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
		_, _ = admin.Exec(ctx, `DROP DATABASE IF EXISTS `+isolated)
		if _, err := admin.Exec(ctx, `CREATE DATABASE `+isolated); err != nil {
			admin.Close()
			isoErr = err
			return
		}
		isoDSN := rewriteDatabaseName(dsn, isolated)
		poolA, err := pgxpool.New(ctx, isoDSN)
		if err != nil {
			admin.Close()
			isoErr = err
			return
		}
		if err := db.RunMigrations(ctx, poolA, db.ResolveMigrationsDir()); err != nil {
			poolA.Close()
			admin.Close()
			isoErr = err
			return
		}
		poolB, err := pgxpool.New(ctx, isoDSN)
		if err != nil {
			poolA.Close()
			admin.Close()
			isoErr = err
			return
		}
		isoPoolA, isoPoolB = poolA, poolB
		isoCleanup = func() {
			poolA.Close()
			poolB.Close()
			_, _ = admin.Exec(context.Background(), `SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname=$1 AND pid <> pg_backend_pid()`, isolated)
			_, _ = admin.Exec(context.Background(), `DROP DATABASE IF EXISTS `+isolated)
			admin.Close()
		}
	})
	if isoErr != nil {
		t.Fatalf("isolated db: %v", isoErr)
	}
	return isoPoolA, isoPoolB
}

func seedUserAndProject(t *testing.T, pool *pgxpool.Pool) (userID, projectID int64) {
	t.Helper()
	ctx := context.Background()
	err := pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID)
	if err != nil {
		t.Fatalf("seed user: %v", err)
	}
	err = pool.QueryRow(ctx, `SELECT id FROM projects ORDER BY id LIMIT 1`).Scan(&projectID)
	if err != nil {
		t.Fatalf("seed project: %v", err)
	}
	return userID, projectID
}

func insertRunning(t *testing.T, pool *pgxpool.Pool, rec *taskrun.RunRecord) {
	t.Helper()
	_, err := pool.Exec(context.Background(), `
		INSERT INTO task_runs (
			run_id, session_id, project_id, user_id, client_turn_id, task_key, status,
			workspace_id, assistant_message_id, created_at, updated_at
		) VALUES ($1,$2,$3,$4,$5,$6,'running',$7,$8,now(),now())
	`, rec.RunID, rec.SessionID, rec.ProjectID, rec.UserID, rec.ClientTurnID, rec.TaskKey, rec.WorkspaceID, rec.AssistantMsgID)
	if err != nil {
		t.Fatalf("insert run: %v", err)
	}
}

func TestPG_DuplicateFinalizeIdempotent(t *testing.T) {
	poolA, poolB := openIsolatedM3(t)
	userID, projectID := seedUserAndProject(t, poolA)
	ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
	sessionID := "sess_u8_dup"
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
	if err != nil {
		t.Fatal(err)
	}
	rec := &taskrun.RunRecord{
		RunID: "run_u8_dup", SessionID: sessionID, ProjectID: projectID, UserID: userID,
		TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: "ws_u8_dup",
	}
	insertRunning(t, poolA, rec)
	body := []byte("# committed intake body\n")
	res := &bizdev.TaskRunResult{
		Success: true, TaskKey: "intake", ArtifactTitle: "用户初始需求收集表", MarkdownBody: string(body),
		DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_dup", ContentSHA256: artifactvault.SHA256Hex(body), Version: 1, ProjectID: projectID},
	}
	cmdsA := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
	cmdsB := &taskrun.Commands{UoW: uow.NewPG(poolB), Ctx: ctxmgr.NewContextManager(poolB, nil, ctxmgr.TimeDecayConfig{}), Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
	st1, err := cmdsA.Finalize(context.Background(), rec, res)
	if err != nil {
		t.Fatal(err)
	}
	st2, err := cmdsB.Finalize(context.Background(), rec, res)
	if err != nil {
		t.Fatal(err)
	}
	if st1 != "completed" || st2 != "completed" {
		t.Fatalf("status %s %s", st1, st2)
	}
	var msgs, mans, envs, stgs int
	_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM chat_messages WHERE session_id=$1 AND role='assistant'`, sessionID).Scan(&msgs)
	_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&mans)
	_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM event_envelopes WHERE aggregate_id=$1`, rec.RunID).Scan(&envs)
	_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM artifact_payload_staging WHERE source_event_id=$1`, "taskrun:"+rec.RunID+":completed").Scan(&stgs)
	if msgs != 1 || mans != 1 || envs != 1 || stgs != 1 {
		t.Fatalf("counts messages=%d manifests=%d envelopes=%d staging=%d", msgs, mans, envs, stgs)
	}
	var sha1, sha2 string
	_ = poolA.QueryRow(context.Background(), `SELECT sha256 FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&sha1)
	_ = poolB.QueryRow(context.Background(), `SELECT sha256 FROM task_run_manifests WHERE run_id=$1`, rec.RunID).Scan(&sha2)
	if sha1 == "" || sha1 != sha2 {
		t.Fatalf("manifest sha %q %q", sha1, sha2)
	}
}

func TestPG_FinalizeCancelRaceOneTerminal(t *testing.T) {
	poolA, poolB := openIsolatedM3(t)
	userID, projectID := seedUserAndProject(t, poolA)
	ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
	sessionID := "sess_u8_race"
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
	if err != nil {
		t.Fatal(err)
	}
	rec := &taskrun.RunRecord{
		RunID: "run_u8_race", SessionID: sessionID, ProjectID: projectID, UserID: userID,
		TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: "ws_u8_race",
	}
	insertRunning(t, poolA, rec)
	body := []byte("# race body\n")
	res := &bizdev.TaskRunResult{
		Success: true, TaskKey: "intake", MarkdownBody: string(body),
		DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_race", ContentSHA256: artifactvault.SHA256Hex(body), Version: 1},
	}
	cmdsA := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
	cmdsB := &taskrun.Commands{UoW: uow.NewPG(poolB), Ctx: ctxmgr.NewContextManager(poolB, nil, ctxmgr.TimeDecayConfig{}), Env: "dev", StorageIdentity: "env:dev"}
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = cmdsA.Finalize(context.Background(), rec, res)
	}()
	go func() {
		defer wg.Done()
		_, _ = cmdsB.Cancel(context.Background(), rec)
	}()
	wg.Wait()
	var status string
	if err := poolA.QueryRow(context.Background(), `SELECT status FROM task_runs WHERE run_id=$1`, rec.RunID).Scan(&status); err != nil {
		t.Fatal(err)
	}
	if status != "completed" && status != "cancelled" {
		t.Fatalf("illegal terminal %q", status)
	}
	var n int
	_ = poolA.QueryRow(context.Background(), `SELECT COUNT(*) FROM chat_messages WHERE session_id=$1 AND role='assistant'`, sessionID).Scan(&n)
	if n != 1 {
		t.Fatalf("assistant turns=%d", n)
	}
}

func TestPG_SSEDoneNotEmittedWhenCommitFails(t *testing.T) {
	poolA, _ := openIsolatedM3(t)
	userID, projectID := seedUserAndProject(t, poolA)
	ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
	sessionID := "sess_u8_sse"
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
	if err != nil {
		t.Fatal(err)
	}
	rec := &taskrun.RunRecord{
		RunID: "run_u8_sse", SessionID: sessionID, ProjectID: projectID, UserID: userID,
		TaskKey: "intake", AssistantMsgID: asstID, Status: "pending",
	}
	unit := uow.NewPG(poolA)
	unit.BeforeCommit = func(uow.TxScope) error { return errors.New("inject commit fail") }
	cmds := &taskrun.Commands{UoW: unit, Ctx: ctxMgr, Env: "dev", StorageIdentity: "env:dev"}
	mgr := taskrun.NewManager(nil, ctxMgr)
	mgr.SetStore(taskrun.NewPGStore(poolA))
	mgr.SetCommands(cmds)
	mgr.SetExecutor(func(context.Context, string, bizdev.TaskRunRequest, *stream.Streamer) (*bizdev.TaskRunResult, error) {
		return &bizdev.TaskRunResult{Success: true, MarkdownBody: "# sse\n"}, nil
	}, nil, nil)
	if _, err := mgr.CreateAndStart(context.Background(), rec.RunID, rec, bizdev.TaskRunRequest{TaskKey: "intake"}, bizdev.StageRunRequest{}, ""); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(8 * time.Second)
	for time.Now().Before(deadline) {
		got, err := mgr.GetRun(rec.RunID)
		if err == nil && got.Status == "finalizing_failed" {
			evs, listErr := taskrun.NewPGStore(poolA).ListEvents(context.Background(), rec.RunID, 0)
			if listErr != nil {
				t.Fatal(listErr)
			}
			for _, ev := range evs {
				if ev.Event.Type == stream.EventDone {
					t.Fatal("done emitted on failed commit")
				}
			}
			msg, _ := ctxMgr.GetMessageByID(context.Background(), asstID)
			if msg != nil && msg.Content != "任务执行中…" {
				t.Fatalf("placeholder mutated after rollback: %q", msg.Content)
			}
			return
		}
		time.Sleep(30 * time.Millisecond)
	}
	t.Fatal("expected finalizing_failed")
}

func TestPG_CleanupWithoutReceiptFailsAndStagingSurvives(t *testing.T) {
	poolA, _ := openIsolatedM3(t)
	userID, projectID := seedUserAndProject(t, poolA)
	ctxMgr := ctxmgr.NewContextManager(poolA, nil, ctxmgr.TimeDecayConfig{})
	sessionID := "sess_u8_clean"
	asstID, err := ctxMgr.RecordUserMessageWithID(context.Background(), sessionID, userID, "assistant", "任务执行中…", "TRIGGER_TASK", "", map[string]any{"project_id": projectID})
	if err != nil {
		t.Fatal(err)
	}
	dataRoot := t.TempDir()
	cap := workspace.NewCapability(dataRoot, nil)
	ref, err := cap.Issue(workspace.WorkspaceRef{ID: "ws_u8_clean", ProjectScope: projectID, RunID: "run_u8_clean"})
	if err != nil {
		t.Fatal(err)
	}
	body := []byte("# vaulted body for cleanup\n")
	if err := cap.WriteArtifact(context.Background(), ref, workspace.DeliverableKey, body); err != nil {
		t.Fatal(err)
	}
	rec := &taskrun.RunRecord{
		RunID: "run_u8_clean", SessionID: sessionID, ProjectID: projectID, UserID: userID,
		TaskKey: "intake", AssistantMsgID: asstID, WorkspaceID: ref.ID,
	}
	insertRunning(t, poolA, rec)
	res := &bizdev.TaskRunResult{
		Success: true, TaskKey: "intake", MarkdownBody: string(body), WorkDir: ref.ID,
		DocRecord: &docgov.DocumentRecord{DocID: "doc_u8_clean", ContentSHA256: artifactvault.SHA256Hex(body), Version: 2},
	}
	cmds := &taskrun.Commands{UoW: uow.NewPG(poolA), Ctx: ctxMgr, Fav: &artifactvault.Ingest{Env: "dev"}, Env: "dev", StorageIdentity: "env:dev"}
	if _, err := cmds.Finalize(context.Background(), rec, res); err != nil {
		t.Fatal(err)
	}
	var manSHA, storageID string
	var members []byte
	if err := poolA.QueryRow(context.Background(), `SELECT sha256, storage_identity, members FROM task_run_manifests WHERE run_id=$1`, rec.RunID).
		Scan(&manSHA, &storageID, &members); err != nil {
		t.Fatal(err)
	}
	var parsed []taskrun.ManifestMember
	if err := json.Unmarshal(members, &parsed); err != nil {
		t.Fatal(err)
	}
	frozen := workspace.FrozenManifest{SHA256: manSHA, StorageIdentity: storageID, WorkspaceID: ref.ID}
	for _, m := range parsed {
		frozen.Members = append(frozen.Members, workspace.FrozenMember{LogicalID: m.LogicalID, Digest: m.Digest, Kind: m.Kind})
	}
	if err := cap.Cleanup(context.Background(), ref, frozen, nil); err == nil {
		t.Fatal("cleanup without receipts must fail")
	}
	if _, err := os.Stat(filepath.Join(dataRoot, "workspaces", ref.ID, workspace.DeliverableKey)); err != nil {
		t.Fatalf("workspace removed without receipts: %v", err)
	}
	var receipts []workspace.RetentionReceipt
	for _, m := range parsed {
		receipts = append(receipts, workspace.RetentionReceipt{
			LogicalID: m.LogicalID, Digest: m.Digest, StorageIdentity: storageID, Verified: true,
		})
	}
	if err := cap.Cleanup(context.Background(), ref, frozen, receipts); err != nil {
		t.Fatal(err)
	}
	eventID := "taskrun:" + rec.RunID + ":completed"
	reg := artifactvault.NewPGRegistry(poolA)
	st, err := reg.LoadStaging(context.Background(), eventID)
	if err != nil {
		t.Fatal(err)
	}
	if artifactvault.SHA256Hex(st.Payload) != artifactvault.SHA256Hex(body) {
		t.Fatalf("staging sha drifted")
	}
}