package hash

import (
	"crypto/sha256"
	"encoding/hex"
)

// MerkleTree 表示一棵基于 SHA-256 的 Merkle 树。
// 叶子为各分段（或各集）的哈希，根用于整体内容的聚合锚定。
// 对应需求：需求1-AC2、需求16-AC4（按集定位篡改）。
type MerkleTree struct {
	Leaves []string   // 叶子哈希（十六进制）
	Levels [][]string // 自底向上的各层，Levels[0] 为叶子层
}

// BuildMerkleTree 由叶子哈希构建 Merkle 树。
// 当某层节点数为奇数时，复制最后一个节点与自身配对（标准做法）。
func BuildMerkleTree(leaves []string) *MerkleTree {
	mt := &MerkleTree{Leaves: append([]string(nil), leaves...)}
	if len(leaves) == 0 {
		mt.Levels = [][]string{{}}
		return mt
	}

	level := append([]string(nil), leaves...)
	mt.Levels = [][]string{level}

	for len(level) > 1 {
		next := make([]string, 0, (len(level)+1)/2)
		for i := 0; i < len(level); i += 2 {
			left := level[i]
			right := left // 奇数个时与自身配对
			if i+1 < len(level) {
				right = level[i+1]
			}
			next = append(next, hashPair(left, right))
		}
		mt.Levels = append(mt.Levels, next)
		level = next
	}
	return mt
}

// Root 返回 Merkle 根哈希；空树返回空字符串。
func (mt *MerkleTree) Root() string {
	if len(mt.Levels) == 0 {
		return ""
	}
	top := mt.Levels[len(mt.Levels)-1]
	if len(top) == 0 {
		return ""
	}
	return top[0]
}

// hashPair 将两个十六进制哈希拼接后再次 SHA-256。
func hashPair(left, right string) string {
	lb, _ := hex.DecodeString(left)
	rb, _ := hex.DecodeString(right)
	h := sha256.New()
	h.Write(lb)
	h.Write(rb)
	return hex.EncodeToString(h.Sum(nil))
}

// LocateChangedLeaves 比较两组叶子哈希，返回发生变化的叶子索引。
// 用于"定位被篡改的具体集"（需求12-AC3）。
func LocateChangedLeaves(oldLeaves, newLeaves []string) []int {
	var changed []int
	max := len(oldLeaves)
	if len(newLeaves) > max {
		max = len(newLeaves)
	}
	for i := 0; i < max; i++ {
		var o, n string
		if i < len(oldLeaves) {
			o = oldLeaves[i]
		}
		if i < len(newLeaves) {
			n = newLeaves[i]
		}
		if o != n {
			changed = append(changed, i)
		}
	}
	return changed
}