package smithy

import (
	"fmt"
	"strings"
	"sync/atomic"
	"unsafe"
)

// ShapeType is a type of Smithy shape.
// See https://smithy.io/2.0/spec/idl.html#defining-shapes.
type ShapeType int

// Enumerates ShapeType per the Smithy IDL.
const (
	ShapeTypeBlob ShapeType = iota
	ShapeTypeBoolean
	ShapeTypeString
	ShapeTypeTimestamp
	ShapeTypeByte
	ShapeTypeShort
	ShapeTypeInteger
	ShapeTypeLong
	ShapeTypeFloat
	ShapeTypeDocument
	ShapeTypeDouble
	ShapeTypeBigDecimal
	ShapeTypeBigInteger
	ShapeTypeEnum
	ShapeTypeIntEnum
	ShapeTypeList
	ShapeTypeSet
	ShapeTypeMap
	ShapeTypeStructure
	ShapeTypeUnion
	ShapeTypeMember
	ShapeTypeService
	ShapeTypeResource
	ShapeTypeOperation
)

// ShapeID fields of a Smithy shape ID.
type ShapeID struct {
	Namespace, Name, Member string
}

// String returns the IDL microformat for the shape ID.
func (s ShapeID) String() string {
	if s.Member == "" {
		return fmt.Sprintf("%s#%s", s.Namespace, s.Name)
	}
	return fmt.Sprintf("%s#%s$%s", s.Namespace, s.Name, s.Member)
}

func stoid(s string) ShapeID {
	ns, n, _ := strings.Cut(s, "#")
	n, m, _ := strings.Cut(n, "$")
	return ShapeID{ns, n, m}
}

// Schema encodes information about a shape from a Smithy model.
//
// Generated clients use schemas at runtime to dynamically (de)serialize
// request/responses.
type Schema struct {
	id         ShapeID
	typ        ShapeType
	members    map[string]*Schema // member name -> schema
	traits     map[ShapeID]Trait  // trait ID -> non-indexed traits only
	indexed    []Trait            // indexed trait slots, sized to max index present
	directMask uint64             // bitmask: bit i set means indexed[i] was declared directly on this schema
	targetID   ShapeID            // for member schemas, the target's shape ID

	// resolved on the fly and cached
	listMember       atomic.Pointer[Schema]
	mapKey, mapValue atomic.Pointer[Schema]

	ext [numExtensionSlots]unsafe.Pointer // lazily-computed codec extensions, accessed atomically
}

// NewSchema creates a new Schema with the given shape ID and traits.
func NewSchema(id ShapeID, typ ShapeType, numMembers int, ts ...Trait) *Schema {
	s := &Schema{
		id:      id,
		typ:     typ,
		members: make(map[string]*Schema, numMembers),
	}
	for _, t := range ts {
		s.addTrait(t, true)
	}
	return s
}

func (s *Schema) addTrait(t Trait, direct bool) {
	if it, ok := t.(IndexableTrait); ok {
		idx := it.TraitIndex()
		if idx >= len(s.indexed) {
			s.indexed = append(s.indexed, make([]Trait, idx-len(s.indexed)+1)...)
		}
		s.indexed[idx] = t
		if direct {
			s.directMask |= 1 << uint(idx)
		}
		return
	}

	if s.traits == nil {
		s.traits = map[ShapeID]Trait{}
	}
	s.traits[t.TraitID()] = t
}

// AddMember adds a member to the schema derived from the target, with
// optional trait overrides. The member schema is returned for caller
// reference.
//
// The member schema's effective trait view (accessed via [SchemaTrait])
// inherits all of the target's traits, then applies the overrides. The
// member's direct trait view (accessed via [SchemaDirectTrait]) contains
// only the overrides, i.e. the traits declared directly on the member.
func (s *Schema) AddMember(name string, target *Schema, ts ...Trait) *Schema {
	m := &Schema{
		id:         ShapeID{Member: name},
		typ:        target.typ,
		members:    target.members,
		indexed:    cloneIndexed(target.indexed),
		traits:     cloneTraits(target.traits),
		directMask: 0, // inherited traits are not direct
		targetID:   target.id,
	}

	// member-declared traits override and are direct
	for _, t := range ts {
		m.addTrait(t, true)
	}

	s.members[name] = m

	// Invalidate cached extensions, schema structure changed.
	for i := range s.ext {
		atomic.StorePointer(&s.ext[i], nil)
	}

	return m
}

func cloneIndexed(src []Trait) []Trait {
	if src == nil {
		return nil
	}
	dst := make([]Trait, len(src))
	copy(dst, src)
	return dst
}

