package main

import (
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"math/big"
	"os"
	"path/filepath"
)

func hash1(curve elliptic.Curve, ID []byte, Qx, Qy, PKx, PKy *big.Int) *big.Int {
	// 创建 SHA-256 散列器
	h := sha256.New()

	// 更新散列器以包含字符串 ID 的内容
	h.Write(ID)

	// 更新散列器以包含点 Q 的坐标信息
	h.Write(Qx.Bytes())
	h.Write(Qy.Bytes())

	// 更新散列器以包含公钥 PK_pub 的坐标信息
	h.Write(PKx.Bytes())
	h.Write(PKy.Bytes())

	// 计算 SHA-256 散列值
	hash := h.Sum(nil)

	// 将散列值转换为大整数
	hBig := new(big.Int).SetBytes(hash)

	// 对大整数 hBig 进行幂运算，幂为 1，模数为 n
	hBig.Mod(hBig, curve.Params().N)

	return hBig
}

func hash2(curve elliptic.Curve, ID []byte, Xx, Xy *big.Int) *big.Int {
	// 计算 hash 值 H_2(ID, X)
	hash := sha256.New()

	hash.Write(ID)

	hash.Write(Xx.Bytes())
	hash.Write(Xy.Bytes())

	// 计算 SHA-256 散列值
	h2 := hash.Sum(nil)

	// 转换为大整数
	h2Big := new(big.Int).SetBytes(h2)

	// 对大整数 h2Big 进行幂运算，幂为 1，模数为 n
	h2Big.Mod(h2Big, curve.Params().N)

	return h2Big
}

func hash3(curve elliptic.Curve, ID, msg []byte, Qx, Qy, Ux, Uy, PKx, PKy *big.Int) *big.Int {
	// 创建 SHA-256 散列器
	h := sha256.New()

	// 更新散列器以包含字符串 ID 的内容
	h.Write(ID)

	// 更新散列器以包含字符串 msg 的内容
	h.Write(msg)

	// 更新散列器以包含点 Q 的坐标信息
	h.Write(Qx.Bytes())
	h.Write(Qy.Bytes())

	// 更新散列器以包含点 U 的坐标信息
	h.Write(Ux.Bytes())
	h.Write(Uy.Bytes())

	// 更新散列器以包含公钥 PK_pub 的坐标信息
	h.Write(PKx.Bytes())
	h.Write(PKy.Bytes())

	// 计算 SHA-256 散列值
	hash := h.Sum(nil)

	// 将散列值转换为大整数
	hBig := new(big.Int).SetBytes(hash)
	hBig.Mod(hBig, curve.Params().N)

	return hBig
}

func equal(lx, ly, rx, ry *big.Int) error {
	if lx.Cmp(rx) == 0 && ly.Cmp(ry) == 0 {
		return nil
	}
	return errors.New("point compare failed ")

}

func genKGCKey(curve elliptic.Curve) (*big.Int, *big.Int, *big.Int) {
	// 随机生成 msk
	msk, _ := rand.Int(rand.Reader, curve.Params().N)
	PKx, PKy := curve.ScalarBaseMult(msk.Bytes())
	return PKx, PKy, msk
}

func genSecret(curve elliptic.Curve) (*big.Int, *big.Int, *big.Int) {
	// 随机生成 x
	x, _ := rand.Int(rand.Reader, curve.Params().N)
	Xx, Xy := curve.ScalarBaseMult(x.Bytes())
	return Xx, Xy, x
}

func genPartialKey(curve elliptic.Curve, ID []byte, msk, Xx, Xy, PKx, PKy *big.Int) (*big.Int, *big.Int, *big.Int) {
	// 随机生成 r
	r, _ := rand.Int(rand.Reader, curve.Params().N)
	Rx, Ry := curve.ScalarBaseMult(r.Bytes())

	//计算h2的hash值
	h2Big := hash2(curve, ID, Xx, Xy)

	//计算 h2Big * X
	tx, ty := curve.ScalarMult(Xx, Xy, h2Big.Bytes())

	//计算 Q = R + h2 * X
	Qx, Qy := curve.Add(Rx, Ry, tx, ty)

	//计算hash1的值
	h1Big := hash1(curve, ID, Qx, Qy, PKx, PKy)

	//计算 d = r + msk * h1 mod n
	temp := new(big.Int).Mul(msk, h1Big)
	d := new(big.Int).Add(r, temp)
	d.Mod(d, curve.Params().N)
	return d, Qx, Qy

}

