package store

import (
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"time"

	"github.com/niko/qcc/pkg/types"
)

type FrameHeader struct {
	Magic       [4]byte
	Version     uint16
	FrameID     uint64
	PrevFrameID uint64
	PrevHash    [32]byte
	Timestamp   int64
	NumRecords  uint32
}

const FrameHeaderSize = 4 + 2 + 8 + 8 + 32 + 8 + 4
const RecordHeaderSize = 1 + 1 + 2 + 8 + 4 + 12 + 16

type RecordHeader struct {
	RecType     types.RecType
	Flags       uint8
	TableID     uint16
	KeyHash     uint64
	PayloadSize uint32
	Nonce       [12]byte
	AuthTag     [16]byte
}

type Frame struct {
	Header  FrameHeader
	Records []Record
}

type Record struct {
	Header  RecordHeader
	Payload []byte
}

func NewFrameHeader(prevFrameID uint64, prevHash [32]byte) FrameHeader {
	var magic [4]byte
	copy(magic[:], []byte("QCCF"))
	return FrameHeader{
		Magic:       magic,
		Version:     1,
		FrameID:     prevFrameID + 1,
		PrevFrameID: prevFrameID,
		PrevHash:    prevHash,
		Timestamp:   time.Now().UnixNano(),
	}
}

func MarshalFrame(f *Frame) ([]byte, error) {
	size := FrameHeaderSize
	for _, rec := range f.Records {
		size += RecordHeaderSize + len(rec.Payload)
	}
	padding := (8 - (size % 8)) % 8
	size += padding + 32

	buf := make([]byte, size)
	off := 0

	copy(buf[off:], f.Header.Magic[:])
	off += 4
	binary.BigEndian.PutUint16(buf[off:], f.Header.Version)
	off += 2
	binary.BigEndian.PutUint64(buf[off:], f.Header.FrameID)
	off += 8
	binary.BigEndian.PutUint64(buf[off:], f.Header.PrevFrameID)
	off += 8
	copy(buf[off:], f.Header.PrevHash[:])
	off += 32
	binary.BigEndian.PutUint64(buf[off:], uint64(f.Header.Timestamp))
	off += 8
	binary.BigEndian.PutUint32(buf[off:], uint32(len(f.Records)))
	off += 4

	for _, rec := range f.Records {
		buf[off] = byte(rec.Header.RecType)
		off++
		buf[off] = rec.Header.Flags
		off++
		binary.BigEndian.PutUint16(buf[off:], rec.Header.TableID)
		off += 2
		binary.BigEndian.PutUint64(buf[off:], rec.Header.KeyHash)
		off += 8
		binary.BigEndian.PutUint32(buf[off:], rec.Header.PayloadSize)
		off += 4
		copy(buf[off:], rec.Header.Nonce[:])
		off += 12
		copy(buf[off:], rec.Header.AuthTag[:])
		off += 16
		copy(buf[off:], rec.Payload)
		off += len(rec.Payload)
	}

	for i := 0; i < padding; i++ {
		buf[off] = 0
		off++
	}

	hash := sha256.Sum256(buf[:off])
	copy(buf[off:], hash[:])
	off += 32

	return buf[:off], nil
}

func UnmarshalFrame(data []byte) (*Frame, error) {
	if len(data) < FrameHeaderSize+32 {
		return nil, fmt.Errorf("frame too short: %d", len(data))
	}
	dataLen := len(data)
	checksumStart := dataLen - 32
	storedChecksum := data[checksumStart:]

	computed := sha256.Sum256(data[:checksumStart])
	if !constantTimeEqual(storedChecksum, computed[:]) {
		return nil, fmt.Errorf("frame checksum mismatch")
	}

	f := &Frame{}
	off := 0

	copy(f.Header.Magic[:], data[off:off+4])
	off += 4
	if string(f.Header.Magic[:]) != "QCCF" {
		return nil, fmt.Errorf("invalid frame magic: %s", hex.EncodeToString(f.Header.Magic[:]))
	}
	f.Header.Version = binary.BigEndian.Uint16(data[off:])
	off += 2
	f.Header.FrameID = binary.BigEndian.Uint64(data[off:])
	off += 8
	f.Header.PrevFrameID = binary.BigEndian.Uint64(data[off:])
	off += 8
	copy(f.Header.PrevHash[:], data[off:off+32])
	off += 32
	f.Header.Timestamp = int64(binary.BigEndian.Uint64(data[off:]))
	off += 8
	f.Header.NumRecords = binary.BigEndian.Uint32(data[off:])
	off += 4

	for i := uint32(0); i < f.Header.NumRecords; i++ {
		if off+RecordHeaderSize > checksumStart {
			return nil, fmt.Errorf("record header %d exceeds frame bounds", i)
		}
		var rec Record
		rec.Header.RecType = types.RecType(data[off])
		off++
		rec.Header.Flags = data[off]
		off++
		rec.Header.TableID = binary.BigEndian.Uint16(data[off:])
		off += 2
		rec.Header.KeyHash = binary.BigEndian.Uint64(data[off:])
		off += 8
		rec.Header.PayloadSize = binary.BigEndian.Uint32(data[off:])
		off += 4
		copy(rec.Header.Nonce[:], data[off:off+12])
		off += 12
		copy(rec.Header.AuthTag[:], data[off:off+16])
		off += 16

		if off+int(rec.Header.PayloadSize) > checksumStart {
			return nil, fmt.Errorf("record payload %d exceeds frame bounds", i)
		}
		rec.Payload = make([]byte, rec.Header.PayloadSize)
		copy(rec.Payload, data[off:off+int(rec.Header.PayloadSize)])
		off += int(rec.Header.PayloadSize)
		f.Records = append(f.Records, rec)
	}

	return f, nil
}

func ComputeFrameChecksum(frameData []byte) [32]byte {
	return sha256.Sum256(frameData)
}

func constantTimeEqual(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	var v byte
	for i := 0; i < len(a); i++ {
		v |= a[i] ^ b[i]
	}
	return v == 0
}