package loadharness

import (
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sort"
	"strings"
	"time"
)

const (
	DatasetVersion       = 1
	maxDatasetGroups     = 10000
	maxDatasetMembership = 2_000_000
	maxDatasetMessages   = 2_000_000
)

// DatasetConfig describes a deterministic social graph whose durable facts are
// later materialized exclusively through real MTProto RPCs. Planning does not
// contact the server and never embeds auth/session material.
type DatasetConfig struct {
	Accounts      int   `json:"accounts"`
	Seed          int64 `json:"seed"`
	PrivateFanout int   `json:"private_fanout"`
	HotGroups     int   `json:"hot_groups"`
	HotMembers    int   `json:"hot_members"`
	HotHistory    int   `json:"hot_history"`
	MediumGroups  int   `json:"medium_groups"`
	MediumMembers int   `json:"medium_members"`
	MediumHistory int   `json:"medium_history"`
	SmallGroups   int   `json:"small_groups"`
	SmallMembers  int   `json:"small_members"`
	SmallHistory  int   `json:"small_history"`
	HeavyGroups   int   `json:"heavy_groups"`
	HeavyAccounts int   `json:"heavy_accounts"`
	HeavyHistory  int   `json:"heavy_history"`
}

func DefaultDatasetConfig(accounts int) DatasetConfig {
	return DatasetConfig{
		Accounts: accounts, Seed: 20260827, PrivateFanout: min(10, max(accounts-1, 0)),
		HotGroups: 10, HotMembers: accounts, HotHistory: 100,
		MediumGroups: 100, MediumMembers: min(100, accounts), MediumHistory: 30,
		SmallGroups: 200, SmallMembers: min(20, accounts), SmallHistory: 10,
		HeavyGroups: 200, HeavyAccounts: min(100, accounts), HeavyHistory: 30,
	}
}

type DatasetPrivateEdge struct {
	SenderAccount    int    `json:"sender_account"`
	RecipientAccount int    `json:"recipient_account"`
	RandomID         int64  `json:"random_id"`
	Marker           string `json:"marker"`
}

type DatasetGroup struct {
	Index           int    `json:"index"`
	Tier            string `json:"tier"`
	Title           string `json:"title"`
	About           string `json:"about"`
	CreatorAccount  int    `json:"creator_account"`
	MemberAccounts  []int  `json:"member_accounts"`
	HistoryMessages int    `json:"history_messages"`
}

// Dataset is the immutable plan. Real server identities and resumable progress
// live in the separate compact DatasetSeedState so a large plan is not rewritten
// after every reconciled RPC batch.
type Dataset struct {
	Version      int                  `json:"version"`
	CreatedAt    time.Time            `json:"created_at"`
	RunID        string               `json:"run_id"`
	Config       DatasetConfig        `json:"config"`
	PlanSHA256   string               `json:"plan_sha256"`
	PrivateEdges []DatasetPrivateEdge `json:"private_edges"`
	Groups       []DatasetGroup       `json:"groups"`
}

type DatasetSeedGroupState struct {
	GroupIndex       int   `json:"group_index"`
	ChannelID        int64 `json:"channel_id,omitempty"`
	AccessHash       int64 `json:"access_hash,omitempty"`
	CreatePending    bool  `json:"create_pending,omitempty"`
	InviteCursor     int   `json:"invite_cursor,omitempty"`
	InvitePendingEnd int   `json:"invite_pending_end,omitempty"`
}

type DatasetSeedState struct {
	Version              int                     `json:"version"`
	PlanSHA256           string                  `json:"plan_sha256"`
	UpdatedAt            time.Time               `json:"updated_at"`
	PrivateSentByAccount []int                   `json:"private_sent_by_account"`
	HistorySentByAccount []int                   `json:"history_sent_by_account"`
	RichStateByAccount   []bool                  `json:"rich_state_by_account,omitempty"`
	Groups               []DatasetSeedGroupState `json:"groups"`
}

