package main

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

type ctrlPacket struct {
	ID     string                     `json:"id"`
	Topic  string                     `json:"topic"`
	Code   int                        `json:"code"`
	Text   string                     `json:"text"`
	Params map[string]json.RawMessage `json:"params"`
}

type dataPacket struct {
	Topic     string          `json:"topic"`
	From      string          `json:"from"`
	Timestamp time.Time       `json:"ts"`
	Seq       int64           `json:"seq"`
	Head      map[string]any  `json:"head"`
	Content   json.RawMessage `json:"content"`
}

type ctrlError struct {
	Code int
	Text string
}

func (e *ctrlError) Error() string {
	return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
}

type WSClient struct {
	account        AccountRecord
	conn           *websocket.Conn
	commandTimeout time.Duration
	onData         func(AccountRecord, dataPacket)
	onDisconnect   func(AccountRecord, error)

	sendMu sync.Mutex
	mu     sync.Mutex
	seq    uint64
	waits  map[string]chan ctrlPacket
	done   chan struct{}
	closed atomic.Bool
}

func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
	onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error)) (*WSClient, error) {
	timeout, err := configDuration(cfg.CommandTimeout)
	if err != nil {
		return nil, err
	}
	endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
	if err != nil {
		return nil, err
	}
	dialer := *websocket.DefaultDialer
	if cfg.InsecureTLS {
		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
	}
	conn, response, err := dialer.DialContext(ctx, endpoint, http.Header{})
	if err != nil {
		if response != nil {
			return nil, fmt.Errorf("websocket dial status=%d: %w", response.StatusCode, err)
		}
		return nil, err
	}
	client := &WSClient{
		account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
		waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
	}
	go client.readLoop()
	if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
		"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure/1.0",
		"platf": "loadtest", "tenant": cfg.Tenant,
	}}); err != nil {
		client.Close()
		return nil, fmt.Errorf("hi %s: %w", account.Username, err)
	}
	login := map[string]any{"id": client.nextID("login")}
	if account.Token != "" {
		login["scheme"] = "token"
		login["secret"] = account.Token
	} else {
		login["scheme"] = "basic"
		login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
		login["login_type"] = 0
		login["tenant_code"] = cfg.Tenant
	}
	ctrl, err := client.command(ctx, "login", map[string]any{"login": login})
	if err != nil && account.Token != "" && password != "" {
		login["id"] = client.nextID("login-basic")
		login["scheme"] = "basic"
		login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
		login["login_type"] = 0
		login["tenant_code"] = cfg.Tenant
		ctrl, err = client.command(ctx, "login", map[string]any{"login": login})
	}
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("login %s: %w", account.Username, err)
	}
	client.applyLogin(ctrl)
	pingInterval, err := configDuration(cfg.WSPingInterval)
	if err != nil {
		client.Close()
		return nil, err
	}
	client.startKeepalive(pingInterval)
	return client, nil
}

func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient, password, verificationCode string) (*WSClient, bool, error) {
	timeout, err := configDuration(cfg.CommandTimeout)
	if err != nil {
		return nil, false, err
	}
	endpoint, err := websocketEndpoint(account.Endpoint, api.apiKey)
	if err != nil {
		return nil, false, err
	}
	dialer := *websocket.DefaultDialer
	if cfg.InsecureTLS {
		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
	}
	conn, _, err := dialer.DialContext(ctx, endpoint, http.Header{})
	if err != nil {
		return nil, false, err
	}
	client := &WSClient{account: account, conn: conn, commandTimeout: timeout, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
	go client.readLoop()
	if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
		"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure-prepare/1.0", "tenant": cfg.Tenant,
	}}); err != nil {
		client.Close()
		return nil, false, err
	}
	secret := base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
	ctrl, loginErr := client.command(ctx, "login", map[string]any{"login": map[string]any{
		"id": client.nextID("login"), "scheme": "basic", "secret": secret,
		"login_type": 0, "tenant_code": cfg.Tenant,
	}})
	created := false
	if loginErr != nil {
		var ce *ctrlError
		if !errors.As(loginErr, &ce) || ce.Code != http.StatusUnauthorized {
			client.Close()
			return nil, false, loginErr
		}
		credential, credentialErr := api.createRegistrationCredential(ctx, account, verificationCode)
		if credentialErr != nil {
			client.Close()
			return nil, false, credentialErr
		}
		acc := map[string]any{
			"id": client.nextID("acc"), "user": "new", "scheme": "basic", "username": account.Username,
			"registration_type": 0, "tenant_code": cfg.Tenant, "secret": secret, "login": true,
			"desc": map[string]any{"public": map[string]any{"fn": strings.ToUpper(account.Username)}},
		}
		if credential != nil {
			acc["phone"] = account.Phone
			acc["cred"] = []map[string]any{{"meth": "tel", "val": account.Phone,
				"resp": credential.VerificationToken, "params": map[string]any{"challenge_id": credential.ChallengeID}}}
		}
		ctrl, err = client.command(ctx, "acc", map[string]any{"acc": acc})
		if err != nil {
			client.Close()
			return nil, false, err
		}
		created = true
	}
	client.applyLogin(ctrl)
	return client, created, nil
}

