package main

import (
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"codex-pool/internal/auth"
	"codex-pool/internal/config"
	"codex-pool/internal/invite"
	"codex-pool/internal/mail"
	"codex-pool/internal/register"
)

type Account struct {
	Account  string `json:"account"`
	Password string `json:"password"`
	Token    string `json:"token"`
}

type MemberAccount struct {
	Email    string
	Password string
	Success  bool
}

const (
	MembersPerTeam = 4 // 每个 team 注册的成员数
	NumTeams       = 2 // 并发运行的 team 数量
)

// ANSI 颜色码
const (
	ColorReset   = "\033[0m"
	ColorRed     = "\033[31m"
	ColorGreen   = "\033[32m"
	ColorYellow  = "\033[33m"
	ColorBlue    = "\033[34m"
	ColorMagenta = "\033[35m"
	ColorCyan    = "\033[36m"
	ColorWhite   = "\033[37m"
	ColorBold    = "\033[1m"
)

// Team 颜色
var teamColors = []string{
	ColorCyan,    // Team 1
	ColorMagenta, // Team 2
	ColorYellow,  // Team 3
	ColorBlue,    // Team 4
}

// TeamLogger 带颜色的Team日志
type TeamLogger struct {
	prefix string
	color  string
	mu     sync.Mutex
}

func NewTeamLogger(teamIdx int) *TeamLogger {
	color := teamColors[teamIdx%len(teamColors)]
	return &TeamLogger{
		prefix: fmt.Sprintf("[Team %d]", teamIdx+1),
		color:  color,
	}
}

func (l *TeamLogger) Log(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("%s%s%s %s\n", l.color, l.prefix, ColorReset, msg)
}

func (l *TeamLogger) Success(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("%s%s%s %s✓%s %s\n", l.color, l.prefix, ColorReset, ColorGreen, ColorReset, msg)
}

func (l *TeamLogger) Error(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("%s%s%s %s✗%s %s\n", l.color, l.prefix, ColorReset, ColorRed, ColorReset, msg)
}

func (l *TeamLogger) Info(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("%s%s%s %s→%s %s\n", l.color, l.prefix, ColorReset, ColorYellow, ColorReset, msg)
}

// Highlight 整行绿色高亮（用于重要成功信息）
func (l *TeamLogger) Highlight(format string, args ...interface{}) {
	l.mu.Lock()
	defer l.mu.Unlock()
	msg := fmt.Sprintf(format, args...)
	fmt.Printf("%s%s %s✓ %s%s\n", ColorGreen, l.prefix, ColorBold, msg, ColorReset)
}

func main() {
	fmt.Printf("%s%s=================================================================%s\n", ColorBold, ColorWhite, ColorReset)
	fmt.Printf("%s  Multi-Team Concurrent Test (Chromedp)%s\n", ColorBold, ColorReset)
	fmt.Printf("  - %d Teams running concurrently\n", NumTeams)
	fmt.Printf("  - %d Members per team\n", MembersPerTeam)
	fmt.Printf("%s%s=================================================================%s\n", ColorBold, ColorWhite, ColorReset)
	fmt.Println()

	// 加载配置
	configPath := config.FindPath()
	cfg, err := config.Load(configPath)
	if err != nil {
		fmt.Printf("%s[Error]%s Failed to load config: %v\n", ColorRed, ColorReset, err)
		os.Exit(1)
	}

	// 初始化邮箱服务
	if len(cfg.MailServices) > 0 {
		mail.Init(cfg.MailServices)
	}

	// 加载账号
	accountsFile := "accounts-3-20260130-052841.json"
	data, err := os.ReadFile(accountsFile)
	if err != nil {
		fmt.Printf("%s[Error]%s Failed to read accounts file: %v\n", ColorRed, ColorReset, err)
		os.Exit(1)
	}

	var accounts []Account
	if err := json.Unmarshal(data, &accounts); err != nil {
		fmt.Printf("%s[Error]%s Failed to parse accounts file: %v\n", ColorRed, ColorReset, err)
		os.Exit(1)
	}

	if len(accounts) < NumTeams {
		fmt.Printf("%s[Error]%s Need at least %d owner accounts, got %d\n", ColorRed, ColorReset, NumTeams, len(accounts))
		os.Exit(1)
	}

	proxy := cfg.DefaultProxy
	if proxy == "" {
		proxy = "http://127.0.0.1:7890"
	}
	fmt.Printf("[Proxy] %s\n", proxy)
	fmt.Println()

	// 显示 Owner 列表
	fmt.Printf("%s========================================%s\n", ColorBold, ColorReset)
	fmt.Printf("%s[Owners]%s\n", ColorBold, ColorReset)
	fmt.Printf("%s========================================%s\n", ColorBold, ColorReset)
	for i := 0; i < NumTeams; i++ {
		color := teamColors[i%len(teamColors)]
		fmt.Printf("  %sTeam %d:%s %s\n", color, i+1, ColorReset, accounts[i].Account)
	}
	fmt.Println()

	// 并发运行多个 Team
	var wg sync.WaitGroup
	var totalRegistered int32
	var totalS2A int32
	startTime := time.Now()

	for teamIdx := 0; teamIdx < NumTeams; teamIdx++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			registered, s2a := runTeam(idx, accounts[idx], cfg, proxy)
			atomic.AddInt32(&totalRegistered, int32(registered))
			atomic.AddInt32(&totalS2A, int32(s2a))
		}(teamIdx)
	}

	wg.Wait()
	totalDuration := time.Since(startTime)

	// 总结
	fmt.Println()
	fmt.Printf("%s%s=================================================================%s\n", ColorBold, ColorWhite, ColorReset)
	fmt.Printf("%s  All Teams Complete%s\n", ColorBold, ColorReset)
	fmt.Printf("%s=================================================================%s\n", ColorBold, ColorReset)
	fmt.Printf("  Total Registered: %s%d/%d%s\n", ColorGreen, totalRegistered, NumTeams*MembersPerTeam, ColorReset)
	fmt.Printf("  Total Added to S2A: %s%d%s\n", ColorGreen, totalS2A, ColorReset)
	fmt.Printf("  Total Duration: %v\n", totalDuration)
	fmt.Printf("%s=================================================================%s\n", ColorBold, ColorReset)
}

