package nezha

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"sort"
	"strconv"
	"strings"
)

// buildSignSource 实现哪吒 sign_note 的签名串规则:收集所有*非空*参数,排除
// sign/sign_type/数组参数(调用方在构造 params 时就不应放入非标量值),按参数名
// ASCII 升序排序,拼成 "key=value&key2=value2"(不拼商户密钥、不做任何 URL 编码)。
// sort.Strings 是逐字节比较,对本协议里全 ASCII 的参数名等价于 ASCII 排序。
func buildSignSource(params map[string]string) string {
	keys := make([]string, 0, len(params))
	for k, v := range params {
		if k == "sign" || k == "sign_type" || v == "" {
			continue
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	parts := make([]string, 0, len(keys))
	for _, k := range keys {
		parts = append(parts, k+"="+params[k])
	}
	return strings.Join(parts, "&")
}

// signRSA 对签名串做 SHA256WithRSA + Base64(sign_note 规定算法)。
func signRSA(priv *rsa.PrivateKey, source string) (string, error) {
	h := sha256.Sum256([]byte(source))
	sig, err := rsa.SignPKCS1v15(rand.Reader, priv, crypto.SHA256, h[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(sig), nil
}

// verifyRSA 验 SHA256WithRSA + Base64 签名。签名不合法/验签失败均返回 error。
func verifyRSA(pub *rsa.PublicKey, source, sigB64 string) error {
	sig, err := base64.StdEncoding.DecodeString(sigB64)
	if err != nil {
		return fmt.Errorf("nezha: sign 不是合法 base64: %w", err)
	}
	h := sha256.Sum256([]byte(source))
	return rsa.VerifyPKCS1v15(pub, crypto.SHA256, h[:], sig)
}

// parsePrivateKey 解析商户 RSA 私钥,兼容 PKCS1/PKCS8、PEM 头或裸 base64 DER
// (与 alipay 账户凭证"无 PEM 头亦可"的既有惯例对齐,详见 config.AlipaySandboxConfig
// 字段注释)。
func parsePrivateKey(s string) (*rsa.PrivateKey, error) {
	der, err := pemOrBase64ToDER(s)
	if err != nil {
		return nil, err
	}
	if key, perr := x509.ParsePKCS8PrivateKey(der); perr == nil {
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("nezha: 私钥不是 RSA 类型(PKCS8)")
		}
		return rsaKey, nil
	}
	if key, perr := x509.ParsePKCS1PrivateKey(der); perr == nil {
		return key, nil
	}
	return nil, fmt.Errorf("nezha: 无法解析商户私钥(需 PKCS1/PKCS8,PEM 或裸 base64 DER)")
}

// parsePublicKey 解析平台 RSA 公钥,兼容 PKIX/PKCS1、PEM 头或裸 base64 DER。
func parsePublicKey(s string) (*rsa.PublicKey, error) {
	der, err := pemOrBase64ToDER(s)
	if err != nil {
		return nil, err
	}
	if key, perr := x509.ParsePKIXPublicKey(der); perr == nil {
		rsaKey, ok := key.(*rsa.PublicKey)
		if !ok {
			return nil, fmt.Errorf("nezha: 公钥不是 RSA 类型(PKIX)")
		}
		return rsaKey, nil
	}
	if key, perr := x509.ParsePKCS1PublicKey(der); perr == nil {
		return key, nil
	}
	return nil, fmt.Errorf("nezha: 无法解析平台公钥(需 PKIX/PKCS1,PEM 或裸 base64 DER)")
}

func pemOrBase64ToDER(s string) ([]byte, error) {
	s = strings.TrimSpace(s)
	if block, _ := pem.Decode([]byte(s)); block != nil {
		return block.Bytes, nil
	}
	clean := strings.Join(strings.Fields(s), "")
	der, err := base64.StdEncoding.DecodeString(clean)
	if err != nil {
		return nil, fmt.Errorf("nezha: 密钥既非 PEM 也非合法 base64: %w", err)
	}
	return der, nil
}

// stringifyParams 把渠道 JSON 响应(map[string]any,数字统一是 float64)拍平成
// map[string]string 供 buildSignSource 用;数组/对象/nil 一律丢弃(sign_note 排除
// 数组参数,响应里也不应出现需要参与签名的复合结构)。
func stringifyParams(raw map[string]any) map[string]string {
	out := make(map[string]string, len(raw))
	for k, v := range raw {
		switch t := v.(type) {
		case string:
			out[k] = t
		case float64:
			out[k] = strconv.FormatFloat(t, 'f', -1, 64)
		case bool:
			out[k] = strconv.FormatBool(t)
		}
	}
	return out
}

// verifyResponseSign 用平台公钥验渠道 JSON 响应(下单/查单)自带的 sign。
func verifyResponseSign(pub *rsa.PublicKey, raw map[string]any) error {
	sig, _ := raw["sign"].(string)
	if sig == "" {
		return fmt.Errorf("响应缺少 sign 字段")
	}
	source := buildSignSource(stringifyParams(raw))
	return verifyRSA(pub, source, sig)
}