package sshmanager

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
)

func MarshalED25519PublicKey(pub ed25519.PublicKey) (string, error) {
	sshPubKey, err := ssh.NewPublicKey(pub)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPubKey))), nil
}

func RegeneratePublicKeyFromPrivateKeyFile(privPath, label string) (string, string, error) {
	sshPubKey, err := PublicKeyFromPrivateKeyFile(privPath)
	if err != nil {
		return "", "", err
	}
	pubKey := sshPubKey + " " + label
	fp, _ := Fingerprint(pubKey)
	return pubKey, fp, nil
}

func PublicKeyFromPrivateKeyFile(privPath string) (string, error) {
	data, err := os.ReadFile(privPath)
	if err != nil {
		return "", err
	}
	block, _ := pem.Decode(data)
	if block == nil {
		return "", fmt.Errorf("no PEM block in private key file")
	}
	signer, err := ssh.ParsePrivateKey(data)
	if err != nil {
		return "", fmt.Errorf("parsing private key: %w", err)
	}
	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(signer.PublicKey()))), nil
}

func Fingerprint(publicKey string) (string, error) {
	pubKey := strings.TrimSpace(publicKey)
	parts := strings.Fields(pubKey)
	if len(parts) < 2 {
		return "", fmt.Errorf("invalid public key format")
	}
	keyData, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return "", fmt.Errorf("decoding public key: %w", err)
	}
	if len(keyData) == ed25519.PublicKeySize {
		h := sha256.Sum256(keyData)
		return "SHA256:" + base64.RawStdEncoding.EncodeToString(h[:]), nil
	}
	h := sha256.Sum256(keyData)
	return "SHA256:" + base64.RawStdEncoding.EncodeToString(h[:]), nil
}

func ParsePublicKey(data []byte) ([]byte, string, error) {
	block, _ := pem.Decode(data)
	if block == nil {
		return nil, "", fmt.Errorf("no PEM block found")
	}
	var pubKey []byte
	var err error
	switch block.Type {
	case "PUBLIC KEY":
		pubKey = block.Bytes
	case "OPENSSH KEY":
		parts := strings.Fields(string(block.Bytes))
		if len(parts) < 2 {
			return nil, "", fmt.Errorf("invalid openssh key format")
		}
		pubKey, err = base64.StdEncoding.DecodeString(parts[1])
		if err != nil {
			return nil, "", err
		}
	default:
		return nil, "", fmt.Errorf("unknown PEM type: %s", block.Type)
	}
	h := sha256.Sum256(pubKey)
	return pubKey, "SHA256:" + base64.RawStdEncoding.EncodeToString(h[:]), nil
}

func GenerateKeyPair(label string, sshDir string) (privPath, pubPath, pubKey, fingerprint string, err error) {
	if err := os.MkdirAll(sshDir, 0700); err != nil {
		return "", "", "", "", fmt.Errorf("creating ssh dir: %w", err)
	}
	privPath = filepath.Join(sshDir, label+".key")
	pubPath = privPath + ".pub"
	if _, err := os.Stat(privPath); err == nil {
		return "", "", "", "", fmt.Errorf("key already exists")
	}
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return "", "", "", "", fmt.Errorf("generating ed25519 key: %w", err)
	}
	privFile, err := os.OpenFile(privPath, os.O_CREATE|os.O_WRONLY, 0600)
	if err != nil {
		return "", "", "", "", fmt.Errorf("creating private key file: %w", err)
	}
	defer privFile.Close()
	privBytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return "", "", "", "", fmt.Errorf("marshaling private key: %w", err)
	}
	pem.Encode(privFile, &pem.Block{Type: "PRIVATE KEY", Bytes: privBytes})
	sshPubKey, err := ssh.NewPublicKey(pub)
	if err != nil {
		return "", "", "", "", fmt.Errorf("wrapping public key: %w", err)
	}
	pubKey = strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPubKey))) + " " + label
	if err := os.WriteFile(pubPath, []byte(pubKey), 0644); err != nil {
		return "", "", "", "", fmt.Errorf("writing public key: %w", err)
	}
	fp, _ := Fingerprint(pubKey)
	return privPath, pubPath, pubKey, fp, nil
}