package relay

import (
	"WhspBrd/menc"
	"WhspBrd/owner"
	"WhspBrd/thrembio"
	"WhspBrd/typio/splco"
	"crypto/mlkem"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"
)

type Client interface {
	Register(tokenId uint32, token []byte) error
	Login() error
	PublishKem() error
	GetKem(whos owner.Identity) (*mlkem.EncapsulationKey1024, error)
	Send(to owner.Identity, encap *mlkem.EncapsulationKey1024, data []byte) error
	MessageLen() uint32
	MessageGet(id uint32) (sender owner.Identity, timestamp uint32, data []byte, err error)
}

type client struct {
	th    thrembio.Client
	sc    owner.Secret
	spuha owner.Identity
}

func NewClient(addr string, secret owner.Secret, spuha owner.Identity) (Client, error) {
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, fmt.Errorf("invalid address, expected ip:port: %w", err)
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return nil, fmt.Errorf("invalid port: %w", err)
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return nil, fmt.Errorf("invalid ip: %s", host)
	}
	th, err := thrembio.NewClient(&net.UDPAddr{IP: ip, Port: port}, secret)
	if err != nil {
		return nil, err
	}
	return &client{
		th:    th,
		sc:    secret,
		spuha: spuha,
	}, nil
}

func (c *client) Register(tokenId uint32, token []byte) error {
	return c.th.Register(tokenId, token)
}

func (c *client) Login() error {
	return c.th.Login()
}

func (c *client) PublishKem() error {
	// [serverPuha] [kek] [osign(serverPuha | kek)]
	data := append(c.spuha[:], c.sc.EncapKey()...)
	osign, err := c.sc.Sign(data)
	if err != nil {
		return err
	}
	data = append(data, osign...)
	data = append([]byte("pubkem"), data...)
	err = c.th.Write(data)
	if err != nil {
		return err
	}
	res, err := c.th.Read()
	if err != nil {
		return err
	}
	if string(res) != "done" {
		return err
	}
	return nil
}

func (c *client) GetKem(whos owner.Identity) (*mlkem.EncapsulationKey1024, error) {
	// [receiverPuha]
	// <- ([serverPuha] [kek] [osign(serverPuha | kek)])(of that receiverPuha)
	err := c.th.Write(append([]byte("getkem"), whos[:]...))
	if err != nil {
		return nil, err
	}
	res, err := c.th.Read()
	if err != nil {
		return nil, err
	}
	if len(res) < owner.IdentitySize+mlkem.EncapsulationKeySize1024 {
		return nil, errors.New("invalid response length")
	}
	serverPuha := res[:owner.IdentitySize]
	kek := res[owner.IdentitySize : owner.IdentitySize+mlkem.EncapsulationKeySize1024]
	osign := res[owner.IdentitySize+mlkem.EncapsulationKeySize1024:]
	whosver, err := owner.Verify(append(serverPuha, kek...), osign)
	if err != nil {
		return nil, err
	}
	if whosver != whos {
		return nil, errors.New("not signed by the expected owner")
	}
	return mlkem.NewEncapsulationKey1024(kek)
}

func (c *client) Send(to owner.Identity, encap *mlkem.EncapsulationKey1024, data []byte) error {
	// -> [serverPuha] [receiverPuha] [kct] [time] [encData] [osign(serverPuha | receiverPuha | kct | time | encData)]
	// <- "send"
	shared, kct := encap.Encapsulate()
	key := sha256.Sum256(shared)
	timee := make([]byte, 8)
	binary.BigEndian.PutUint64(timee, uint64(time.Now().Unix()))
	encData, err := menc.AESGCM_Quick_Encrypt(key[:], data, timee)
	if err != nil {
		return err
	}
	payload := splco.Append(c.spuha[:], to[:], kct, timee, encData)
	osign, err := c.sc.Sign(payload)
	if err != nil {
		return err
	}
	payload = append(payload, osign...)
	err = c.th.Write(append([]byte("send"), payload...))
	if err != nil {
		return err
	}
	res, err := c.th.Read()
	if err != nil {
		return err
	}
	if string(res) != "send" {
		return errors.New("unexpected response")
	}
	return nil
}

func (c *client) MessageLen() uint32 {
	// -> "msglen"
	// <- [len]
	err := c.th.Write([]byte("msglen"))
	if err != nil {
		return 0
	}
	res, err := c.th.Read()
	if err != nil {
		return 0
	}
	if len(res) != 4 {
		return 0
	}
	return binary.BigEndian.Uint32(res)
}

func (c *client) MessageGet(id uint32) (sender owner.Identity, timestamp uint32, data []byte, err error) {
	// -> [id]
	// <- [serverPuha] [receiverPuha] [kct] [time] [encData] [osign(serverPuha | senderPuha | kct | time | encData)]
	idB := make([]byte, 4)
	binary.BigEndian.PutUint32(idB, id)
	err = c.th.Write(append([]byte("msgget"), idB...))
	if err != nil {
		return owner.Identity{}, 0, nil, err
	}
	res, err := c.th.Read()
	if err != nil {
		return owner.Identity{}, 0, nil, err
	}
	if len(res) < owner.IdentitySize*2+mlkem.CiphertextSize1024+8 {
		return owner.Identity{}, 0, nil, errors.New("invalid response length")
	}
	serverPuha := owner.Identity(res[:owner.IdentitySize])
	//receiverPuha := owner.Identity(res[owner.IdentitySize : owner.IdentitySize*2])
	kct := res[owner.IdentitySize*2 : owner.IdentitySize*2+mlkem.CiphertextSize1024]
	timee := res[owner.IdentitySize*2+mlkem.CiphertextSize1024 : owner.IdentitySize*2+mlkem.CiphertextSize1024+8]
	encData := res[owner.IdentitySize*2+mlkem.CiphertextSize1024+8 : owner.IdentitySize*2+mlkem.CiphertextSize1024+8+len(res)-owner.IdentitySize*2-mlkem.CiphertextSize1024-8]
	osign := res[len(res)-owner.IdentitySize:]
	verData := res[:len(res)-owner.IdentitySize]
	whosver, err := owner.Verify(verData, osign)
	if err != nil {
		return owner.Identity{}, 0, nil, err
	}
	shared, err := c.sc.Decapsulate(kct)
	if err != nil {
		return owner.Identity{}, 0, nil, err
	}
	key := sha256.Sum256(shared)
	data, err = menc.AESGCM_Quick_Decrypt(key[:], encData, timee)
	if err != nil {
		return owner.Identity{}, 0, nil, err
	}
	if owner.IdentityEq(c.spuha, serverPuha) {
		return owner.Identity{}, 0, nil, errors.New("unexpected sender")
	}
	return whosver, binary.BigEndian.Uint32(timee), data, nil

}