// Package erasure provides Reed–Solomon k-of-n shard encode/decode for spread payloads.
package erasure

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"fmt"

	"github.com/klauspost/reedsolomon"
)

const (
	// SchemeReedSolomonV1 is the deploy-plan erasure scheme identifier.
	SchemeReedSolomonV1 = "reed_solomon_v1"
	// DefaultDataShards is the data shard count for spread payload encoding.
	DefaultDataShards = 4
	// DefaultParityShards is the parity shard count (any DefaultDataShards of total reconstruct).
	DefaultParityShards = 2
)

// Params describes a Reed–Solomon split.
type Params struct {
	DataShards   int
	ParityShards int
}

// DefaultParams returns the standard 4+2 erasure split.
func DefaultParams() Params {
	return Params{DataShards: DefaultDataShards, ParityShards: DefaultParityShards}
}

// MinShards returns the minimum shard count required for reconstruction.
func (p Params) MinShards() int {
	if p.DataShards <= 0 {
		return 0
	}
	return p.DataShards
}

// TotalShards returns data + parity shard count.
func (p Params) TotalShards() int {
	return p.DataShards + p.ParityShards
}

// Normalize fills zero values with defaults and validates counts.
func (p Params) Normalize() (Params, error) {
	if p.DataShards <= 0 {
		p.DataShards = DefaultDataShards
	}
	if p.ParityShards <= 0 {
		p.ParityShards = DefaultParityShards
	}
	if p.DataShards < 1 || p.ParityShards < 1 {
		return Params{}, fmt.Errorf("erasure: invalid shard counts data=%d parity=%d", p.DataShards, p.ParityShards)
	}
	if p.DataShards+p.ParityShards > 256 {
		return Params{}, fmt.Errorf("erasure: too many shards")
	}
	return p, nil
}

// Encode splits payload into equal-sized Reed–Solomon shards.
// The returned size is the original payload length (for Join on decode).
func Encode(data []byte, p Params) ([][]byte, int, error) {
	p, err := p.Normalize()
	if err != nil {
		return nil, 0, err
	}
	if len(data) == 0 {
		return nil, 0, fmt.Errorf("erasure: empty payload")
	}
	origSize := len(data)
	enc, err := reedsolomon.New(p.DataShards, p.ParityShards)
	if err != nil {
		return nil, 0, err
	}
	shards, err := enc.Split(data)
	if err != nil {
		return nil, 0, err
	}
	if err := enc.Encode(shards); err != nil {
		return nil, 0, err
	}
	return shards, origSize, nil
}

// Decode reconstructs payload from at least MinShards() shards (nil entries allowed for missing).
func Decode(shards [][]byte, payloadSize int, p Params) ([]byte, error) {
	p, err := p.Normalize()
	if err != nil {
		return nil, err
	}
	if len(shards) < p.TotalShards() {
		return nil, fmt.Errorf("erasure: shard slice too short")
	}
	present := 0
	for i := 0; i < p.TotalShards(); i++ {
		if len(shards[i]) > 0 {
			present++
		}
	}
	if present < p.MinShards() {
		return nil, fmt.Errorf("erasure: need %d shards, have %d", p.MinShards(), present)
	}
	enc, err := reedsolomon.New(p.DataShards, p.ParityShards)
	if err != nil {
		return nil, err
	}
	if err := enc.Reconstruct(shards); err != nil {
		return nil, err
	}
	ok, err := enc.Verify(shards)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("erasure: shard verification failed")
	}
	if payloadSize <= 0 {
		payloadSize = len(shards[0]) * p.DataShards
	}
	var buf bytes.Buffer
	if err := enc.Join(&buf, shards, payloadSize); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// PayloadSHA256 returns hex SHA256 of the original payload.
func PayloadSHA256(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:])
}