package mtproto

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"io"
	"sync"

	"github.com/xtls/xray-core/v1/common"
)

const (
	HeaderSize = 64
)

type SessionContext struct {
	ConnectionType [4]byte
	DataCenterID   uint16
}

func DefaultSessionContext() SessionContext {
	return SessionContext{
		ConnectionType: [4]byte{0xef, 0xef, 0xef, 0xef},
		DataCenterID:   0,
	}
}

type contextKey int32

const (
	sessionContextKey contextKey = iota
)

func ContextWithSessionContext(ctx context.Context, c SessionContext) context.Context {
	return context.WithValue(ctx, sessionContextKey, c)
}

func SessionContextFromContext(ctx context.Context) SessionContext {
	if c := ctx.Value(sessionContextKey); c != nil {
		return c.(SessionContext)
	}
	return DefaultSessionContext()
}

type Authentication struct {
	Header        [HeaderSize]byte
	DecodingKey   [32]byte
	EncodingKey   [32]byte
	DecodingNonce [16]byte
	EncodingNonce [16]byte
}

func (a *Authentication) DataCenterID() uint16 {
	x := ((int16(a.Header[61]) << 8) | int16(a.Header[60]))
	if x < 0 {
		x = -x
	}
	return uint16(x) - 1
}

func (a *Authentication) ConnectionType() [4]byte {
	var x [4]byte
	copy(x[:], a.Header[56:60])
	return x
}

func (a *Authentication) ApplySecret(b []byte) {
	a.DecodingKey = sha256.Sum256(append(a.DecodingKey[:], b...))
	a.EncodingKey = sha256.Sum256(append(a.EncodingKey[:], b...))
}

func generateRandomBytes(random []byte, connType [4]byte) {
	for {
		common.Must2(rand.Read(random))

		if random[0] == 0xef {
			continue
		}

		val := (uint32(random[3]) << 24) | (uint32(random[2]) << 16) | (uint32(random[1]) << 8) | uint32(random[0])
		if val == 0x44414548 || val == 0x54534f50 || val == 0x20544547 || val == 0x4954504f || val == 0xeeeeeeee {
			continue
		}

		if (uint32(random[7])<<24)|(uint32(random[6])<<16)|(uint32(random[5])<<8)|uint32(random[4]) == 0x00000000 {
			continue
		}

		copy(random[56:60], connType[:])

		return
	}
}

func NewAuthentication(sc SessionContext) *Authentication {
	auth := getAuthenticationObject()
	random := auth.Header[:]
	generateRandomBytes(random, sc.ConnectionType)
	copy(auth.EncodingKey[:], random[8:])
	copy(auth.EncodingNonce[:], random[8+32:])
	keyivInverse := Inverse(random[8 : 8+32+16])
	copy(auth.DecodingKey[:], keyivInverse)
	copy(auth.DecodingNonce[:], keyivInverse[32:])
	return auth
}

func ReadAuthentication(reader io.Reader) (*Authentication, error) {
	auth := getAuthenticationObject()

	if _, err := io.ReadFull(reader, auth.Header[:]); err != nil {
		putAuthenticationObject(auth)
		return nil, err
	}

	copy(auth.DecodingKey[:], auth.Header[8:])
	copy(auth.DecodingNonce[:], auth.Header[8+32:])
	keyivInverse := Inverse(auth.Header[8 : 8+32+16])
	copy(auth.EncodingKey[:], keyivInverse)
	copy(auth.EncodingNonce[:], keyivInverse[32:])

	return auth, nil
}

// Inverse returns a new byte array. It is a sequence of bytes when the input is read from end to beginning.Inverse
// Visible for testing only.
func Inverse(b []byte) []byte {
	lenb := len(b)
	b2 := make([]byte, lenb)
	for i, v := range b {
		b2[lenb-i-1] = v
	}
	return b2
}

var (
	authPool = sync.Pool{
		New: func() interface{} {
			return new(Authentication)
		},
	}
)

func getAuthenticationObject() *Authentication {
	return authPool.Get().(*Authentication)
}

func putAuthenticationObject(auth *Authentication) {
	authPool.Put(auth)
}