package main

import (
	"91porn-server/common"
	"91porn-server/common/conf"
	"91porn-server/common/constant"
	"91porn-server/common/db"
	"91porn-server/common/log"
	"91porn-server/models"
	"91porn-server/models/Init"
	"91porn-server/models/commod"
	"91porn-server/skd/job/dailyAdverCalc"
	"91porn-server/skd/skdg"
	"context"
	"flag"
	"fmt"
	"math"
	"os"
	"regexp"
	"strings"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	defaultConfigPath   = "config/skd.json"
	defaultSeedTag      = "abtest-seed"
	defaultDistrictCode = "ABTEST_SEED"
	defaultVersion      = "1.0.6"
	defaultStartDate    = "2030-01-01"
	defaultUIDBase      = 990000000000000
	batchSize           = 1000
)

type seedConfig struct {
	ConfigPath       string
	SeedTag          string
	DistrictCode     string
	Version          string
	StartDate        string
	CohortDays       int
	ActiveDays       int
	UsersPerGroup    int
	UIDBase          int64
	Cleanup          bool
	RebuildRetention bool
	DryRun           bool
}

type seedSummary struct {
	RegDay       time.Time
	AdGroup      commod.AdGroup
	CohortSize   int
	ActiveCounts []int
}

func main() {
	time.Local, _ = time.LoadLocation("Asia/Shanghai")

	cfg := parseFlags()
	if err := validateConfig(cfg); err != nil {
		fmt.Fprintf(os.Stderr, "invalid config: %v\n", err)
		os.Exit(1)
	}

	log.Init(log.Options{Level: "info"})

	loc, _ := time.LoadLocation("Asia/Shanghai")
	startDate, err := time.ParseInLocation(time.DateOnly, cfg.StartDate, loc)
	if err != nil {
		fmt.Fprintf(os.Stderr, "parse start date failed: %v\n", err)
		os.Exit(1)
	}
	startDate = common.NormalizeDate(startDate)

	if err := initMongo(cfg.ConfigPath); err != nil {
		fmt.Fprintf(os.Stderr, "init mongo failed: %v\n", err)
		os.Exit(1)
	}
	defer db.CloseDS()
	Init.InitMongo()

	users, stats, summaries := buildSeedDocs(cfg, startDate, loc)
	endActiveDate := startDate.AddDate(0, 0, cfg.CohortDays+cfg.ActiveDays-2)

	printPlan(cfg, startDate, endActiveDate, summaries)
	if cfg.DryRun {
		return
	}

	if cfg.Cleanup {
		if err := cleanupSeedData(cfg, startDate, endActiveDate); err != nil {
			fmt.Fprintf(os.Stderr, "cleanup failed: %v\n", err)
			os.Exit(1)
		}
	}

	if err := insertInBatches(models.UserTable, users); err != nil {
		fmt.Fprintf(os.Stderr, "insert users failed: %v\n", err)
		os.Exit(1)
	}
	if err := insertInBatches(models.UserAdverStat, stats); err != nil {
		fmt.Fprintf(os.Stderr, "insert user_adver_stat failed: %v\n", err)
		os.Exit(1)
	}

	if cfg.RebuildRetention {
		if err := rebuildRetentionRange(startDate, endActiveDate); err != nil {
			fmt.Fprintf(os.Stderr, "rebuild retention failed: %v\n", err)
			os.Exit(1)
		}
	}

	fmt.Printf("\nseed completed: users=%d stats=%d rebuildRetention=%v\n", len(users), len(stats), cfg.RebuildRetention)
}

func parseFlags() seedConfig {
	cfg := seedConfig{}
	flag.StringVar(&cfg.ConfigPath, "config", defaultConfigPath, "path to skd json config")
	flag.StringVar(&cfg.SeedTag, "seed-tag", defaultSeedTag, "seed tag used in devID/name prefix")
	flag.StringVar(&cfg.DistrictCode, "district-code", defaultDistrictCode, "district code written into seeded docs")
	flag.StringVar(&cfg.Version, "version", defaultVersion, "android version written into seeded docs")
	flag.StringVar(&cfg.StartDate, "start-date", defaultStartDate, "cohort start date in YYYY-MM-DD")
	flag.IntVar(&cfg.CohortDays, "cohort-days", 5, "number of registration days to generate")
	flag.IntVar(&cfg.ActiveDays, "active-days", 5, "number of active days to generate for each cohort")
	flag.IntVar(&cfg.UsersPerGroup, "users-per-group", 30, "base cohort size per ad group per day")
	flag.Int64Var(&cfg.UIDBase, "uid-base", defaultUIDBase, "starting uid for seeded users")
	flag.BoolVar(&cfg.Cleanup, "cleanup", true, "cleanup previous seeded docs in the target window")
	flag.BoolVar(&cfg.RebuildRetention, "rebuild-retention", true, "run CalcRetention for the generated active date range")
	flag.BoolVar(&cfg.DryRun, "dry-run", false, "print the plan without writing data")
	flag.Parse()
	return cfg
}

