package main

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

const manifestVersion = 1

type Config struct {
	Tenant           string              `json:"tenant"`
	HTTPBase         string              `json:"http_base"`
	APIKeyEnv        string              `json:"api_key_env"`
	PasswordEnv      string              `json:"password_env"`
	UsersPerMachine  int                 `json:"users_per_machine"`
	Endpoints        map[string][]string `json:"endpoints"`
	InsecureTLS      bool                `json:"insecure_tls"`
	PrepareWorkers   int                 `json:"prepare_workers"`
	ConnectRamp      string              `json:"connect_ramp"`
	CommandTimeout   string              `json:"command_timeout"`
	P2PMessages      int                 `json:"p2p_messages_per_user"`
	P2PInterval      string              `json:"p2p_interval"`
	DrainP2P         string              `json:"drain_p2p"`
	DrainGroup       string              `json:"drain_group"`
	PhonePrefix      string              `json:"registration_phone_prefix"`
	VerifyCodeEnv    string              `json:"verification_code_env"`
	HistoryWorkers   int                 `json:"history_workers"`
	HistoryTimeout   string              `json:"history_timeout"`
	PreflightWorkers int                 `json:"preflight_workers"`
	SkipP2PPreflight bool                `json:"skip_p2p_preflight"`
}

func loadConfig(path string) (*Config, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var cfg Config
	if err = json.Unmarshal(raw, &cfg); err != nil {
		return nil, fmt.Errorf("decode config: %w", err)
	}
	if cfg.Tenant == "" {
		cfg.Tenant = "006"
	}
	if cfg.UsersPerMachine <= 0 {
		cfg.UsersPerMachine = 1000
	}
	if cfg.PrepareWorkers <= 0 {
		cfg.PrepareWorkers = 20
	}
	if cfg.ConnectRamp == "" {
		cfg.ConnectRamp = "20s"
	}
	if cfg.CommandTimeout == "" {
		cfg.CommandTimeout = "20s"
	}
	if cfg.P2PMessages <= 0 {
		cfg.P2PMessages = 10
	}
	if cfg.P2PInterval == "" {
		cfg.P2PInterval = "1s"
	}
	if cfg.DrainP2P == "" {
		cfg.DrainP2P = "30s"
	}
	if cfg.DrainGroup == "" {
		cfg.DrainGroup = "60s"
	}
	if cfg.HistoryWorkers <= 0 {
		cfg.HistoryWorkers = 48
	}
	if cfg.HistoryWorkers > 128 {
		return nil, errors.New("history_workers must not exceed 128")
	}
	if cfg.HistoryTimeout == "" {
		cfg.HistoryTimeout = "3m"
	}
	if _, err = configDuration(cfg.HistoryTimeout); err != nil {
		return nil, fmt.Errorf("history_timeout: %w", err)
	}
	if cfg.PreflightWorkers <= 0 {
		cfg.PreflightWorkers = 32
	}
	if cfg.PreflightWorkers > 128 {
		return nil, errors.New("preflight_workers must not exceed 128")
	}
	if strings.TrimSpace(cfg.HTTPBase) == "" {
		return nil, errors.New("http_base is required")
	}
	for _, machine := range []string{"a", "b"} {
		if len(cfg.Endpoints[machine]) == 0 {
			return nil, fmt.Errorf("at least one endpoint is required for machine %s", machine)
		}
	}
	if cfg.PhonePrefix != "" {
		if len(cfg.PhonePrefix) != 7 || strings.Trim(cfg.PhonePrefix, "0123456789") != "" {
			return nil, errors.New("registration_phone_prefix must contain exactly 7 digits")
		}
		if cfg.UsersPerMachine > 4999 {
			return nil, errors.New("registration phone suffix supports at most 4999 users per machine")
		}
	}
	return &cfg, nil
}

func (c *Config) apiKey() (string, error) {
	name := c.APIKeyEnv
	if name == "" {
		name = "IM_API_KEY"
	}
	value := strings.TrimSpace(os.Getenv(name))
	if value == "" {
		return "", fmt.Errorf("environment variable %s is required", name)
	}
	return value, nil
}

func (c *Config) password() (string, error) {
	name := c.PasswordEnv
	if name == "" {
		name = "IM_LOAD_PASSWORD"
	}
	value := os.Getenv(name)
	if value == "" {
		return "", fmt.Errorf("environment variable %s is required", name)
	}
	return value, nil
}

func (c *Config) verificationCode() (string, error) {
	name := c.VerifyCodeEnv
	if name == "" {
		name = "IM_VERIFICATION_CODE"
	}
	value := strings.TrimSpace(os.Getenv(name))
	if c.PhonePrefix != "" && value == "" {
		return "", fmt.Errorf("environment variable %s is required", name)
	}
	return value, nil
}

