package main

import (
"encoding/hex"
"flag"
"fmt"
"os"

"himbocrypt/pkg/engine"
)

func main() {
	if len(os.Args) < 2 {
		printUsage()
		os.Exit(1)
	}

	command := os.Args[1]
	e := engine.NewEngine()

	switch command {
	case "keygen":
		handleKeygen(e)
	case "encrypt":
		handleEncrypt(e)
	case "decrypt":
		handleDecrypt(e)
	default:
		fmt.Printf("Unknown command: %s\n", command)
		printUsage()
		os.Exit(1)
	}
}

func printUsage() {
	fmt.Println("Usage:")
	fmt.Println("  himbocrypt keygen")
	fmt.Println("  himbocrypt encrypt -sender-priv <key> -recipient-pub <key> -msg <message>")
	fmt.Println("  himbocrypt decrypt -recipient-priv <key> -sender-pub <key> -ciphertext <hex>")
}

func handleKeygen(e *engine.Engine) {
	kp, err := e.GenerateKeyPair()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error generating keys: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("Generated KeyPair:")
	fmt.Printf("Private Key: %s\n", kp.EncodePrivateKey())
	fmt.Printf("Public Key:  %s\n", kp.EncodePublicKey())
}

func handleEncrypt(e *engine.Engine) {
	encryptCmd := flag.NewFlagSet("encrypt", flag.ExitOnError)
	senderPrivHex := encryptCmd.String("sender-priv", "", "Sender's Private Key (Hex)")
	recipientPubHex := encryptCmd.String("recipient-pub", "", "Recipient's Public Key (Hex)")
	message := encryptCmd.String("msg", "", "Message to encrypt")

	encryptCmd.Parse(os.Args[2:])

	if *senderPrivHex == "" || *recipientPubHex == "" || *message == "" {
		encryptCmd.Usage()
		os.Exit(1)
	}

	senderPriv, err := e.DecodePrivateKey(*senderPrivHex)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid sender private key: %v\n", err)
		os.Exit(1)
	}

	recipientPub, err := e.DecodePublicKey(*recipientPubHex)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid recipient public key: %v\n", err)
		os.Exit(1)
	}

	ciphertext, err := e.Encrypt(senderPriv, recipientPub, []byte(*message))
	if err != nil {
		fmt.Fprintf(os.Stderr, "Encryption failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Ciphertext: %s\n", hex.EncodeToString(ciphertext))
}

func handleDecrypt(e *engine.Engine) {
	decryptCmd := flag.NewFlagSet("decrypt", flag.ExitOnError)
	recipientPrivHex := decryptCmd.String("recipient-priv", "", "Recipient's Private Key (Hex)")
	senderPubHex := decryptCmd.String("sender-pub", "", "Sender's Public Key (Hex)")
	ciphertextHex := decryptCmd.String("ciphertext", "", "Ciphertext (Hex)")

	decryptCmd.Parse(os.Args[2:])

	if *recipientPrivHex == "" || *senderPubHex == "" || *ciphertextHex == "" {
		decryptCmd.Usage()
		os.Exit(1)
	}

	recipientPriv, err := e.DecodePrivateKey(*recipientPrivHex)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid recipient private key: %v\n", err)
		os.Exit(1)
	}

	senderPub, err := e.DecodePublicKey(*senderPubHex)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid sender public key: %v\n", err)
		os.Exit(1)
	}

	ciphertextBytes, err := hex.DecodeString(*ciphertextHex)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Invalid ciphertext hex: %v\n", err)
		os.Exit(1)
	}

	plaintext, err := e.Decrypt(recipientPriv, senderPub, ciphertextBytes)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Decryption failed: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("Plaintext: %s\n", string(plaintext))
}
