package vipcardexperimentmod

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/bsontype"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

// OrderedProductIDs exposes productIds as an order-number-to-product-ID map
// while keeping a slice internally so every backend consumer has deterministic
// numeric order. BSON is stored as an embedded document with keys 1..n.
// Legacy JSON/BSON arrays remain readable for existing experiment records.
type OrderedProductIDs []primitive.ObjectID

func (ids OrderedProductIDs) MarshalJSON() ([]byte, error) {
	var buffer bytes.Buffer
	buffer.WriteByte('{')
	for index, productID := range ids {
		if index > 0 {
			buffer.WriteByte(',')
		}
		key, _ := json.Marshal(strconv.Itoa(index + 1))
		value, err := json.Marshal(productID)
		if err != nil {
			return nil, err
		}
		buffer.Write(key)
		buffer.WriteByte(':')
		buffer.Write(value)
	}
	buffer.WriteByte('}')
	return buffer.Bytes(), nil
}

func (ids *OrderedProductIDs) UnmarshalJSON(data []byte) error {
	data = bytes.TrimSpace(data)
	if len(data) == 0 {
		return fmt.Errorf("productIds is required")
	}
	if bytes.Equal(data, []byte("null")) {
		*ids = nil
		return nil
	}
	switch data[0] {
	case '[':
		var legacy []primitive.ObjectID
		if err := json.Unmarshal(data, &legacy); err != nil {
			return fmt.Errorf("invalid legacy productIds: %w", err)
		}
		*ids = legacy
		return nil
	case '{':
		var values map[string]json.RawMessage
		if err := json.Unmarshal(data, &values); err != nil {
			return fmt.Errorf("invalid productIds map: %w", err)
		}
		ordered := make(OrderedProductIDs, len(values))
		occupied := make([]bool, len(values))
		seenProducts := make(map[primitive.ObjectID]struct{}, len(values))
		for key, rawProductID := range values {
			index, err := parseProductOrder(key, len(values))
			if err != nil {
				return err
			}
			var productID primitive.ObjectID
			if err = json.Unmarshal(rawProductID, &productID); err != nil || productID.IsZero() {
				return fmt.Errorf("productIds.%s must be a valid product ID", key)
			}
			if occupied[index] {
				return fmt.Errorf("productIds order %s is duplicated", key)
			}
			if _, exists := seenProducts[productID]; exists {
				return fmt.Errorf("productIds contains duplicate product ID: %s", productID.Hex())
			}
			occupied[index] = true
			seenProducts[productID] = struct{}{}
			ordered[index] = productID
		}
		*ids = ordered
		return nil
	default:
		return fmt.Errorf("productIds must be an order map")
	}
}

func (ids OrderedProductIDs) MarshalBSONValue() (bsontype.Type, []byte, error) {
	document := make(bson.D, 0, len(ids))
	for index, productID := range ids {
		document = append(document, bson.E{
			Key:   strconv.Itoa(index + 1),
			Value: productID,
		})
	}
	return bson.MarshalValue(document)
}

func (ids *OrderedProductIDs) UnmarshalBSONValue(valueType bsontype.Type, data []byte) error {
	rawValue := bson.RawValue{Type: valueType, Value: data}
	switch valueType {
	case bsontype.Array:
		var legacy []primitive.ObjectID
		if err := rawValue.Unmarshal(&legacy); err != nil {
			return fmt.Errorf("invalid legacy productIds: %w", err)
		}
		*ids = legacy
		return nil
	case bsontype.EmbeddedDocument:
		var document bson.Raw
		if err := rawValue.Unmarshal(&document); err != nil {
			return fmt.Errorf("invalid productIds document: %w", err)
		}
		elements, err := document.Elements()
		if err != nil {
			return fmt.Errorf("invalid productIds document: %w", err)
		}
		ordered := make(OrderedProductIDs, len(elements))
		occupied := make([]bool, len(elements))
		seenProducts := make(map[primitive.ObjectID]struct{}, len(elements))
		for _, element := range elements {
			key := element.Key()
			index, parseErr := parseProductOrder(key, len(elements))
			if parseErr != nil {
				return parseErr
			}
			productID, ok := element.Value().ObjectIDOK()
			if !ok || productID.IsZero() {
				return fmt.Errorf("productIds.%s must be a valid product ID", key)
			}
			if occupied[index] {
				return fmt.Errorf("productIds order %s is duplicated", key)
			}
			if _, exists := seenProducts[productID]; exists {
				return fmt.Errorf("productIds contains duplicate product ID: %s", productID.Hex())
			}
			occupied[index] = true
			seenProducts[productID] = struct{}{}
			ordered[index] = productID
		}
		*ids = ordered
		return nil
	case bsontype.Null, bsontype.Undefined:
		*ids = nil
		return nil
	default:
		return fmt.Errorf("productIds must be an ordered document")
	}
}

func parseProductOrder(key string, size int) (int, error) {
	order, err := strconv.Atoi(key)
	if err != nil || order <= 0 || strconv.Itoa(order) != key {
		return 0, fmt.Errorf("productIds order must be a positive integer: %s", key)
	}
	if order > size {
		return 0, fmt.Errorf("productIds order must be continuous from 1")
	}
	return order - 1, nil
}