// runTeam 运行单个 Team 的流程
func runTeam(teamIdx int, owner Account, cfg *config.Config, proxy string) (registered, s2a int) {
	log := NewTeamLogger(teamIdx)

	log.Log("Starting with owner: %s", owner.Account)

	// Step 1: 获取 Team ID
	log.Info("Fetching Team ID...")
	inviter := invite.NewWithProxy(owner.Token, proxy)
	teamID, err := inviter.GetAccountID()
	if err != nil {
		log.Error("Failed to get Team ID: %v", err)
		return 0, 0
	}
	log.Success("Team ID: %s", teamID)

	// Step 2: 生成成员邮箱
	log.Info("Generating %d member emails...", MembersPerTeam)
	children := make([]MemberAccount, MembersPerTeam)
	for i := 0; i < MembersPerTeam; i++ {
		children[i].Email = mail.GenerateEmail()
		children[i].Password = register.GeneratePassword()
		log.Log("[Member %d] Email: %s", i+1, children[i].Email)
	}

	// 批量发送邀请
	log.Info("Sending invites...")
	inviteEmails := make([]string, MembersPerTeam)
	for i, c := range children {
		inviteEmails[i] = c.Email
	}
	if err := inviter.SendInvites(inviteEmails); err != nil {
		log.Error("Failed to send invites: %v", err)
		return 0, 0
	}
	log.Success("Sent %d invite(s)", len(inviteEmails))

	// Step 3: 并发注册成员
	log.Info("Starting member registration...")
	var memberWg sync.WaitGroup
	var successCount int32
	memberMutex := sync.Mutex{}

	for i := range children {
		memberWg.Add(1)
		go func(memberIdx int) {
			defer memberWg.Done()

			memberMutex.Lock()
			email := children[memberIdx].Email
			password := children[memberIdx].Password
			memberMutex.Unlock()

			name := register.GenerateName()
			birthdate := register.GenerateBirthdate()

			// 最多重试3次
			for attempt := 0; attempt < 3; attempt++ {
				if attempt > 0 {
					email = mail.GenerateEmail()
					password = register.GeneratePassword()
					log.Log("[Member %d] Retry %d - New email: %s", memberIdx+1, attempt, email)

					// 发送新邀请
					if err := inviter.SendInvites([]string{email}); err != nil {
						log.Error("[Member %d] Failed to send retry invite: %v", memberIdx+1, err)
						continue
					}
					log.Success("[Member %d] Sent retry invite", memberIdx+1)
				}

				// 详细注册流程
				if err := registerMemberDetailed(log, memberIdx+1, email, password, name, birthdate, proxy); err != nil {
					if strings.Contains(err.Error(), "验证码") {
						log.Error("[Member %d] OTP timeout, will retry...", memberIdx+1)
						continue
					}
					log.Error("[Member %d] Registration failed: %v", memberIdx+1, err)
					return
				}

				// 成功
				memberMutex.Lock()
				children[memberIdx].Email = email
				children[memberIdx].Password = password
				children[memberIdx].Success = true
				memberMutex.Unlock()

				atomic.AddInt32(&successCount, 1)
				log.Success("[Member %d] Registration complete!", memberIdx+1)
				return
			}

			log.Error("[Member %d] Failed after 3 retries", memberIdx+1)
		}(i)
	}

	memberWg.Wait()
	registered = int(successCount)
	log.Success("Registration phase complete: %d/%d", registered, MembersPerTeam)

	// 收集成功的成员
	var registeredChildren []MemberAccount
	for _, c := range children {
		if c.Success {
			registeredChildren = append(registeredChildren, c)
		}
	}

	if len(registeredChildren) == 0 {
		log.Error("No members registered")
		return registered, 0
	}

	// Step 4: 串行入库
	log.Info("Starting S2A authorization...")

	for i, child := range registeredChildren {
		log.Log("[Member %d] Getting S2A auth URL...", i+1)
		s2aResp, err := auth.GenerateS2AAuthURL(cfg.S2AApiBase, cfg.S2AAdminKey, cfg.ProxyID)
		if err != nil {
			log.Error("[Member %d] Auth URL failed: %v", i+1, err)
			continue
		}
		log.Success("[Member %d] Got auth URL", i+1)

		log.Log("[Member %d] Running browser automation (Chromedp)...", i+1)
		code, err := auth.CompleteWithChromedp(s2aResp.Data.AuthURL, child.Email, child.Password, teamID, true, proxy)
		if err != nil {
			log.Error("[Member %d] Browser auth failed: %v", i+1, err)
			continue
		}
		log.Success("[Member %d] Browser auth complete", i+1)

		log.Log("[Member %d] Submitting to S2A...", i+1)
		result, err := auth.SubmitS2AOAuth(
			cfg.S2AApiBase,
			cfg.S2AAdminKey,
			s2aResp.Data.SessionID,
			code,
			child.Email,
			cfg.Concurrency,
			cfg.Priority,
			cfg.GroupIDs,
			cfg.ProxyID,
		)
		if err != nil {
			log.Error("[Member %d] S2A submit failed: %v", i+1, err)
			continue
		}

		log.Highlight("[Member %d] Added to S2A! ID=%d, Status=%s", i+1, result.Data.ID, result.Data.Status)
		s2a++

		time.Sleep(500 * time.Millisecond)
	}

	log.Success("Team complete: %d registered, %d in S2A", registered, s2a)
	return registered, s2a
}