func validateConfig(cfg seedConfig) error {
	switch {
	case cfg.CohortDays <= 0:
		return fmt.Errorf("cohort-days must be > 0")
	case cfg.ActiveDays <= 0:
		return fmt.Errorf("active-days must be > 0")
	case cfg.UsersPerGroup <= 0:
		return fmt.Errorf("users-per-group must be > 0")
	case cfg.UIDBase <= 0:
		return fmt.Errorf("uid-base must be > 0")
	case strings.TrimSpace(cfg.SeedTag) == "":
		return fmt.Errorf("seed-tag must not be empty")
	}
	return nil
}

func initMongo(configPath string) error {
	cfg := skdg.GlobalConfig{}
	if err := conf.LoadJSON(configPath, &cfg); err != nil {
		return err
	}

	db.InitDS([]db.DBmap{
		{Key: models.VideoDb, URL: cfg.Mongo.VideoDbUrl},
		{Key: models.StatDb, URL: cfg.Mongo.StatDbUrl},
		{Key: models.LogDb, URL: cfg.Mongo.LogDbUrl},
	}, registerPool())
	return nil
}

func registerPool() []db.Register {
	res := make([]db.Register, len(models.RegisterPool))
	for i, v := range models.RegisterPool {
		res[i] = db.Register{Key: v.Key, Table: v.Table}
	}
	return res
}

func buildSeedDocs(cfg seedConfig, startDate time.Time, loc *time.Location) ([]bson.M, []bson.M, []seedSummary) {
	var (
		users     []bson.M
		stats     []bson.M
		summaries []seedSummary
		nextUID   = cfg.UIDBase
	)

	for dayIdx := 0; dayIdx < cfg.CohortDays; dayIdx++ {
		regDay := startDate.AddDate(0, 0, dayIdx)
		for _, group := range []commod.AdGroup{commod.AdGroupA, commod.AdGroupB, commod.AdGroupC} {
			cohortSize := cohortSizeFor(group, dayIdx, cfg.UsersPerGroup)
			activeCounts := make([]int, cfg.ActiveDays)
			cohortUIDs := make([]int64, 0, cohortSize)

			for userIdx := 0; userIdx < cohortSize; userIdx++ {
				uid := nextUID
				nextUID++
				cohortUIDs = append(cohortUIDs, uid)

				createdAt := time.Date(regDay.Year(), regDay.Month(), regDay.Day(), 10, userIdx%50, 0, 0, loc)
				devID := fmt.Sprintf("%s-%s-%02d-%03d", cfg.SeedTag, group, dayIdx, userIdx)
				users = append(users, bson.M{
					"uid":          uint64(uid),
					"name":         fmt.Sprintf("%s_%s_d%02d_u%03d", cfg.SeedTag, group, dayIdx, userIdx),
					"devID":        devID,
					"districtCode": cfg.DistrictCode,
					"sysType":      constant.SysTypeAndroid,
					"oriVer":       cfg.Version,
					"adGroup":      group,
					"createdAt":    createdAt,
					"updatedAt":    createdAt,
				})
			}

			for activeOffset := 0; activeOffset < cfg.ActiveDays; activeOffset++ {
				activeCount := activeUserCount(cohortSize, groupActivityRate(group, activeOffset))
				activeCounts[activeOffset] = activeCount
				activeDay := common.NormalizeDate(regDay.AddDate(0, 0, activeOffset))

				for userIdx := 0; userIdx < activeCount; userIdx++ {
					uid := cohortUIDs[userIdx]
					adClick := int64(groupWeight(group) + 1 + activeOffset%2)
					appClick := int64(1)
					payCount, payTotal := payForUser(group, activeOffset, userIdx)

					stats = append(stats, bson.M{
						"userId":       uid,
						"date":         activeDay,
						"districtCode": cfg.DistrictCode,
						"sysType":      constant.SysTypeAndroid,
						"adGroup":      group,
						"adClick":      adClick,
						"appClick":     appClick,
						"totalClick":   adClick + appClick,
						"payCount":     payCount,
						"payTotal":     payTotal,
						"regDay":       common.NormalizeDate(regDay),
						"version":      cfg.Version,
						"createdAt":    activeDay.Add(2 * time.Hour),
						"updatedAt":    activeDay.Add(2 * time.Hour),
					})
				}
			}

			summaries = append(summaries, seedSummary{
				RegDay:       common.NormalizeDate(regDay),
				AdGroup:      group,
				CohortSize:   cohortSize,
				ActiveCounts: activeCounts,
			})
		}
	}

	return users, stats, summaries
}

func cohortSizeFor(group commod.AdGroup, dayIdx, base int) int {
	return base + dayIdx*2 + groupWeight(group)*3
}

func groupWeight(group commod.AdGroup) int {
	switch group {
	case commod.AdGroupA:
		return 0
	case commod.AdGroupB:
		return 1
	default:
		return 2
	}
}