func PlanDataset(cfg DatasetConfig) (*Dataset, error) {
	if err := cfg.validate(); err != nil {
		return nil, err
	}
	runID := fmt.Sprintf("rpc-startup-%d-%016x", cfg.Accounts, uint64(cfg.Seed))
	dataset := &Dataset{
		Version: DatasetVersion, CreatedAt: time.Now().UTC(),
		RunID: runID, Config: cfg,
		PrivateEdges: make([]DatasetPrivateEdge, 0, cfg.Accounts*cfg.PrivateFanout),
	}
	for account := 0; account < cfg.Accounts; account++ {
		for offset := 1; offset <= cfg.PrivateFanout; offset++ {
			recipient := (account + offset) % cfg.Accounts
			dataset.PrivateEdges = append(dataset.PrivateEdges, DatasetPrivateEdge{
				SenderAccount: account, RecipientAccount: recipient,
				RandomID: stableDatasetID(cfg.Seed, "private", cfg.Accounts, account, offset),
				Marker:   fmt.Sprintf("[%s private %04d/%02d]", runID, account, offset),
			})
		}
	}
	groupIndex := 0
	appendTier := func(tier string, count, members, history int, memberSet func(int) []int) {
		for i := 0; i < count; i++ {
			groupMembers := memberSet(i)
			creator := groupMembers[i%len(groupMembers)]
			dataset.Groups = append(dataset.Groups, DatasetGroup{
				Index: groupIndex, Tier: tier,
				Title:          fmt.Sprintf("%s %s %04d", runID, tier, i+1),
				About:          fmt.Sprintf("telesrv real-RPC load dataset %s group %d", tier, i+1),
				CreatorAccount: creator, MemberAccounts: groupMembers, HistoryMessages: history,
			})
			groupIndex++
		}
	}
	appendTier("hot", cfg.HotGroups, cfg.HotMembers, cfg.HotHistory, func(i int) []int {
		return cyclicMembers(cfg.Accounts, i*max(cfg.HotMembers, 1), cfg.HotMembers)
	})
	appendTier("medium", cfg.MediumGroups, cfg.MediumMembers, cfg.MediumHistory, func(i int) []int {
		return cyclicMembers(cfg.Accounts, i*max(cfg.MediumMembers, 1), cfg.MediumMembers)
	})
	appendTier("small", cfg.SmallGroups, cfg.SmallMembers, cfg.SmallHistory, func(i int) []int {
		return cyclicMembers(cfg.Accounts, i*max(cfg.SmallMembers, 1), cfg.SmallMembers)
	})
	appendTier("heavy", cfg.HeavyGroups, cfg.HeavyAccounts, cfg.HeavyHistory, func(int) []int {
		members := make([]int, cfg.HeavyAccounts)
		for i := range members {
			members[i] = i
		}
		return members
	})
	planHash, err := dataset.planHash()
	if err != nil {
		return nil, err
	}
	dataset.PlanSHA256 = planHash
	if err := dataset.Validate(); err != nil {
		return nil, err
	}
	return dataset, nil
}

func (c DatasetConfig) validate() error {
	if c.Accounts < 2 || c.Accounts > 100000 {
		return errors.New("dataset accounts must be between 2 and 100000")
	}
	if c.PrivateFanout < 0 || c.PrivateFanout >= c.Accounts {
		return errors.New("private fanout must be non-negative and smaller than accounts")
	}
	groups := c.HotGroups + c.MediumGroups + c.SmallGroups + c.HeavyGroups
	if groups <= 0 || groups > maxDatasetGroups {
		return fmt.Errorf("dataset group count must be between 1 and %d", maxDatasetGroups)
	}
	tiers := []struct {
		name                     string
		groups, members, history int
	}{
		{"hot", c.HotGroups, c.HotMembers, c.HotHistory},
		{"medium", c.MediumGroups, c.MediumMembers, c.MediumHistory},
		{"small", c.SmallGroups, c.SmallMembers, c.SmallHistory},
		{"heavy", c.HeavyGroups, c.HeavyAccounts, c.HeavyHistory},
	}
	memberships := 0
	messages := c.Accounts * c.PrivateFanout
	for _, tier := range tiers {
		if tier.groups < 0 || tier.members < 0 || tier.members > c.Accounts || tier.history < 0 {
			return fmt.Errorf("invalid %s dataset tier", tier.name)
		}
		if tier.groups > 0 && tier.members == 0 {
			return fmt.Errorf("%s dataset tier has groups without members", tier.name)
		}
		memberships += tier.groups * tier.members
		messages += tier.groups * tier.history
	}
	if memberships > maxDatasetMembership {
		return fmt.Errorf("dataset memberships %d exceed hard limit %d", memberships, maxDatasetMembership)
	}
	if messages > maxDatasetMessages {
		return fmt.Errorf("dataset messages %d exceed hard limit %d", messages, maxDatasetMessages)
	}
	return nil
}

