package fqapi

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"math/big"
)

// CM 加密管理器
type CM struct {
	bigC  *big.Int    // 大素数C
	bigB  int64       // 底数B
	arryD []byte      // 固定密钥
	thisD [2]*big.Int // 密钥对(私钥, 公钥)
	thisC []byte      // 随机IV
}

// NewCM 创建新的CM实例
func NewCM() *CM {
	// 大素数C
	bigCStr := "2410312426921032588552076022197566074856950548502459942654116941958108831682612228890093858261341614673227141477904012196503648957050582631942730706805009223062734745341073406696246014589361659774041027169249453200378729434170325843778659198143763193776859869524088940195577346119843545301547043747207749969763750084308926339295559968882457872412993810129130294592999947926365264059284647209730384947211681434464714438488520940127459844288859336526896320919633919"
	bigC := new(big.Int)
	bigC.SetString(bigCStr, 10)

	// 固定密钥（base64解码）
	arryD, _ := base64.StdEncoding.DecodeString("rCXGfd2POMGzeiNIgo4iLg==")

	return &CM{
		bigC:  bigC,
		bigB:  2,
		arryD: arryD,
	}
}

// generateKeyPair 生成密钥对
func (cm *CM) generateKeyPair() [2]*big.Int {
	// 生成随机私钥
	randomBytes := make([]byte, 32)
	rand.Read(randomBytes)

	x := new(big.Int).SetBytes(randomBytes)
	x = x.Mod(x, new(big.Int).Sub(cm.bigC, big.NewInt(1)))

	// 计算公钥: y = b^x mod c
	y := new(big.Int).Exp(big.NewInt(cm.bigB), x, cm.bigC)

	cm.thisD = [2]*big.Int{x, y}
	return cm.thisD
}

// pairToBytes 将公钥转换为字节数组
func (cm *CM) pairToBytes(keyPair [2]*big.Int) []byte {
	return keyPair[1].Bytes()
}

// randIV 生成随机IV
func (cm *CM) randIV() []byte {
	iv := make([]byte, 16)
	rand.Read(iv)
	return iv
}

// encodeAES AES-CBC加密
func (cm *CM) encodeAES(key, iv, data []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// PKCS7填充
	padding := aes.BlockSize - len(data)%aes.BlockSize
	padText := make([]byte, padding)
	for i := range padText {
		padText[i] = byte(padding)
	}
	paddedData := append(data, padText...)

	ciphertext := make([]byte, len(paddedData))
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext, paddedData)

	return ciphertext, nil
}

// decodeAES AES-CBC解密
func (cm *CM) decodeAES(key, iv, data []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	decrypted := make([]byte, len(data))
	mode.CryptBlocks(decrypted, data)

	// 移除PKCS7填充
	padding := int(decrypted[len(decrypted)-1])
	if padding > len(decrypted) || padding > aes.BlockSize {
		return nil, fmt.Errorf("无效的填充")
	}

	return decrypted[:len(decrypted)-padding], nil
}

// ClientHandshake 客户端握手，生成请求头y
func (cm *CM) ClientHandshake() string {
	keyPair := cm.generateKeyPair()
	cm.thisC = cm.randIV()

	publicKeyBytes := cm.pairToBytes(keyPair)
	encrypted, _ := cm.encodeAES(cm.arryD, cm.thisC, publicKeyBytes)

	combined := append(cm.thisC, encrypted...)
	return base64.StdEncoding.EncodeToString(combined)
}

// decryptF 解密函数
func (cm *CM) decryptF(peerPublicKey []byte, keyPair [2]*big.Int, iv, data []byte) ([]byte, error) {
	x := keyPair[0]

	// 计算共享密钥: shared_secret = peer_public_key^x mod c
	peerPubInt := new(big.Int).SetBytes(peerPublicKey)
	sharedSecret := new(big.Int).Exp(peerPubInt, x, cm.bigC)

	key := sharedSecret.Bytes()
	// 确保key长度为32字节
	if len(key) < 32 {
		// 前面补0
		paddedKey := make([]byte, 32)
		copy(paddedKey[32-len(key):], key)
		key = paddedKey
	} else if len(key) > 32 {
		key = key[:32]
	}

	return cm.decodeAES(key, iv, data)
}

// Decrypt 解密响应内容
func (cm *CM) Decrypt(c, y, content string) (string, error) {
	// c 应该为 "1" 表示启用加密
	// y 是服务器返回的公钥
	// content 是加密的内容

	if c != "1" || y == "" {
		return content, nil
	}

	// 解码content
	contentDecoded, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		return "", err
	}

	if len(contentDecoded) < 16 {
		return "", fmt.Errorf("内容太短")
	}

	// 提取IV和数据
	iv := contentDecoded[:16]
	data := contentDecoded[16:]

	// 解码服务器公钥
	peerPublicKey, err := base64.StdEncoding.DecodeString(y)
	if err != nil {
		return "", err
	}

	// 检查密钥对是否已初始化
	if cm.thisD[0] == nil {
		return "", fmt.Errorf("密钥对未初始化")
	}

	// 解密
	decrypted, err := cm.decryptF(peerPublicKey, cm.thisD, iv, data)
	if err != nil {
		return "", err
	}

	return string(decrypted), nil
}