package plan

import (
	"bytes"
	"context"
	"fmt"
	"time"

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

type SaveDraftParams struct {
	WorkspaceID      string
	PlanID           string
	ExpectedRevision int64
	Source           string
	Layout           []byte
	PreserveLayout   bool
	Actor            Actor
	Now              time.Time
}

func (store *Store) LoadDraft(ctx context.Context, workspaceID string, planID string) (PlanDraft, error) {
	if err := store.Validate(); err != nil {
		return PlanDraft{}, err
	}
	workspaceID, planID = trim(workspaceID), trim(planID)
	if workspaceID == "" || planID == "" {
		return PlanDraft{}, storeerr.BadRequest("workspaceId and planId are required")
	}
	return loadDraft(ctx, store.db, workspaceID, planID, false)
}

// LoadDraftForUpdate locks one Plan draft through the caller-owned transaction.
func (store *TxStore) LoadDraftForUpdate(ctx context.Context, workspaceID string, planID string) (PlanDraft, error) {
	if err := store.validate(); err != nil {
		return PlanDraft{}, err
	}
	workspaceID, planID = trim(workspaceID), trim(planID)
	if workspaceID == "" || planID == "" {
		return PlanDraft{}, storeerr.BadRequest("workspaceId and planId are required")
	}
	return loadDraft(ctx, store.tx, workspaceID, planID, true)
}

func (store *Store) SaveDraft(ctx context.Context, params SaveDraftParams) (PlanDraft, error) {
	if err := store.Validate(); err != nil {
		return PlanDraft{}, err
	}
	params.WorkspaceID, params.PlanID = trim(params.WorkspaceID), trim(params.PlanID)
	if params.WorkspaceID == "" || params.PlanID == "" || params.ExpectedRevision <= 0 {
		return PlanDraft{}, storeerr.BadRequest("workspaceId, planId, and a positive expectedRevision are required")
	}
	actor, err := validateActor(params.Actor)
	if err != nil {
		return PlanDraft{}, err
	}
	var updated PlanDraft
	err = coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		plan, err := loadPlan(ctx, tx, params.WorkspaceID, params.PlanID, true)
		if err != nil {
			return err
		}
		if plan.ArchivedAt != nil {
			return ErrArchived
		}
		draft, err := loadDraft(ctx, tx, params.WorkspaceID, params.PlanID, true)
		if err != nil {
			return err
		}
		if draft.Revision != params.ExpectedRevision {
			return RevisionConflictError{Resource: "plan draft", Expected: params.ExpectedRevision, Actual: draft.Revision}
		}
		layout := params.Layout
		if params.PreserveLayout {
			layout = draft.Layout
		}
		if draft.Source == params.Source && bytes.Equal(draft.Layout, layout) {
			updated = draft
			return nil
		}
		result, err := tx.Exec(ctx, `
UPDATE workflow_plan_drafts
SET source = $1,
    layout = $2,
    revision = revision + 1,
    updated_by_member_id = $3,
    updated_by_actor_type = $4,
    updated_by_actor_id = $5,
    updated_at = $6
WHERE workspace_id = $7 AND plan_id = $8 AND revision = $9
`, params.Source, string(layout), actor.MemberID, actor.Type, actor.ID,
			utcOrNow(params.Now), params.WorkspaceID, params.PlanID, params.ExpectedRevision)
		if err != nil {
			return fmt.Errorf("save workflow plan draft: %w", mapWriteError(err))
		}
		if result.RowsAffected() == 0 {
			return RevisionConflictError{Resource: "plan draft", Expected: params.ExpectedRevision, Actual: draft.Revision}
		}
		updated, err = loadDraft(ctx, tx, params.WorkspaceID, params.PlanID, false)
		return err
	})
	return updated, err
}

func loadDraft(ctx context.Context, queryer coredb.QueryRower, workspaceID string, planID string, forUpdate bool) (PlanDraft, error) {
	query := `
SELECT plan_id, workspace_id, source, layout, revision,
       updated_by_actor_type, updated_by_actor_id, updated_by_member_id,
       created_at, updated_at
FROM workflow_plan_drafts
WHERE workspace_id = $1 AND plan_id = $2
`
	if forUpdate {
		query += "FOR UPDATE\n"
	}
	var (
		draft      PlanDraft
		layoutText string
	)
	if err := queryer.QueryRow(ctx, query, workspaceID, planID).Scan(
		&draft.PlanID, &draft.WorkspaceID, &draft.Source, &layoutText, &draft.Revision,
		&draft.UpdatedBy.Type, &draft.UpdatedBy.ID, &draft.UpdatedBy.MemberID,
		&draft.CreatedAt, &draft.UpdatedAt,
	); err != nil {
		return PlanDraft{}, mapWriteError(err)
	}
	draft.Layout = []byte(layoutText)
	draft.CreatedAt = draft.CreatedAt.UTC()
	draft.UpdatedAt = draft.UpdatedAt.UTC()
	return draft, nil
}