func (d *Dataset) Validate() error {
	if d == nil || d.Version != DatasetVersion {
		return errors.New("invalid dataset version")
	}
	if strings.TrimSpace(d.RunID) == "" || strings.TrimSpace(d.PlanSHA256) == "" {
		return errors.New("dataset is missing run id or plan hash")
	}
	if err := d.Config.validate(); err != nil {
		return err
	}
	if len(d.PrivateEdges) != d.Config.Accounts*d.Config.PrivateFanout {
		return errors.New("dataset private graph size does not match config")
	}
	seenGroups := make(map[int]struct{}, len(d.Groups))
	for _, group := range d.Groups {
		if group.Index < 0 || group.CreatorAccount < 0 || group.CreatorAccount >= d.Config.Accounts || group.HistoryMessages < 0 {
			return fmt.Errorf("invalid group %d", group.Index)
		}
		if _, exists := seenGroups[group.Index]; exists {
			return fmt.Errorf("duplicate group index %d", group.Index)
		}
		seenGroups[group.Index] = struct{}{}
		members := make(map[int]struct{}, len(group.MemberAccounts))
		for _, account := range group.MemberAccounts {
			if account < 0 || account >= d.Config.Accounts {
				return fmt.Errorf("group %d has invalid account %d", group.Index, account)
			}
			if _, exists := members[account]; exists {
				return fmt.Errorf("group %d has duplicate account %d", group.Index, account)
			}
			members[account] = struct{}{}
		}
		if _, ok := members[group.CreatorAccount]; !ok {
			return fmt.Errorf("group %d creator is not a member", group.Index)
		}
	}
	wantHash, err := d.planHash()
	if err != nil {
		return err
	}
	if wantHash != d.PlanSHA256 {
		return errors.New("dataset immutable plan hash mismatch")
	}
	return nil
}

func WriteDataset(path string, dataset *Dataset) error {
	if err := dataset.Validate(); err != nil {
		return err
	}
	data, err := json.MarshalIndent(dataset, "", "  ")
	if err != nil {
		return fmt.Errorf("encode dataset: %w", err)
	}
	return writeFileAtomic(path, append(data, '\n'), 0o600)
}

func NewDatasetSeedState(dataset *Dataset) (*DatasetSeedState, error) {
	if err := dataset.Validate(); err != nil {
		return nil, err
	}
	state := &DatasetSeedState{
		Version: DatasetVersion, PlanSHA256: dataset.PlanSHA256, UpdatedAt: time.Now().UTC(),
		PrivateSentByAccount: make([]int, dataset.Config.Accounts),
		HistorySentByAccount: make([]int, dataset.Config.Accounts),
		RichStateByAccount:   make([]bool, dataset.Config.Accounts),
		Groups:               make([]DatasetSeedGroupState, len(dataset.Groups)),
	}
	for i, group := range dataset.Groups {
		state.Groups[i].GroupIndex = group.Index
	}
	return state, nil
}

func (s *DatasetSeedState) Validate(dataset *Dataset) error {
	if s == nil || s.Version != DatasetVersion || dataset == nil || s.PlanSHA256 != dataset.PlanSHA256 {
		return errors.New("seed state does not match dataset plan")
	}
	if len(s.PrivateSentByAccount) != dataset.Config.Accounts || len(s.HistorySentByAccount) != dataset.Config.Accounts || len(s.Groups) != len(dataset.Groups) {
		return errors.New("seed state dimensions do not match dataset plan")
	}
	// A nil rich-state vector is accepted only for seed journals created before
	// the comprehensive startup workload added this phase. New journals always
	// allocate the full vector and therefore cannot silently skip it.
	if len(s.RichStateByAccount) != 0 && len(s.RichStateByAccount) != dataset.Config.Accounts {
		return errors.New("seed rich-state dimensions do not match dataset plan")
	}
	historyTasks := datasetHistoryTaskCounts(dataset)
	for account := 0; account < dataset.Config.Accounts; account++ {
		if s.PrivateSentByAccount[account] < 0 || s.PrivateSentByAccount[account] > dataset.Config.PrivateFanout {
			return fmt.Errorf("invalid private seed progress for account %d", account)
		}
		if s.HistorySentByAccount[account] < 0 || s.HistorySentByAccount[account] > historyTasks[account] {
			return fmt.Errorf("invalid history seed progress for account %d", account)
		}
	}
	for i, groupState := range s.Groups {
		group := dataset.Groups[i]
		invitees := len(group.MemberAccounts) - 1
		if groupState.GroupIndex != group.Index || groupState.InviteCursor < 0 || groupState.InviteCursor > invitees {
			return fmt.Errorf("invalid seed state for group %d", group.Index)
		}
		if (groupState.ChannelID == 0) != (groupState.AccessHash == 0) {
			return fmt.Errorf("group %d has partial channel identity", group.Index)
		}
		if groupState.ChannelID != 0 && groupState.CreatePending {
			return fmt.Errorf("group %d has both channel identity and pending create", group.Index)
		}
		if groupState.InvitePendingEnd < groupState.InviteCursor || groupState.InvitePendingEnd > invitees {
			return fmt.Errorf("group %d has invalid pending invite range", group.Index)
		}
		if groupState.ChannelID == 0 && (groupState.InviteCursor != 0 || groupState.InvitePendingEnd != 0) {
			return fmt.Errorf("group %d has invite progress without a channel identity", group.Index)
		}
	}
	return nil
}