func VerifyEquation(curve elliptic.Curve, ID []byte, d, Qx, Qy, Xx, Xy, PKx, PKy *big.Int) error {
	//计算等式左值 + h2 * X
	lx, ly := curve.ScalarBaseMult(d.Bytes())
	h2Big := hash2(curve, ID, Xx, Xy)
	h2Xx, h2Xy := curve.ScalarMult(Xx, Xy, h2Big.Bytes())
	lx, ly = curve.Add(lx, ly, h2Xx, h2Xy)

	//计算 Q  + h1 * PK(t)
	h1Big := hash1(curve, ID, Qx, Qy, PKx, PKy)
	tx, ty := curve.ScalarMult(PKx, PKy, h1Big.Bytes())

	rx, ry := curve.Add(Qx, Qy, tx, ty)

	err := equal(lx, ly, rx, ry)
	if err != nil {
		return err
	}

	return nil
}

func genPriKey(curve elliptic.Curve, ID []byte, x, d, Xx, Xy *big.Int) *big.Int {
	//计算h2Big
	h2Big := hash2(curve, ID, Xx, Xy)

	//计算s = d + h2Big * x
	temp := new(big.Int).Mul(h2Big, x)
	sUser := new(big.Int).Add(d, temp)

	//s = s mod n
	sUser = sUser.Mod(sUser, curve.Params().N)

	return sUser
}

func SignThumb(curve elliptic.Curve, ID, msg []byte, s, Qx, Qy, PKx, PKy *big.Int) (*big.Int, *big.Int, *big.Int) {
	// 生成小于素数 的随机数 u
	u, _ := rand.Int(rand.Reader, curve.Params().N)

	// 生成点U
	Ux, Uy := curve.ScalarMult(curve.Params().Gx, curve.Params().Gy, u.Bytes())

	//计算h3
	h3 := hash3(curve, ID, msg, Qx, Qy, Ux, Uy, PKx, PKy)

	// 计算 v = u + h3 * s mod n
	// 计算 tmp = sa * h_3_big
	tmp := new(big.Int).Mul(h3, s)

	// 计算 v = u + tmp
	v := new(big.Int).Add(u, tmp)

	// 计算 v = v mod N
	v.Mod(v, curve.Params().N)

	return Ux, Uy, v
}

func VerifyThumb(curve elliptic.Curve, ID, msg []byte, Qx, Qy, PKx, PKy, Ux, Uy, v *big.Int) error {
	//计算h1和h3
	h1Big := hash1(curve, ID, Qx, Qy, PKx, PKy)
	h3Big := hash3(curve, ID, msg, Qx, Qy, Ux, Uy, PKx, PKy)

	//验证等式 v*P = U + h_3(Q + h_1 * PK_pub)
	//等式左边
	lx, ly := curve.ScalarBaseMult(v.Bytes())

	// 计算 t = h_1_big * PK_pub
	tx, ty := curve.ScalarMult(PKx, PKy, h1Big.Bytes())

	// 计算 out = Q + t
	tx, ty = curve.Add(Qx, Qy, tx, ty)

	// 计算 t = h_3_big * t
	tx, ty = curve.ScalarMult(tx, ty, h3Big.Bytes())

	// 计算 t = U + tmp_p
	rx, ry := curve.Add(Ux, Uy, tx, ty)

	//比较等式左右两值是否相等
	err := equal(lx, ly, rx, ry)
	if err != nil {
		fmt.Println("invalid!")
		return err
	}
	fmt.Println("valid!")
	return nil
}

func writeToFolder(directory, filename string, data []byte) error {
	// 检查 cert 文件夹是否存在
	if _, err := os.Stat(directory); os.IsNotExist(err) {
		// 如果文件夹不存在，则创建它
		err := os.Mkdir(directory, 0755)
		if err != nil {
			return fmt.Errorf("创建文件夹时出错: %v", err)
		}
	}

	// 构造文件的完整路径
	filePath := filepath.Join(directory, filename)

	err := os.WriteFile(filePath, []byte(data), 0644)
	if err != nil {
		return fmt.Errorf("写入文件时出错: %v", err)
	}
	fmt.Println("字符串已成功写入文件:", filePath)
	return nil
}

func mergeBigInts(bigInts ...*big.Int) []byte {
	var mergedBytes []byte

	for _, num := range bigInts {
		numBytes := []byte(num.Text(16))
		mergedBytes = append(mergedBytes, numBytes...)
	}

	return mergedBytes
}

func readFileFromDirectory(directory, filename string) ([]byte, error) {
	// 构造文件的完整路径
	filePath := filepath.Join(directory, filename)

	// 读取文件内容
	data, err := os.ReadFile(filePath)
	if err != nil {
		return nil, err
	}

	return data, nil
}

func BytesToBig(bytes []byte) *big.Int {
	//不能直接转为big，转为字符串后用16进制解码
	big := new(big.Int)
	big.SetString(string(bytes), 16)
	return big
}

func bytes_halve(slice []byte) ([]byte, []byte) {
	length := len(slice)

	// 计算每个部分的长度
	halfLength := length / 2

	// 切分切片
	part1 := slice[:halfLength]
	part2 := slice[halfLength:]

	return part1, part2
}
