package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"sort"
	"sync"
	"sync/atomic"
	"time"
)

func runPrepare(ctx context.Context, cfg *Config, manifestPath string) error {
	apiKey, err := cfg.apiKey()
	if err != nil {
		return err
	}
	password, err := cfg.password()
	if err != nil {
		return err
	}
	verificationCode, err := cfg.verificationCode()
	if err != nil {
		return err
	}
	manifest := newManifest(cfg)
	if existing, loadErr := loadManifest(manifestPath); loadErr == nil {
		manifest, err = reconcileManifest(cfg, existing)
		if err != nil {
			return err
		}
		log.Printf("resuming manifest with %d existing accounts, target=%d", len(existing.Accounts), len(manifest.Accounts))
	}
	if err = applyConfigEndpoints(cfg, manifest); err != nil {
		return err
	}
	api := newAPIClient(cfg, apiKey)
	if err = prepareAccounts(ctx, cfg, manifest, api, password, verificationCode); err != nil {
		return err
	}
	if err = saveManifest(manifestPath, manifest); err != nil {
		return err
	}
	if err = prepareFriendships(ctx, cfg, manifest, api); err != nil {
		return err
	}
	if err = saveManifest(manifestPath, manifest); err != nil {
		return err
	}
	if err = prepareGroup(ctx, cfg, manifest, api); err != nil {
		return err
	}
	if err = saveManifest(manifestPath, manifest); err != nil {
		return err
	}
	return runVerify(ctx, cfg, manifest, api)
}

func prepareAccounts(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient, password, verificationCode string) error {
	var created atomic.Int64
	var existing atomic.Int64
	var mu sync.Mutex
	err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
		if account.UserID != "" && account.Token != "" && account.TokenExpires.After(time.Now().Add(5*time.Minute)) {
			existing.Add(1)
			return nil
		}
		client, wasCreated, err := connectForPrepareWithRetry(ctx, cfg, account, api, password, verificationCode)
		if err != nil {
			return fmt.Errorf("account %s: %w", account.Username, err)
		}
		updated := client.Account()
		client.Close()
		mu.Lock()
		defer mu.Unlock()
		for i := range manifest.Accounts {
			if manifest.Accounts[i].Machine == account.Machine && manifest.Accounts[i].Index == account.Index {
				manifest.Accounts[i].UserID = updated.UserID
				manifest.Accounts[i].Token = updated.Token
				manifest.Accounts[i].TokenExpires = updated.TokenExpires
				break
			}
		}
		if wasCreated {
			created.Add(1)
		} else {
			existing.Add(1)
		}
		return nil
	})
	if err == nil {
		log.Printf("accounts ready=%d created=%d existing=%d", len(manifest.Accounts), created.Load(), existing.Load())
	}
	return err
}

func connectForPrepareWithRetry(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient,
	password, verificationCode string) (*WSClient, bool, error) {
	var lastErr error
	for attempt := 1; attempt <= 6; attempt++ {
		client, created, err := connectForPrepare(ctx, cfg, account, api, password, verificationCode)
		if err == nil {
			return client, created, nil
		}
		lastErr = err
		if ctx.Err() != nil {
			return nil, false, ctx.Err()
		}
		delay := time.Duration(attempt*5) * time.Second
		log.Printf("account=%s prepare attempt=%d failed; retry_in=%s: %v", account.Username, attempt, delay, err)
		if err = waitContext(ctx, delay); err != nil {
			return nil, false, err
		}
	}
	return nil, false, fmt.Errorf("account %s prepare failed after retries: %w", account.Username, lastErr)
}

func prepareFriendships(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
	accounts := manifestAccountMap(manifest)
	pairs := make([]AccountRecord, 0, cfg.UsersPerMachine)
	for i := 1; i <= cfg.UsersPerMachine; i++ {
		pairs = append(pairs, *accounts[accountKey("a", i)])
	}
	var mu sync.Mutex
	err := parallelAccounts(ctx, pairs, cfg.PrepareWorkers, func(a AccountRecord) error {
		return retryPrepareOperation(ctx, fmt.Sprintf("friend-pair-%04d", a.Index), func() error {
			b := *accounts[accountKey("b", a.Index)]
			if a.UserID == "" || b.UserID == "" {
				return fmt.Errorf("pair %d has missing user ID", a.Index)
			}
			conversation, conversationErr := api.ensureP2P(ctx, a, b)
			if conversationErr == nil && conversation.Relationship == "friend" && conversation.Topic != "" {
				mu.Lock()
				accounts[accountKey("a", a.Index)].P2PTopic = conversation.Topic
				accounts[accountKey("b", a.Index)].P2PTopic = conversation.Topic
				mu.Unlock()
				return nil
			}
			if conversationErr != nil {
				var apiErr *apiError
				if !errors.As(conversationErr, &apiErr) || apiErr.Status != http.StatusForbidden {
					return fmt.Errorf("check friend pair %d: %w", a.Index, conversationErr)
				}
			}
			request, requestErr := api.createFriendRequest(ctx, a, b)
			if requestErr == nil && request.Status == "pending" {
				if _, requestErr = api.acceptFriendRequest(ctx, b, request.RequestID); requestErr != nil {
					return fmt.Errorf("accept pair %d: %w", a.Index, requestErr)
				}
			} else if requestErr != nil {
				var apiErr *apiError
				if !errors.As(requestErr, &apiErr) || apiErr.Status != http.StatusConflict {
					return fmt.Errorf("friend pair %d: %w", a.Index, requestErr)
				}
				pendingID, findErr := api.findPendingFriendRequest(ctx, b, a.UserID)
				if findErr != nil {
					return fmt.Errorf("find pending friend pair %d: %w", a.Index, findErr)
				}
				if pendingID != "" {
					if _, acceptErr := api.acceptFriendRequest(ctx, b, pendingID); acceptErr != nil {
						return fmt.Errorf("resume friend pair %d: %w", a.Index, acceptErr)
					}
				}
			}
			conversation, err := api.ensureP2P(ctx, a, b)
			if err != nil {
				return fmt.Errorf("ensure P2P pair %d: %w", a.Index, err)
			}
			if conversation.Relationship != "friend" || conversation.Topic == "" {
				return fmt.Errorf("pair %d relationship=%s topic=%q", a.Index, conversation.Relationship, conversation.Topic)
			}
			mu.Lock()
			accounts[accountKey("a", a.Index)].P2PTopic = conversation.Topic
			accounts[accountKey("b", a.Index)].P2PTopic = conversation.Topic
			mu.Unlock()
			return nil
		})
	})
	if err == nil {
		log.Printf("friendships ready=%d", cfg.UsersPerMachine)
	}
	return err
}