func WriteDatasetSeedState(path string, dataset *Dataset, state *DatasetSeedState) error {
	if err := state.Validate(dataset); err != nil {
		return err
	}
	state.UpdatedAt = time.Now().UTC()
	data, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return fmt.Errorf("encode dataset seed state: %w", err)
	}
	return writeFileAtomic(path, append(data, '\n'), 0o600)
}

func LoadDatasetSeedState(path string, dataset *Dataset) (*DatasetSeedState, error) {
	data, err := os.ReadFile(path)
	if os.IsNotExist(err) {
		return NewDatasetSeedState(dataset)
	}
	if err != nil {
		return nil, fmt.Errorf("read dataset seed state: %w", err)
	}
	var state DatasetSeedState
	decoder := json.NewDecoder(strings.NewReader(string(data)))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&state); err != nil {
		return nil, fmt.Errorf("decode dataset seed state: %w", err)
	}
	if err := state.Validate(dataset); err != nil {
		return nil, err
	}
	return &state, nil
}

func datasetHistoryTaskCounts(dataset *Dataset) []int {
	counts := make([]int, dataset.Config.Accounts)
	for _, group := range dataset.Groups {
		for message := 0; message < group.HistoryMessages; message++ {
			counts[group.MemberAccounts[message%len(group.MemberAccounts)]]++
		}
	}
	return counts
}

func LoadDataset(path string) (*Dataset, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read dataset: %w", err)
	}
	var dataset Dataset
	decoder := json.NewDecoder(strings.NewReader(string(data)))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&dataset); err != nil {
		return nil, fmt.Errorf("decode dataset: %w", err)
	}
	if err := dataset.Validate(); err != nil {
		return nil, err
	}
	return &dataset, nil
}

func (d *Dataset) planHash() (string, error) {
	type immutableDataset struct {
		Version      int                  `json:"version"`
		RunID        string               `json:"run_id"`
		Config       DatasetConfig        `json:"config"`
		PrivateEdges []DatasetPrivateEdge `json:"private_edges"`
		Groups       []struct {
			Index           int    `json:"index"`
			Tier            string `json:"tier"`
			Title           string `json:"title"`
			About           string `json:"about"`
			CreatorAccount  int    `json:"creator_account"`
			MemberAccounts  []int  `json:"member_accounts"`
			HistoryMessages int    `json:"history_messages"`
		} `json:"groups"`
	}
	immutable := immutableDataset{Version: d.Version, RunID: d.RunID, Config: d.Config, PrivateEdges: d.PrivateEdges}
	for _, group := range d.Groups {
		immutable.Groups = append(immutable.Groups, struct {
			Index           int    `json:"index"`
			Tier            string `json:"tier"`
			Title           string `json:"title"`
			About           string `json:"about"`
			CreatorAccount  int    `json:"creator_account"`
			MemberAccounts  []int  `json:"member_accounts"`
			HistoryMessages int    `json:"history_messages"`
		}{group.Index, group.Tier, group.Title, group.About, group.CreatorAccount, group.MemberAccounts, group.HistoryMessages})
	}
	data, err := json.Marshal(immutable)
	if err != nil {
		return "", err
	}
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:]), nil
}

func cyclicMembers(accounts, start, count int) []int {
	if count <= 0 {
		return nil
	}
	members := make([]int, count)
	for i := range members {
		members[i] = (start + i) % accounts
	}
	sort.Ints(members)
	return members
}

func stableDatasetID(seed int64, namespace string, values ...int) int64 {
	h := sha256.New()
	var buf [8]byte
	binary.LittleEndian.PutUint64(buf[:], uint64(seed))
	_, _ = h.Write(buf[:])
	_, _ = h.Write([]byte(namespace))
	for _, value := range values {
		binary.LittleEndian.PutUint64(buf[:], uint64(value))
		_, _ = h.Write(buf[:])
	}
	sum := h.Sum(nil)
	id := int64(binary.LittleEndian.Uint64(sum[:8]) & ^(uint64(1) << 63))
	if id == 0 {
		return 1
	}
	return id
}