package main

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/c0mm4nd/go-ripemd"
	"github.com/go-piv/piv-go/v2/piv"
	"github.com/martinlindhe/gogost/gost34112012256"
	"github.com/tjfoc/gmsm/sm3"
)

// Ed25519 constants (EXACTLY from yubisigner.go:95-99)
const (
	Ed25519SignatureSize = 64
	Ed25519PublicKeySize = 32
	Ed25519CombinedSize  = Ed25519SignatureSize + Ed25519PublicKeySize
)

// Algorithm constant (EXACTLY from yubisigner.go:69-73)
const (
	AlgorithmED25519 = "ED25519"
)

// normalizeToCRLF converts any line ending to RFC-compliant CRLF (EXACTLY from yubisigner.go:538-544)
func normalizeToCRLF(data []byte) []byte {
	s := string(data)
	s = strings.ReplaceAll(s, "\r\n", "\n")
	s = strings.ReplaceAll(s, "\r", "\n")
	s = strings.ReplaceAll(s, "\n", "\r\n")
	return []byte(s)
}

// ensureUTF8 ensures the string is valid UTF-8 (EXACTLY from yubisigner.go:547-552)
func ensureUTF8(s string) string {
	if utf8.ValidString(s) {
		return s
	}
	return strings.ToValidUTF8(s, " ")
}

// calculateHashesRAM calculates all hashes for files that fit in RAM (EXACTLY from yubisigner.go:1621-1640)
func calculateHashesRAM(data []byte) map[string]string {
	hashes := make(map[string]string)

	gostHasher := gost34112012256.New()
	gostHasher.Write(data)
	hashes["Streebog-256"] = hex.EncodeToString(gostHasher.Sum(nil))

	ripemdHasher := ripemd.New256()
	ripemdHasher.Write(data)
	hashes["RIPEMD-256"] = hex.EncodeToString(ripemdHasher.Sum(nil))

	sha256Hash := sha256.Sum256(data)
	hashes["SHA-256"] = hex.EncodeToString(sha256Hash[:])

	sm3Hasher := sm3.New()
	sm3Hasher.Write(data)
	hashes["SM3"] = hex.EncodeToString(sm3Hasher.Sum(nil))

	return hashes
}

