// rsa.go
package crypt

import (
	"context"
	"fmt"

	"github.com/sashabaranov/go-openai"
	"github.com/sashabaranov/go-openai/jsonschema"
)

// ========== RSA签名工具 ==========
type RSASignTool struct{}

func (t *RSASignTool) Name() string { return "rsa_sign" }

func (t *RSASignTool) Description() string {
	return "RSA私钥签名实现指南。当文档要求使用RSA算法对请求进行数字签名时使用，包含签名串拼接、私钥加载、PKCS1/PKCS8格式处理"
}

func (t *RSASignTool) GetFunctionDefinition() openai.FunctionDefinition {
	return openai.FunctionDefinition{
		Name:        t.Name(),
		Description: t.Description(),
		Parameters: jsonschema.Definition{
			Type: jsonschema.Object,
			Properties: map[string]jsonschema.Definition{
				"hash_algorithm": {
					Type:        jsonschema.String,
					Description: "哈希算法，可选 SHA256, SHA384, SHA512，默认 SHA256",
					Enum:        []string{"SHA256", "SHA384", "SHA512"},
				},
				"encoding": {
					Type:        jsonschema.String,
					Description: "输出编码方式，base64 或 hex，默认 base64",
					Enum:        []string{"base64", "hex"},
				},
			},
			Required: []string{},
		},
	}
}

func (t *RSASignTool) GetDetail(ctx context.Context, params map[string]string) (string, error) {
	hashAlgo := "SHA256"
	if v, ok := params["hash_algorithm"]; ok && v != "" {
		hashAlgo = v
	}
	encoding := "base64"
	if v, ok := params["encoding"]; ok && v != "" {
		encoding = v
	}

	return fmt.Sprintf(`
### RSA签名完整实现指南

**适用场景**：文档要求使用RSA私钥对请求参数进行签名

**配置参数**：
- 哈希算法: %s
- 编码方式: %s

**完整代码模板**：
`+"```go"+`
package crypto

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

// RSAConfig RSA签名配置
type RSAConfig struct {
    PrivateKeyPEM string // PEM格式的私钥
    KeyID         string // 密钥ID（如文档要求携带）
}

// SignWithRSA 使用RSA私钥对请求进行签名
func SignWithRSA(params map[string]string, config *RSAConfig) (string, error) {
    // 1. 拼接参数（按字典序排序）
    keys := make([]string, 0, len(params))
    for k := range params {
        if k != "sign" && k != "signature" {
            keys = append(keys, k)
        }
    }
    sort.Strings(keys)
    
    var sb strings.Builder
    for _, k := range keys {
        if sb.Len() > 0 {
            sb.WriteString("&")
        }
        sb.WriteString(k)
        sb.WriteString("=")
        sb.WriteString(params[k])
    }
    signStr := sb.String()
    
    // 2. 加载私钥
    block, _ := pem.Decode([]byte(config.PrivateKeyPEM))
    if block == nil {
        return "", fmt.Errorf("failed to decode PEM private key")
    }
    
    privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
    if err != nil {
        // 尝试PKCS1格式
        privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
        if err != nil {
            return "", fmt.Errorf("failed to parse private key: %%v", err)
        }
    }
    
    rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey)
    if !ok {
        return "", fmt.Errorf("not a RSA private key")
    }
    
    // 3. 计算哈希
    var hashed []byte
    var hashFunc crypto.Hash
    
    switch "%s" {
    case "SHA384":
        h := sha512.New384()
        h.Write([]byte(signStr))
        hashed = h.Sum(nil)
        hashFunc = crypto.SHA384
    case "SHA512":
        h := sha512.New()
        h.Write([]byte(signStr))
        hashed = h.Sum(nil)
        hashFunc = crypto.SHA512
    default: // SHA256
        h := sha256.New()
        h.Write([]byte(signStr))
        hashed = h.Sum(nil)
        hashFunc = crypto.SHA256
    }
    
    // 4. RSA签名
    signature, err := rsa.SignPKCS1v15(rand.Reader, rsaPrivateKey, hashFunc, hashed)
    if err != nil {
        return "", err
    }
    
    // 5. 编码
    if "%s" == "hex" {
        return hex.EncodeToString(signature), nil
    }
    return base64.StdEncoding.EncodeToString(signature), nil
}
`+"```"+`

**注意事项**：
- 确认文档使用的是 PKCS1 还是 PKCS8 格式
- 确认是否需要添加额外的固定盐值
- 确认编码是 Base64 还是 Hex
- 注意排除签名字段本身（sign/signature）
`, hashAlgo, encoding, hashAlgo, encoding), nil
}

func (t *RSASignTool) Execute(ctx context.Context, params map[string]string) (string, error) {
	return t.GetDetail(ctx, params)
}