func cloneTraits(src map[ShapeID]Trait) map[ShapeID]Trait {
	if src == nil {
		return nil
	}
	dst := make(map[ShapeID]Trait, len(src))
	for k, v := range src {
		dst[k] = v
	}
	return dst
}

// ListMember returns the "member" schema for list types.
func (s *Schema) ListMember() *Schema {
	return s.lookup(&s.listMember, "member")
}

// MapKey returns the "key" schema for map types.
func (s *Schema) MapKey() *Schema {
	return s.lookup(&s.mapKey, "key")
}

// MapValue returns the "value" schema for map types.
func (s *Schema) MapValue() *Schema {
	return s.lookup(&s.mapValue, "value")
}

func (s *Schema) lookup(cached *atomic.Pointer[Schema], name string) *Schema {
	if v := cached.Load(); v != nil {
		return v
	}

	m, ok := s.members[name]
	if !ok {
		return nil
	}

	cached.Store(m)
	return m
}

// MemberName returns the member component of the schema's shape ID.
func (s *Schema) MemberName() string {
	return s.id.Member
}

// ID returns the shape ID of the schema.
func (s *Schema) ID() ShapeID {
	return s.id
}

// TargetID returns the shape ID of the member's target shape.
func (s *Schema) TargetID() ShapeID {
	return s.targetID
}

// Type returns the shape type of the schema.
func (s *Schema) Type() ShapeType {
	return s.typ
}

// Member returns the member schema for the given name, or nil.
func (s *Schema) Member(name string) *Schema {
	return s.members[name]
}

// Members returns the schema's members as a map of name to schema.
func (s *Schema) Members() map[string]*Schema {
	return s.members
}

// OperationSchema describes an operation, which is essentially its own schema
// with additional pointers to its input and output.
type OperationSchema struct {
	*Schema
	Input, Output *Schema

	inputStream, outputStream bool
}

// NewOperationSchema returns an OperationSchema for (input, output).
func NewOperationSchema(op, input, output *Schema) *OperationSchema {
	return &OperationSchema{
		Schema:       op,
		Input:        input,
		Output:       output,
		inputStream:  isEventStream(input),
		outputStream: isEventStream(output),
	}
}

// IsInputEventStream reports whether this is an input event stream.
func (s *OperationSchema) IsInputEventStream() bool {
	return s.inputStream
}

// IsOutputEventStream reports whether this is an output event stream.
func (s *OperationSchema) IsOutputEventStream() bool {
	return s.outputStream
}

// ServiceSchema describes a service shape.
type ServiceSchema struct {
	*Schema
	Version string
}

// NewServiceSchema returns a ServiceSchema for the given service shape.
func NewServiceSchema(schema *Schema, version string) *ServiceSchema {
	return &ServiceSchema{Schema: schema, Version: version}
}

// SchemaTrait returns the target trait on the schema if it exists.
//
// For member schemas this returns the effective trait, which is the trait
// declared directly on the member if present, else the trait inherited from
// the target shape.
func SchemaTrait[T Trait](s *Schema) (T, bool) {
	return schemaTrait[T](s, false)
}

// SchemaDirectTrait returns the target trait on the schema if it was
// declared directly on the schema.
//
// For member schemas this returns the trait only if it was declared on the
// member itself, ignoring any trait inherited from the target shape. For
// non-member schemas this is equivalent to [SchemaTrait].
func SchemaDirectTrait[T Trait](s *Schema) (T, bool) {
	return schemaTrait[T](s, true)
}

func schemaTrait[T Trait](s *Schema, directOnly bool) (T, bool) {
	var zero T

	if s == nil {
		return zero, false
	}

	if it, ok := Trait(zero).(IndexableTrait); ok {
		idx := it.TraitIndex()
		if idx >= len(s.indexed) {
			return zero, false
		}
		if directOnly && s.directMask&(1<<uint(idx)) == 0 {
			return zero, false
		}
		tt, ok := s.indexed[idx].(T)
		return tt, ok
	}

	opaque, ok := s.traits[zero.TraitID()]
	if !ok {
		return zero, false
	}

	tt, ok := opaque.(T)
	return tt, ok
}

// indexStreaming is the indexed trait slot for @streaming, mirrored from
// traits.indexStreaming. We can't import the traits package from here due to a
// circular dependency.
const indexStreaming = 17

func isEventStream(s *Schema) bool {
	if s == nil {
		return false
	}
	for _, m := range s.members {
		if m.typ != ShapeTypeUnion {
			continue
		}
		if len(m.indexed) > indexStreaming && m.indexed[indexStreaming] != nil {
			return true
		}
	}
	return false
}