// formatHashes formats hashes with right-aligned names (EXACTLY from yubisigner.go:1746-1766)
func formatHashes(hashes map[string]string) string {
	keys := make([]string, 0, len(hashes))
	for k := range hashes {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	maxLen := 0
	for _, k := range keys {
		if len(k) > maxLen {
			maxLen = len(k)
		}
	}

	var result strings.Builder
	for _, k := range keys {
		paddedKey := fmt.Sprintf("%*s", maxLen, k)
		result.WriteString(fmt.Sprintf("%s: %s\r\n", paddedKey, hashes[k]))
	}
	return result.String()
}

// formatSignatureRFC formats signature in RFC-compliant format with line breaks (EXACTLY from yubisigner.go:1996-2007)
func formatSignatureRFC(sig string) string {
	var result strings.Builder
	for i := 0; i < len(sig); i += 64 {
		end := i + 64
		if end > len(sig) {
			end = len(sig)
		}
		result.WriteString(sig[i:end])
		result.WriteString("\r\n")
	}
	return result.String()
}

// openYubiKey opens YubiKey at specified index (EXACTLY from yubisigner.go:2073-2092)
func openYubiKey(index int) (*piv.YubiKey, error) {
	cards, err := piv.Cards()
	if err != nil {
		return nil, fmt.Errorf("failed to list cards: %v", err)
	}
	if len(cards) == 0 {
		return nil, fmt.Errorf("no smart card found")
	}

	count := 0
	for _, card := range cards {
		if strings.Contains(strings.ToLower(card), "yubikey") {
			if count == index {
				return piv.Open(card)
			}
			count++
		}
	}
	return nil, fmt.Errorf("no YubiKey found at index %d", index)
}

// signEd25519Data handles Ed25519 signing (EXACTLY from yubisigner.go:1843-1862)
func signEd25519Data(pin string, hash []byte, pubKey ed25519.PublicKey, yk *piv.YubiKey) (string, string, error) {
	auth := piv.KeyAuth{PIN: pin}
	priv, err := yk.PrivateKey(piv.SlotSignature, pubKey, auth)
	if err != nil {
		return "", "", fmt.Errorf("failed to get private key: %v", err)
	}

	signer, ok := priv.(crypto.Signer)
	if !ok {
		return "", "", fmt.Errorf("key does not implement crypto.Signer")
	}

	signature, err := signer.Sign(rand.Reader, hash, crypto.Hash(0))
	if err != nil {
		return "", "", fmt.Errorf("Ed25519 signing failed: %v", err)
	}

	combined := append(pubKey, signature...)
	return hex.EncodeToString(combined), AlgorithmED25519, nil
}

// signDataInternal performs the actual signing operation with YubiKey (EXACTLY from yubisigner.go:1769-1840, Ed25519 only)
func signDataInternal(pin, data []byte) (string, string, error) {
	yk, err := openYubiKey(0)
	if err != nil {
		return "", "", err
	}
	defer yk.Close()

	cert, err := yk.Certificate(piv.SlotSignature)
	if err != nil {
		return "", "", fmt.Errorf("failed to get certificate from signature slot: %v", err)
	}

	ed25519PubKey, ok := cert.PublicKey.(ed25519.PublicKey)
	if !ok {
		return "", "", fmt.Errorf("public key is not Ed25519 (this CLI tool only supports Ed25519)")
	}

	hash := sha256.Sum256(data)
	return signEd25519Data(string(pin), hash[:], ed25519PubKey, yk)
}

func main() {
	var (
		inputFile  string
		author     string
		email      string
		url        string
		telefax    string
		comment    string
		pin        string
		outputFile string
	)

	flag.StringVar(&inputFile, "input", "", "File to sign (required)")
	flag.StringVar(&author, "author", "", "Author name (required)")
	flag.StringVar(&email, "email", "n/a", "Email address (optional)")
	flag.StringVar(&url, "url", "n/a", "URL (optional)")
	flag.StringVar(&telefax, "telefax", "n/a", "Telefax (optional)")
	flag.StringVar(&comment, "comment", "n/a", "Comment (optional)")
	flag.StringVar(&pin, "pin", "", "YubiKey PIN (will prompt if not provided)")
	flag.StringVar(&outputFile, "output", "", "Output signature file (default: <input>.sig)")
	flag.Parse()

	// Validate required fields
	if inputFile == "" {
		fmt.Fprintln(os.Stderr, "Error: --input is required")
		flag.Usage()
		os.Exit(1)
	}
	if author == "" {
		fmt.Fprintln(os.Stderr, "Error: --author is required")
		flag.Usage()
		os.Exit(1)
	}

	// Ensure UTF-8 (EXACTLY from yubisigner.go:1010, 1120)
	author = ensureUTF8(author)
	if comment != "n/a" {
		comment = ensureUTF8(comment)
	}

	// Read input file
	data, err := os.ReadFile(inputFile)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading file: %v\n", err)
		os.Exit(1)
	}

	// Get file info
	fileInfo, err := os.Stat(inputFile)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error getting file info: %v\n", err)
		os.Exit(1)
	}
	fileSize := fileInfo.Size()

	// Calculate hashes (EXACTLY from yubisigner.go:1103)
	fmt.Println("Calculating hashes...")
	hashes := calculateHashesRAM(data)

	// Prepare metadata (EXACTLY from yubisigner.go:1124-1132)
	metadata := fmt.Sprintf("Author: %s\r\n", author)
	metadata += fmt.Sprintf("Signed at: %s\r\n", time.Now().UTC().Format("2006-01-02 15:04:05 +0000"))
	metadata += fmt.Sprintf("Filename: %s\r\n", filepath.Base(inputFile))
	metadata += fmt.Sprintf("File size: %d bytes\r\n", fileSize)
	metadata += fmt.Sprintf("Email: %s\r\n", email)
	metadata += fmt.Sprintf("Telefax: %s\r\n", telefax)
	metadata += fmt.Sprintf("URL: %s\r\n", url)
	metadata += fmt.Sprintf("Comment: %s\r\n", comment)
	metadata += formatHashes(hashes)

	// Get PIN if not provided
	if pin == "" {
		fmt.Print("Enter YubiKey PIN: ")
		fmt.Scanln(&pin)
	}

	// Sign metadata (EXACTLY from yubisigner.go:1138)
	fmt.Println("Signing with YubiKey (touch required)...")
	sig, algo, err := signDataInternal([]byte(pin), []byte(metadata))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Signing failed: %v\n", err)
		os.Exit(1)
	}

	// Build result (EXACTLY from yubisigner.go:1144-1148)
	result := metadata
	result += "-----BEGIN YUBISIGNER " + algo + " SIGNATURE-----\r\n"
	result += formatSignatureRFC(sig)
	result += "-----END YUBISIGNER " + algo + " SIGNATURE-----\r\n"

	// Determine output file
	if outputFile == "" {
		outputFile = inputFile + ".sig"
	}

	// Write signature file
	err = os.WriteFile(outputFile, []byte(result), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error writing signature: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("✓ File signed successfully: %s\n", outputFile)
}