package deploy

import (
	"bytes"
	"fmt"

	"github.com/klauspost/reedsolomon"
)

const (
	erasureSchemeReedSolomonV1 = "reed_solomon_v1"
	erasureDefaultDataShards   = 4
	erasureDefaultParityShards = 2
)

type erasureParams struct {
	DataShards   int
	ParityShards int
}

func (p erasureParams) normalize() (erasureParams, error) {
	if p.DataShards <= 0 {
		p.DataShards = erasureDefaultDataShards
	}
	if p.ParityShards <= 0 {
		p.ParityShards = erasureDefaultParityShards
	}
	if p.DataShards < 1 || p.ParityShards < 1 {
		return erasureParams{}, fmt.Errorf("erasure: invalid shard counts")
	}
	return p, nil
}

func (p erasureParams) minShards() int {
	return p.DataShards
}

func (p erasureParams) totalShards() int {
	return p.DataShards + p.ParityShards
}

// decodeErasureShards reconstructs payload from Reed–Solomon shards (nil = missing).
func decodeErasureShards(shards [][]byte, payloadSize int, p erasureParams) ([]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: verification failed")
	}
	if payloadSize <= 0 {
		payloadSize = len(shards[0]) * p.DataShards
	}
	var out bytes.Buffer
	if err := enc.Join(&out, shards, payloadSize); err != nil {
		return nil, err
	}
	return out.Bytes(), nil
}