func configDuration(value string) (time.Duration, error) {
	d, err := time.ParseDuration(value)
	if err != nil || d <= 0 {
		return 0, fmt.Errorf("invalid positive duration %q", value)
	}
	return d, nil
}

type Manifest struct {
	Version    int             `json:"version"`
	Tenant     string          `json:"tenant"`
	GroupTopic string          `json:"group_topic,omitempty"`
	CreatedAt  time.Time       `json:"created_at"`
	Accounts   []AccountRecord `json:"accounts"`
}

type AccountRecord struct {
	Machine      string    `json:"machine"`
	Index        int       `json:"index"`
	Username     string    `json:"username"`
	Phone        string    `json:"phone,omitempty"`
	UserID       string    `json:"user_id"`
	Token        string    `json:"token"`
	TokenExpires time.Time `json:"token_expires,omitempty"`
	P2PTopic     string    `json:"p2p_topic,omitempty"`
	Endpoint     string    `json:"endpoint,omitempty"`
}

func newManifest(cfg *Config) *Manifest {
	m := &Manifest{Version: manifestVersion, Tenant: cfg.Tenant, CreatedAt: time.Now().UTC()}
	for _, machine := range []string{"a", "b"} {
		for i := 1; i <= cfg.UsersPerMachine; i++ {
			phone, _ := registrationPhone(cfg, machine, i)
			m.Accounts = append(m.Accounts, AccountRecord{
				Machine:  machine,
				Index:    i,
				Username: fmt.Sprintf("lt%s%04d", machine, i),
				Phone:    phone,
				Endpoint: cfg.Endpoints[machine][(i-1)%len(cfg.Endpoints[machine])],
			})
		}
	}
	return m
}

func applyConfigEndpoints(cfg *Config, manifest *Manifest) error {
	if manifest.Tenant != cfg.Tenant {
		return fmt.Errorf("manifest tenant=%s config tenant=%s", manifest.Tenant, cfg.Tenant)
	}
	if len(manifest.Accounts) != cfg.UsersPerMachine*2 {
		return fmt.Errorf("manifest accounts=%d expected=%d", len(manifest.Accounts), cfg.UsersPerMachine*2)
	}
	for i := range manifest.Accounts {
		account := &manifest.Accounts[i]
		endpoints := cfg.Endpoints[account.Machine]
		if account.Index <= 0 || len(endpoints) == 0 {
			return fmt.Errorf("invalid endpoint assignment for %s", accountKey(account.Machine, account.Index))
		}
		account.Endpoint = endpoints[(account.Index-1)%len(endpoints)]
		phone, err := registrationPhone(cfg, account.Machine, account.Index)
		if err != nil {
			return err
		}
		account.Phone = phone
	}
	return nil
}

func reconcileManifest(cfg *Config, existing *Manifest) (*Manifest, error) {
	if existing.Tenant != cfg.Tenant {
		return nil, fmt.Errorf("manifest tenant=%s config tenant=%s", existing.Tenant, cfg.Tenant)
	}
	desired := newManifest(cfg)
	if len(existing.Accounts) > len(desired.Accounts) {
		return nil, fmt.Errorf("refusing to shrink manifest accounts from %d to %d", len(existing.Accounts), len(desired.Accounts))
	}
	desired.GroupTopic = existing.GroupTopic
	desired.CreatedAt = existing.CreatedAt
	old := manifestAccountMap(existing)
	for i := range desired.Accounts {
		account := &desired.Accounts[i]
		if previous := old[accountKey(account.Machine, account.Index)]; previous != nil {
			account.UserID = previous.UserID
			account.Token = previous.Token
			account.TokenExpires = previous.TokenExpires
			account.P2PTopic = previous.P2PTopic
		}
	}
	return desired, nil
}

func registrationPhone(cfg *Config, machine string, index int) (string, error) {
	if cfg.PhonePrefix == "" {
		return "", nil
	}
	ordinal := index
	if machine == "b" {
		ordinal += 5000
	} else if machine != "a" {
		return "", fmt.Errorf("invalid machine %q", machine)
	}
	if index <= 0 || ordinal > 9999 {
		return "", fmt.Errorf("invalid phone ordinal %d", ordinal)
	}
	return fmt.Sprintf("%s%04d", cfg.PhonePrefix, ordinal), nil
}

func loadManifest(path string) (*Manifest, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var manifest Manifest
	if err = json.Unmarshal(raw, &manifest); err != nil {
		return nil, err
	}
	if manifest.Version != manifestVersion {
		return nil, fmt.Errorf("unsupported manifest version %d", manifest.Version)
	}
	return &manifest, manifest.validate()
}

