package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/orgmembership"
	"github.com/tidbcloud/AgentTeam/internal/tidb"
)

const defaultWorkspaceOrgMembershipBatchSize = 100

type workspaceOrgMembershipOptions struct {
	OrganizationID string
	WorkspaceID    string
	BatchSize      int
	Apply          bool
}

type workspaceOrgMembershipGap struct {
	OrganizationID       string `json:"organization_id"`
	WorkspaceID          string `json:"workspace_id"`
	WorkspaceMemberID    string `json:"workspace_member_id"`
	UserID               string `json:"user_id"`
	AuditUserID          string `json:"audit_user_id"`
	OrganizationMemberID string `json:"organization_member_id,omitempty"`
	MemberStatus         string `json:"member_status,omitempty"`
}

type workspaceOrgMembershipBackfillFailure struct {
	OrganizationID string `json:"organization_id"`
	WorkspaceID    string `json:"workspace_id"`
	UserID         string `json:"user_id"`
	Error          string `json:"error"`
}

type workspaceOrgMembershipBackfillSkip struct {
	OrganizationID string `json:"organization_id"`
	WorkspaceID    string `json:"workspace_id"`
	UserID         string `json:"user_id"`
	MemberID       string `json:"member_id,omitempty"`
	Reason         string `json:"reason"`
}

type workspaceOrgMembershipReport struct {
	DryRun             bool                                    `json:"dry_run"`
	Apply              bool                                    `json:"apply"`
	GapCount           int64                                   `json:"gap_count"`
	RemainingGapCount  int64                                   `json:"remaining_gap_count"`
	CandidateCount     int                                     `json:"candidate_count"`
	ProcessedCount     int                                     `json:"processed_count"`
	CreatedMemberCount int                                     `json:"created_member_count"`
	SkippedCount       int                                     `json:"skipped_count"`
	ErrorCount         int                                     `json:"error_count"`
	Candidates         []workspaceOrgMembershipGap             `json:"candidates,omitempty"`
	Skips              []workspaceOrgMembershipBackfillSkip    `json:"skips,omitempty"`
	Failures           []workspaceOrgMembershipBackfillFailure `json:"failures,omitempty"`
}

func runWorkspaceOrgMembershipVerify(ctx context.Context, args []string) error {
	options, err := parseWorkspaceOrgMembershipOptions("workspace-org-membership-verify", args, false)
	if err != nil {
		return err
	}
	db, closeDB, err := openWorkspaceOrgMembershipDB(ctx)
	if err != nil {
		return err
	}
	defer closeDB()

	gapCount, err := countWorkspaceOrgMembershipGaps(ctx, db, options)
	if err != nil {
		return err
	}
	candidates, err := listWorkspaceOrgMembershipGaps(ctx, db, options)
	if err != nil {
		return err
	}
	report := workspaceOrgMembershipReport{
		DryRun:         true,
		Apply:          false,
		GapCount:       gapCount,
		CandidateCount: len(candidates),
		Candidates:     candidates,
	}
	if err := encodeWorkspaceOrgMembershipReport(report); err != nil {
		return err
	}
	if gapCount != 0 {
		return fmt.Errorf("workspace org membership verify failed: gap_count=%d", gapCount)
	}
	return nil
}

func runWorkspaceOrgMembershipBackfill(ctx context.Context, args []string) error {
	options, err := parseWorkspaceOrgMembershipOptions("workspace-org-membership-backfill", args, true)
	if err != nil {
		return err
	}
	db, closeDB, err := openWorkspaceOrgMembershipDB(ctx)
	if err != nil {
		return err
	}
	defer closeDB()

	gapCount, err := countWorkspaceOrgMembershipGaps(ctx, db, options)
	if err != nil {
		return err
	}
	candidates, err := listWorkspaceOrgMembershipGaps(ctx, db, options)
	if err != nil {
		return err
	}
	report := workspaceOrgMembershipReport{
		DryRun:         !options.Apply,
		Apply:          options.Apply,
		GapCount:       gapCount,
		CandidateCount: len(candidates),
		Candidates:     candidates,
	}
	if options.Apply {
		applyWorkspaceOrgMembershipBackfill(ctx, db, candidates, &report)
		remaining, countErr := countWorkspaceOrgMembershipGaps(ctx, db, options)
		if countErr != nil {
			report.Failures = append(report.Failures, workspaceOrgMembershipBackfillFailure{Error: countErr.Error()})
			report.ErrorCount++
		} else {
			report.RemainingGapCount = remaining
		}
	}
	if err := encodeWorkspaceOrgMembershipReport(report); err != nil {
		return err
	}
	if report.ErrorCount != 0 {
		return fmt.Errorf("workspace org membership backfill failed: error_count=%d", report.ErrorCount)
	}
	return nil
}

func parseWorkspaceOrgMembershipOptions(command string, args []string, includeApply bool) (workspaceOrgMembershipOptions, error) {
	options := workspaceOrgMembershipOptions{BatchSize: defaultWorkspaceOrgMembershipBatchSize}
	flagSet := flag.NewFlagSet(command, flag.ContinueOnError)
	flagSet.StringVar(&options.OrganizationID, "organization-id", "", "optional owning organization id filter")
	flagSet.StringVar(&options.WorkspaceID, "workspace-id", "", "optional workspace id filter")
	flagSet.IntVar(&options.BatchSize, "batch-size", options.BatchSize, "maximum rows to inspect or repair in this invocation")
	if includeApply {
		flagSet.BoolVar(&options.Apply, "apply", false, "write missing organization membership rows")
	}
	if err := flagSet.Parse(args); err != nil {
		return workspaceOrgMembershipOptions{}, err
	}
	if flagSet.NArg() != 0 {
		return workspaceOrgMembershipOptions{}, fmt.Errorf("unexpected %s arguments: %s", command, strings.Join(flagSet.Args(), " "))
	}
	options.OrganizationID = strings.TrimSpace(options.OrganizationID)
	options.WorkspaceID = strings.TrimSpace(options.WorkspaceID)
	if options.BatchSize <= 0 {
		return workspaceOrgMembershipOptions{}, fmt.Errorf("%s requires a positive --batch-size", command)
	}
	return options, nil
}