func websocketEndpoint(raw, apiKey string) (string, error) {
	u, err := url.Parse(strings.TrimSpace(raw))
	if err != nil {
		return "", err
	}
	if u.Scheme != "ws" && u.Scheme != "wss" {
		return "", fmt.Errorf("unsupported websocket scheme %q", u.Scheme)
	}
	q := u.Query()
	if q.Get("apikey") == "" {
		q.Set("apikey", apiKey)
	}
	u.RawQuery = q.Encode()
	return u.String(), nil
}

func (c *WSClient) applyLogin(ctrl ctrlPacket) {
	if raw := ctrl.Params["user"]; len(raw) > 0 {
		_ = json.Unmarshal(raw, &c.account.UserID)
	}
	if raw := ctrl.Params["token"]; len(raw) > 0 {
		_ = json.Unmarshal(raw, &c.account.Token)
	}
	if raw := ctrl.Params["expires"]; len(raw) > 0 {
		var value string
		if json.Unmarshal(raw, &value) == nil {
			c.account.TokenExpires, _ = time.Parse(time.RFC3339Nano, value)
		}
	}
}

func (c *WSClient) Account() AccountRecord { return c.account }

func (c *WSClient) Subscribe(ctx context.Context, topic string) error {
	if topic == "" {
		return errors.New("empty topic")
	}
	id := c.nextID("sub")
	_, err := c.command(ctx, "sub", map[string]any{"sub": map[string]any{
		"id": id, "topic": topic, "get": map[string]any{"what": "desc"},
	}})
	return err
}

func (c *WSClient) GetData(ctx context.Context, topic string, since, before int64, limit int) error {
	if topic == "" || since <= 0 || before <= since || limit <= 0 {
		return errors.New("invalid get data range")
	}
	id := c.nextID("get-data")
	_, err := c.command(ctx, "get", map[string]any{"get": map[string]any{
		"id": id, "topic": topic, "what": "data",
		"data": map[string]any{"since": since, "before": before, "limit": limit},
	}})
	return err
}

func (c *WSClient) Publish(ctx context.Context, topic, clientMid, runID, scenario string, messageIndex int) (AcceptedMessage, error) {
	id := c.nextID("pub")
	sentAt := time.Now().UTC()
	ctrl, err := c.command(ctx, "pub", map[string]any{"pub": map[string]any{
		"id": id, "topic": topic, "noecho": false,
		"head": map[string]any{"mime": "application/json", "x-im-schema": 1, "x-im-type": "text", "x-im-client-mid": clientMid},
		"content": map[string]any{"text": "loadtest", "run_id": runID, "scenario": scenario,
			"sender_machine": c.account.Machine, "sender_index": c.account.Index,
			"message_index": messageIndex, "sent_at_unix_nano": sentAt.UnixNano()},
	}})
	if err != nil {
		return AcceptedMessage{}, err
	}
	seq, err := ctrlInt64(ctrl.Params["seq"])
	if err != nil || seq <= 0 {
		return AcceptedMessage{}, fmt.Errorf("publish ACK has invalid seq: %w", err)
	}
	kind := "p2p"
	if strings.HasPrefix(topic, "grp") {
		kind = "group"
	}
	return AcceptedMessage{RunID: runID, Scenario: scenario, Kind: kind, ClientMid: clientMid,
		Topic: topic, Seq: seq, SenderMachine: c.account.Machine, SenderIndex: c.account.Index,
		SentAt: sentAt, AckAt: time.Now().UTC()}, nil
}