func (m *Manifest) validate() error {
	seen := make(map[string]bool, len(m.Accounts))
	for _, account := range m.Accounts {
		key := accountKey(account.Machine, account.Index)
		if account.Machine != "a" && account.Machine != "b" {
			return fmt.Errorf("invalid machine %q", account.Machine)
		}
		if account.Index <= 0 || account.Username == "" || seen[key] {
			return fmt.Errorf("invalid or duplicate account %s", key)
		}
		seen[key] = true
	}
	return nil
}

func saveManifest(path string, manifest *Manifest) error {
	if err := manifest.validate(); err != nil {
		return err
	}
	raw, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err = os.WriteFile(tmp, append(raw, '\n'), 0o600); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

func accountKey(machine string, index int) string {
	return strings.ToLower(machine) + fmt.Sprintf("-%04d", index)
}

func accountsForMachine(manifest *Manifest, machine string) []AccountRecord {
	var out []AccountRecord
	for _, account := range manifest.Accounts {
		if account.Machine == machine {
			out = append(out, account)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Index < out[j].Index })
	return out
}

func manifestAccountMap(manifest *Manifest) map[string]*AccountRecord {
	out := make(map[string]*AccountRecord, len(manifest.Accounts))
	for i := range manifest.Accounts {
		account := &manifest.Accounts[i]
		out[accountKey(account.Machine, account.Index)] = account
	}
	return out
}

func topicHash(topic string) string {
	sum := sha256.Sum256([]byte(topic))
	return hex.EncodeToString(sum[:6])
}

type AcceptedMessage struct {
	RunID         string    `json:"run_id"`
	Scenario      string    `json:"scenario"`
	Kind          string    `json:"kind"`
	ClientMid     string    `json:"client_mid"`
	Topic         string    `json:"topic"`
	Seq           int64     `json:"seq"`
	SenderMachine string    `json:"sender_machine"`
	SenderIndex   int       `json:"sender_index"`
	PeerMachine   string    `json:"peer_machine,omitempty"`
	PeerIndex     int       `json:"peer_index,omitempty"`
	SentAt        time.Time `json:"sent_at"`
	AckAt         time.Time `json:"ack_at"`
}

type AgentSummary struct {
	RunID               string              `json:"run_id"`
	Scenario            string              `json:"scenario"`
	Machine             string              `json:"machine"`
	StartedAt           time.Time           `json:"started_at"`
	LoadEndedAt         time.Time           `json:"load_ended_at"`
	FinishedAt          time.Time           `json:"finished_at"`
	ExpectedConnections int                 `json:"expected_connections"`
	ReadyConnections    int                 `json:"ready_connections"`
	PublishAttempted    int64               `json:"publish_attempted"`
	PublishAccepted     int64               `json:"publish_accepted"`
	PublishRejected     int64               `json:"publish_rejected"`
	PublishErrors       map[string]int64    `json:"publish_errors,omitempty"`
	ExpectedPeer        int64               `json:"expected_peer"`
	ReceivedPeer        int64               `json:"received_peer"`
	ExpectedEcho        int64               `json:"expected_echo"`
	ReceivedEcho        int64               `json:"received_echo"`
	ExpectedGroup       int64               `json:"expected_group"`
	ReceivedGroup       int64               `json:"received_group"`
	RealtimeLost        int64               `json:"realtime_lost"`
	RealtimeLostStored  int64               `json:"realtime_lost_but_persisted"`
	Duplicates          int64               `json:"duplicates"`
	OutOfOrder          int64               `json:"out_of_order"`
	OrderingViolations  []OrderingViolation `json:"ordering_violations,omitempty"`
	Disconnects         int64               `json:"disconnects"`
	UnexpectedMessages  int64               `json:"unexpected_messages"`
	HistoryChecked      int64               `json:"history_checked"`
	HistoryFound        int64               `json:"history_found"`
	PersistenceLost     int64               `json:"persistence_lost"`
	AckLatency          HistogramSnapshot   `json:"ack_latency"`
	DeliveryLatency     HistogramSnapshot   `json:"delivery_latency"`
	Errors              []string            `json:"errors,omitempty"`
	Missing             []MissingDelivery   `json:"missing,omitempty"`
	PersistenceMissing  []string            `json:"persistence_missing,omitempty"`
	AcceptedLocal       []AcceptedMessage   `json:"accepted_local,omitempty"`
}

type OrderingViolation struct {
	ReceiverMachine string `json:"receiver_machine"`
	ReceiverIndex   int    `json:"receiver_index"`
	TopicHash       string `json:"topic_hash"`
	PreviousSeq     int64  `json:"previous_seq"`
	CurrentSeq      int64  `json:"current_seq"`
	ClientMid       string `json:"client_mid"`
}

type MissingDelivery struct {
	ClientMid     string `json:"client_mid"`
	Kind          string `json:"kind"`
	TopicHash     string `json:"topic_hash"`
	Seq           int64  `json:"seq"`
	ReceiverIndex int    `json:"receiver_index"`
}