func applyWorkspaceOrgMembershipBackfill(ctx context.Context, db coredb.DB, candidates []workspaceOrgMembershipGap, report *workspaceOrgMembershipReport) {
	for _, candidate := range candidates {
		result, err := applyOneWorkspaceOrgMembershipBackfill(ctx, db, candidate)
		if err != nil {
			report.Failures = append(report.Failures, workspaceOrgMembershipBackfillFailure{
				OrganizationID: candidate.OrganizationID,
				WorkspaceID:    candidate.WorkspaceID,
				UserID:         candidate.UserID,
				Error:          err.Error(),
			})
			report.ErrorCount++
			continue
		}
		if result.Skipped {
			report.SkippedCount++
			report.Skips = append(report.Skips, workspaceOrgMembershipBackfillSkip{
				OrganizationID: candidate.OrganizationID,
				WorkspaceID:    candidate.WorkspaceID,
				UserID:         candidate.UserID,
				MemberID:       result.MemberID,
				Reason:         result.SkipReason,
			})
			continue
		}
		report.ProcessedCount++
		if result.CreatedMember {
			report.CreatedMemberCount++
		}
	}
}

func applyOneWorkspaceOrgMembershipBackfill(ctx context.Context, db coredb.DB, candidate workspaceOrgMembershipGap) (orgmembership.WorkspaceBackfillMemberResult, error) {
	tx, err := db.BeginTx(ctx, coredb.TxOptions{})
	if err != nil {
		return orgmembership.WorkspaceBackfillMemberResult{}, fmt.Errorf("begin workspace org membership backfill tx: %w", err)
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	result, err := orgmembership.EnsureWorkspaceBackfillMemberInTx(ctx, tx, orgmembership.WorkspaceBackfillMemberParams{
		OrganizationID: candidate.OrganizationID,
		UserID:         candidate.UserID,
		AuditUserID:    candidate.AuditUserID,
	})
	if err != nil {
		return orgmembership.WorkspaceBackfillMemberResult{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return orgmembership.WorkspaceBackfillMemberResult{}, fmt.Errorf("commit workspace org membership backfill tx: %w", err)
	}
	return result, nil
}

func countWorkspaceOrgMembershipGaps(ctx context.Context, db coredb.QueryRower, options workspaceOrgMembershipOptions) (int64, error) {
	return orgmembership.CountWorkspaceBackfillGaps(ctx, db, workspaceOrgMembershipGapOptions(options))
}

func listWorkspaceOrgMembershipGaps(ctx context.Context, db coredb.Queryer, options workspaceOrgMembershipOptions) ([]workspaceOrgMembershipGap, error) {
	storeGaps, err := orgmembership.ListWorkspaceBackfillGaps(ctx, db, workspaceOrgMembershipGapOptions(options))
	if err != nil {
		return nil, err
	}
	gaps := make([]workspaceOrgMembershipGap, 0, len(storeGaps))
	for _, gap := range storeGaps {
		gaps = append(gaps, workspaceOrgMembershipGap{
			OrganizationID:       gap.OrganizationID,
			WorkspaceID:          gap.WorkspaceID,
			WorkspaceMemberID:    gap.WorkspaceMemberID,
			UserID:               gap.UserID,
			AuditUserID:          gap.AuditUserID,
			OrganizationMemberID: gap.OrganizationMemberID,
			MemberStatus:         gap.MemberStatus,
		})
	}
	return gaps, nil
}

func workspaceOrgMembershipGapOptions(options workspaceOrgMembershipOptions) orgmembership.WorkspaceBackfillGapOptions {
	return orgmembership.WorkspaceBackfillGapOptions{
		OrganizationID: options.OrganizationID,
		WorkspaceID:    options.WorkspaceID,
		Limit:          options.BatchSize,
	}
}

func encodeWorkspaceOrgMembershipReport(report workspaceOrgMembershipReport) error {
	encoder := json.NewEncoder(os.Stdout)
	encoder.SetEscapeHTML(false)
	if err := encoder.Encode(report); err != nil {
		return fmt.Errorf("encode workspace org membership report: %w", err)
	}
	return nil
}

func openWorkspaceOrgMembershipDB(ctx context.Context) (coredb.DB, func(), error) {
	databaseURL, err := currentDatabaseURL()
	if err != nil {
		return nil, nil, err
	}
	connectTimeout, err := parseConnectTimeout()
	if err != nil {
		return nil, nil, err
	}
	admin := tidb.NewAdmin()
	db, err := admin.Open(ctx, databaseURL, connectTimeout)
	if err != nil {
		return nil, nil, fmt.Errorf("connect core main-path database for workspace org membership repair: %w", err)
	}
	if err := admin.RunStartupPreflight(ctx, db); err != nil {
		_ = db.Close()
		return nil, nil, err
	}
	return coredb.WrapSQLDB(db), func() { _ = db.Close() }, nil
}