func prepareGroup(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
	if err := parallelAccounts(ctx, manifest.Accounts, cfg.PrepareWorkers, func(account AccountRecord) error {
		return retryPrepareOperation(ctx, "privacy-"+account.Username, func() error {
			return api.disableGroupInviteVerification(ctx, account)
		})
	}); err != nil {
		return fmt.Errorf("prepare group invite privacy: %w", err)
	}
	accounts := manifestAccountMap(manifest)
	owner := *accounts[accountKey("a", 1)]
	if manifest.GroupTopic == "" {
		group, err := api.createGroup(ctx, owner)
		if err != nil {
			return err
		}
		manifest.GroupTopic = group.Topic
		log.Printf("created load-test group topic_hash=%s", topicHash(group.Topic))
	}
	memberIDs := make([]string, 0, len(manifest.Accounts)-1)
	for _, account := range manifest.Accounts {
		if account.UserID != owner.UserID {
			memberIDs = append(memberIDs, account.UserID)
		}
	}
	sort.Strings(memberIDs)
	for offset, batch := 0, 1; offset < len(memberIDs); offset, batch = offset+100, batch+1 {
		end := offset + 100
		if end > len(memberIDs) {
			end = len(memberIDs)
		}
		var response *groupInvitationResponse
		err := retryPrepareOperation(ctx, fmt.Sprintf("group-invite-%03d", batch), func() error {
			var inviteErr error
			response, inviteErr = api.inviteGroup(ctx, owner, manifest.GroupTopic, batch, memberIDs[offset:end])
			return inviteErr
		})
		if err != nil {
			return fmt.Errorf("invite group batch %d: %w", batch, err)
		}
		for _, item := range response.Items {
			if item.Status != "joined" && item.Status != "already_member" {
				return fmt.Errorf("invite user %s status=%s reason=%s", item.UserID, item.Status, item.Reason)
			}
		}
	}
	count, err := api.countGroupMembers(ctx, owner, manifest.GroupTopic)
	if err != nil {
		return err
	}
	if count != len(manifest.Accounts) {
		return fmt.Errorf("group has %d members, expected %d", count, len(manifest.Accounts))
	}
	log.Printf("group ready members=%d topic_hash=%s", count, topicHash(manifest.GroupTopic))
	return nil
}

func retryPrepareOperation(ctx context.Context, label string, operation func() error) error {
	var lastErr error
	for attempt := 1; attempt <= 6; attempt++ {
		if err := operation(); err == nil {
			return nil
		} else {
			lastErr = err
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		delay := time.Duration(attempt*5) * time.Second
		log.Printf("operation=%s attempt=%d failed; retry_in=%s: %v", label, attempt, delay, lastErr)
		if err := waitContext(ctx, delay); err != nil {
			return err
		}
	}
	return fmt.Errorf("operation %s failed after retries: %w", label, lastErr)
}

func runVerify(ctx context.Context, cfg *Config, manifest *Manifest, api *apiClient) error {
	if len(manifest.Accounts) != cfg.UsersPerMachine*2 {
		return fmt.Errorf("manifest accounts=%d expected=%d", len(manifest.Accounts), cfg.UsersPerMachine*2)
	}
	for _, account := range manifest.Accounts {
		if account.UserID == "" || account.Token == "" || account.P2PTopic == "" {
			return fmt.Errorf("account %s is incomplete", account.Username)
		}
	}
	if manifest.GroupTopic == "" {
		return errors.New("manifest group topic is empty")
	}
	owner := manifestAccountMap(manifest)[accountKey("a", 1)]
	count, err := api.countGroupMembers(ctx, *owner, manifest.GroupTopic)
	if err != nil {
		return err
	}
	if count != len(manifest.Accounts) {
		return fmt.Errorf("group member count=%d expected=%d", count, len(manifest.Accounts))
	}
	log.Printf("verify passed accounts=%d pairs=%d group_members=%d", len(manifest.Accounts), cfg.UsersPerMachine, count)
	return nil
}

func parallelAccounts(ctx context.Context, accounts []AccountRecord, workers int, fn func(AccountRecord) error) error {
	if workers <= 0 {
		workers = 1
	}
	workCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	jobs := make(chan AccountRecord)
	errCh := make(chan error, 1)
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for account := range jobs {
				if err := fn(account); err != nil {
					select {
					case errCh <- err:
					default:
					}
					cancel()
					return
				}
			}
		}()
	}
	go func() {
		defer close(jobs)
		for _, account := range accounts {
			select {
			case jobs <- account:
			case <-workCtx.Done():
				return
			}
		}
	}()
	wg.Wait()
	select {
	case err := <-errCh:
		return err
	default:
		return ctx.Err()
	}
}