package rpc

import (
	"context"
	"crypto"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"eta/eta_mini_crm_ht/utils"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
	"io"
	"log"
	"os"
	"time"
)

const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"

type RpcService interface {
	GetPool(addr string) utils.ClientPool
}

type DefaultRpcClient struct {
}

func (d *DefaultRpcClient) WrapSign(request interface{}, doHandler func(ctx context.Context, req interface{}) error) {
	timestamp, nonce, sign, _ := d.signature(request, 16)
	fmt.Println(string(sign))
	ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs(
		"nonce", nonce,
		"timestamp", fmt.Sprintf("%d", timestamp),
		"signature", sign,
	))

	err := doHandler(ctx, request)
	if err != nil {
		utils.FileLog.Error("调用PRC服务失败:%v", err)
	}
}
func (d *DefaultRpcClient) getOptions() (opts []grpc.DialOption) {
	opts = make([]grpc.DialOption, 0)
	opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
	return opts
}

func (d *DefaultRpcClient) GetPool(addr string) utils.ClientPool {
	pool, err := utils.GetPoolInstance(addr, d.getOptions()...)
	if err != nil {
		log.Fatal(err)
	}
	return pool
}

// 生成签名
func (d *DefaultRpcClient) generateSignature(req interface{}) string {
	message := req.(proto.Message)
	reqData, _ := anypb.New(message)
	reqBytes, _ := proto.Marshal(reqData)

	key := []byte("secret") // 秘钥应该保密
	mac := hmac.New(sha256.New, key)
	mac.Write(reqBytes)
	return hex.EncodeToString(mac.Sum(nil))
}

func (d *DefaultRpcClient) generateNonceStr(length int) (string, error) {
	result := make([]byte, length)
	n, err := io.ReadFull(rand.Reader, result)
	if err != nil {
		return "", err
	}
	for i := 0; i < n; i++ {
		result[i] = chars[int(result[i])%len(chars)]
	}
	return string(result), nil
}

type encryptedRequest struct {
	Message   interface{} `json:"ciphertext"`
	Nonce     string      `json:"nonce"`     // 添加随机字符串
	Timestamp int64       `json:"timestamp"` // 添加时间戳
}

func (d *DefaultRpcClient) signature(req interface{}, nonceLen int) (timestamp int64, nonceStr string, encodedData string, err error) {
	nonceStr, err = d.generateNonceStr(nonceLen)
	if err != nil {
		utils.FileLog.Error("随机字符串生成失败: %v", err)
		return
	}
	timestamp = time.Now().UnixMilli()
	// 构建加密后的数据结构
	encryptedStruct := encryptedRequest{
		Message:   req,
		Nonce:     nonceStr,  // 添加随机字符串
		Timestamp: timestamp, // 添加时间戳
	}
	entryData, err := json.Marshal(encryptedStruct)

	if err != nil {
		utils.FileLog.Error("序列化加密后的数据失败: %v", err)
		return
	}

	// 使用私钥加密
	var privateKey *rsa.PrivateKey
	privateKey, err = d.parsePrivateKeyFromPEM()
	if err != nil {
		utils.FileLog.Error("私钥解析失败: %v", err)
		return
	}
	var encryptedData []byte
	hash := sha256.Sum256(entryData)
	// RSA加密
	encryptedData, err = rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hash[:])
	if err != nil {
		utils.FileLog.Error("RSA加密失败: %v", err)
		return
	}
	// 编码加密后的数据
	encodedData = base64.StdEncoding.EncodeToString(encryptedData)
	return
}

// 解析RSA公钥
func (d *DefaultRpcClient) parsePrivateKeyFromPEM() (privateKey *rsa.PrivateKey, err error) {
	pemBlock, err := os.ReadFile("./conf/rsa_private_key.pem")
	block, _ := pem.Decode(pemBlock)
	if block == nil {
		utils.FileLog.Error("公钥解析失败")
		return nil, errors.New("公钥解析失败")
	}
	privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return
}