// registerMemberDetailed 详细的注册流程，带日志
func registerMemberDetailed(log *TeamLogger, memberNum int, email, password, name, birthdate, proxy string) error {
	prefix := fmt.Sprintf("[Member %d]", memberNum)

	log.Log("%s Creating TLS client...", prefix)
	reg, err := register.New(proxy)
	if err != nil {
		return err
	}

	log.Log("%s Initializing session...", prefix)
	if err := reg.InitSession(); err != nil {
		return fmt.Errorf("初始化失败: %v", err)
	}
	log.Success("%s Session initialized", prefix)

	log.Log("%s Getting authorize URL...", prefix)
	if err := reg.GetAuthorizeURL(email); err != nil {
		return fmt.Errorf("获取授权URL失败: %v", err)
	}
	log.Success("%s Got authorize URL", prefix)

	log.Log("%s Starting authorize flow...", prefix)
	if err := reg.StartAuthorize(); err != nil {
		return fmt.Errorf("启动授权失败: %v", err)
	}
	log.Success("%s Authorize flow started", prefix)

	log.Log("%s Registering account...", prefix)
	if err := reg.Register(email, password); err != nil {
		return fmt.Errorf("注册失败: %v", err)
	}
	log.Success("%s Account registered", prefix)

	log.Log("%s Sending verification email...", prefix)
	if err := reg.SendVerificationEmail(); err != nil {
		return fmt.Errorf("发送邮件失败: %v", err)
	}
	log.Success("%s Verification email sent", prefix)

	log.Log("%s Waiting for OTP code (5s timeout)...", prefix)
	otpCode, err := mail.GetVerificationCode(email, 5*time.Second)
	if err != nil {
		log.Log("%s OTP not received in 5s, waiting 15s more...", prefix)
		otpCode, err = mail.GetVerificationCode(email, 15*time.Second)
		if err != nil {
			return fmt.Errorf("验证码获取超时")
		}
	}
	log.Success("%s Got OTP: %s", prefix, otpCode)

	log.Log("%s Validating OTP...", prefix)
	if err := reg.ValidateOTP(otpCode); err != nil {
		return fmt.Errorf("OTP验证失败: %v", err)
	}
	log.Success("%s OTP validated", prefix)

	log.Log("%s Creating account (name=%s, birthdate=%s)...", prefix, name, birthdate)
	if err := reg.CreateAccount(name, birthdate); err != nil {
		return fmt.Errorf("创建账户失败: %v", err)
	}
	log.Success("%s Account created", prefix)

	log.Log("%s Getting session token...", prefix)
	_ = reg.GetSessionToken()
	if reg.AccessToken != "" {
		log.Success("%s Got access token: %s...", prefix, truncate(reg.AccessToken, 30))
	}

	return nil
}

func truncate(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen]
}