func groupActivityRate(group commod.AdGroup, dayOffset int) float64 {
	rates := map[commod.AdGroup][]float64{
		commod.AdGroupA: {0.85, 0.60, 0.45, 0.30, 0.20, 0.12, 0.08},
		commod.AdGroupB: {0.70, 0.50, 0.35, 0.22, 0.14, 0.08, 0.05},
		commod.AdGroupC: {0.55, 0.35, 0.20, 0.10, 0.06, 0.03, 0.02},
	}
	if dayOffset < len(rates[group]) {
		return rates[group][dayOffset]
	}
	last := rates[group][len(rates[group])-1]
	decay := math.Pow(0.7, float64(dayOffset-len(rates[group])+1))
	return last * decay
}

func activeUserCount(cohortSize int, rate float64) int {
	if cohortSize <= 0 || rate <= 0 {
		return 0
	}
	count := int(math.Round(float64(cohortSize) * rate))
	if count < 0 {
		return 0
	}
	if count > cohortSize {
		return cohortSize
	}
	return count
}

func payForUser(group commod.AdGroup, dayOffset, userIdx int) (int64, int64) {
	switch group {
	case commod.AdGroupA:
		if dayOffset <= 2 && userIdx%7 == 0 {
			return 1, 1999
		}
	case commod.AdGroupB:
		if dayOffset <= 1 && userIdx%9 == 0 {
			return 1, 1299
		}
	case commod.AdGroupC:
		if dayOffset == 0 && userIdx%12 == 0 {
			return 1, 699
		}
	}
	return 0, 0
}

func cleanupSeedData(cfg seedConfig, startDate, endActiveDate time.Time) error {
	userFilter := bson.M{
		"devID": bson.M{
			"$regex": fmt.Sprintf("^%s-", regexp.QuoteMeta(cfg.SeedTag)),
		},
	}
	statFilter := bson.M{
		"districtCode": cfg.DistrictCode,
		"version":      cfg.Version,
		"regDay": bson.M{
			"$gte": startDate,
			"$lt":  startDate.AddDate(0, 0, cfg.CohortDays),
		},
		"date": bson.M{
			"$gte": startDate,
			"$lte": endActiveDate,
		},
	}

	if _, err := db.BaseDAO(models.UserTable, nil).DeleteMany(userFilter); err != nil {
		return err
	}
	if _, err := db.BaseDAO(models.UserAdverStat, nil).DeleteMany(statFilter); err != nil {
		return err
	}

	if cfg.RebuildRetention {
		retentionFilter := bson.M{
			"date": bson.M{
				"$gte": startDate,
				"$lt":  startDate.AddDate(0, 0, cfg.CohortDays),
			},
			"adGroup": bson.M{
				"$in": []commod.AdGroup{commod.AdGroupA, commod.AdGroupB, commod.AdGroupC},
			},
		}
		if _, err := db.BaseDAO(models.DailyRetention, nil).DeleteMany(retentionFilter); err != nil {
			return err
		}
	}
	return nil
}

func insertInBatches(table string, docs []bson.M) error {
	if len(docs) == 0 {
		return nil
	}
	dao := db.BaseDAO(table, nil)
	for start := 0; start < len(docs); start += batchSize {
		end := start + batchSize
		if end > len(docs) {
			end = len(docs)
		}
		batch := docs[start:end]
		writes := make([]mongo.WriteModel, 0, len(batch))
		for _, doc := range batch {
			writes = append(writes, mongo.NewInsertOneModel().SetDocument(doc))
		}
		if _, err := dao.Bulk(writes, options.BulkWrite().SetOrdered(false)); err != nil {
			return err
		}
	}
	return nil
}

func rebuildRetentionRange(startDate, endActiveDate time.Time) error {
	ctx := context.Background()
	for d := common.NormalizeDate(startDate); !d.After(endActiveDate); d = d.AddDate(0, 0, 1) {
		if err := dailyAdverCalc.CalcRetention(ctx, d); err != nil {
			return fmt.Errorf("calc retention for %s failed: %w", d.Format(time.DateOnly), err)
		}
	}
	return nil
}

func printPlan(cfg seedConfig, startDate, endActiveDate time.Time, summaries []seedSummary) {
	fmt.Printf("seed tag: %s\n", cfg.SeedTag)
	fmt.Printf("date range: cohort=%s ~ %s, active=%s ~ %s\n",
		startDate.Format(time.DateOnly),
		startDate.AddDate(0, 0, cfg.CohortDays-1).Format(time.DateOnly),
		startDate.Format(time.DateOnly),
		endActiveDate.Format(time.DateOnly),
	)
	fmt.Printf("base users/group/day: %d, cleanup=%v, rebuildRetention=%v, dryRun=%v\n\n",
		cfg.UsersPerGroup, cfg.Cleanup, cfg.RebuildRetention, cfg.DryRun,
	)

	for _, item := range summaries {
		rates := make([]string, 0, len(item.ActiveCounts))
		for _, active := range item.ActiveCounts {
			rates = append(rates, fmt.Sprintf("%.4f", float64(active)/float64(item.CohortSize)))
		}
		fmt.Printf("%s group=%s cohort=%d active=%v rates=%v\n",
			item.RegDay.Format(time.DateOnly), item.AdGroup, item.CohortSize, item.ActiveCounts, rates,
		)
	}
}