func ctrlInt64(raw json.RawMessage) (int64, error) {
	if len(raw) == 0 {
		return 0, errors.New("missing value")
	}
	var number json.Number
	if err := json.Unmarshal(raw, &number); err == nil {
		return strconv.ParseInt(number.String(), 10, 64)
	}
	var text string
	if err := json.Unmarshal(raw, &text); err != nil {
		return 0, err
	}
	return strconv.ParseInt(text, 10, 64)
}

func (c *WSClient) command(ctx context.Context, _ string, message map[string]any) (ctrlPacket, error) {
	id := nestedMessageID(message)
	if id == "" {
		return ctrlPacket{}, errors.New("command id is required")
	}
	wait := make(chan ctrlPacket, 1)
	c.mu.Lock()
	c.waits[id] = wait
	c.mu.Unlock()
	c.sendMu.Lock()
	err := c.conn.WriteJSON(message)
	c.sendMu.Unlock()
	if err != nil {
		c.removeWait(id)
		return ctrlPacket{}, err
	}
	timer := time.NewTimer(c.commandTimeout)
	defer timer.Stop()
	select {
	case ctrl := <-wait:
		if ctrl.Code < 200 || ctrl.Code >= 300 {
			return ctrl, &ctrlError{Code: ctrl.Code, Text: ctrl.Text}
		}
		return ctrl, nil
	case <-timer.C:
		c.removeWait(id)
		return ctrlPacket{}, fmt.Errorf("ctrl timeout id=%s", id)
	case <-ctx.Done():
		c.removeWait(id)
		return ctrlPacket{}, ctx.Err()
	case <-c.done:
		c.removeWait(id)
		return ctrlPacket{}, errors.New("websocket closed")
	}
}

func nestedMessageID(message map[string]any) string {
	for _, value := range message {
		if body, ok := value.(map[string]any); ok {
			if id, ok := body["id"].(string); ok {
				return id
			}
		}
	}
	return ""
}

func (c *WSClient) nextID(prefix string) string {
	n := atomic.AddUint64(&c.seq, 1)
	return fmt.Sprintf("%s-%s-%04d-%d", prefix, c.account.Machine, c.account.Index, n)
}

func (c *WSClient) removeWait(id string) {
	c.mu.Lock()
	delete(c.waits, id)
	c.mu.Unlock()
}

func (c *WSClient) readLoop() {
	defer func() {
		_ = c.conn.Close()
		close(c.done)
	}()
	for {
		var envelope map[string]json.RawMessage
		if err := c.conn.ReadJSON(&envelope); err != nil {
			if !c.closed.Load() && c.onDisconnect != nil {
				c.onDisconnect(c.account, err)
			}
			return
		}
		if raw := envelope["ctrl"]; len(raw) > 0 {
			var ctrl ctrlPacket
			if json.Unmarshal(raw, &ctrl) == nil && ctrl.ID != "" {
				c.mu.Lock()
				wait := c.waits[ctrl.ID]
				delete(c.waits, ctrl.ID)
				c.mu.Unlock()
				if wait != nil {
					wait <- ctrl
				}
			}
			continue
		}
		if raw := envelope["data"]; len(raw) > 0 && c.onData != nil {
			var data dataPacket
			if json.Unmarshal(raw, &data) == nil {
				c.onData(c.account, data)
			}
		}
	}
}

func (c *WSClient) startKeepalive(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				deadline := time.Now().Add(interval / 2)
				if err := c.conn.WriteControl(websocket.PingMessage, nil, deadline); err != nil {
					_ = c.conn.Close()
					return
				}
			case <-c.done:
				return
			}
		}
	}()
}

func (c *WSClient) alive() bool {
	if c == nil || c.closed.Load() {
		return false
	}
	select {
	case <-c.done:
		return false
	default:
		return true
	}
}

func (c *WSClient) Close() {
	if c == nil || !c.closed.CompareAndSwap(false, true) {
		return
	}
	_ = c.conn.Close()
}