package main

import (
	"crypto/rand"
	"fmt"
	"os"

	"github.com/spf13/cobra"

	"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/keymanager"
	"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/mlkem768"
	"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/x25519"
	"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)

func newRegistry() crypto.Registry {
	reg := crypto.NewRegistry()
	_ = reg.Register(0x0006, func() crypto.KEM { return mlkem768.New() })
	_ = reg.Register(0x0007, func() crypto.KEM { return x25519.New() })
	return reg
}

func newKeygenCmd() *cobra.Command {
	return newKeygenCmdWithDeps(newRegistry())
}

func newKeygenCmdWithDeps(reg crypto.Registry) *cobra.Command {
	var (
		keyType   string
		outPrefix string
		force     bool
	)

	cmd := &cobra.Command{
		Use:   "keygen",
		Short: "Generate encryption key pairs",
		RunE: func(cmd *cobra.Command, args []string) error {
			schemes := []struct {
				name     string
				schemeID uint16
			}{}

			switch keyType {
			case "pq":
				schemes = append(schemes, struct {
					name     string
					schemeID uint16
				}{"pq", 0x0006})
			case "classical":
				schemes = append(schemes, struct {
					name     string
					schemeID uint16
				}{"classical", 0x0007})
			case "both":
				schemes = append(
					schemes,
					struct {
						name     string
						schemeID uint16
					}{"pq", 0x0006},
					struct {
						name     string
						schemeID uint16
					}{"classical", 0x0007},
				)
			default:
				return fmt.Errorf("invalid --type %q; must be pq, classical, or both", keyType)
			}

			if !force {
				for _, s := range schemes {
					pubPath := outPrefix + "." + s.name + ".pub.pem"
					privPath := outPrefix + "." + s.name + ".priv.pem"
					for _, p := range []string{pubPath, privPath} {
						if _, err := os.Stat(p); err == nil {
							return fmt.Errorf("file already exists: %s (use --force to overwrite)", p)
						}
					}
				}
			}

			km := keymanager.NewKeyManager(reg)

			for _, s := range schemes {
				pubPath := outPrefix + "." + s.name + ".pub.pem"
				privPath := outPrefix + "." + s.name + ".priv.pem"

				pubFile, err := os.OpenFile(pubPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
				if err != nil {
					return fmt.Errorf("create public key file %s: %w", pubPath, err)
				}

				privFile, err := os.OpenFile(privPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
				if err != nil {
					_ = pubFile.Close()
					return fmt.Errorf("create private key file %s: %w", privPath, err)
				}

				if err := km.Generate(s.schemeID, pubFile, privFile, rand.Reader); err != nil {
					_ = pubFile.Close()
					_ = privFile.Close()
					return fmt.Errorf("generate %s keys: %w", s.name, err)
				}

				if err := pubFile.Close(); err != nil {
					return fmt.Errorf("close public key file %s: %w", pubPath, err)
				}
				if err := privFile.Close(); err != nil {
					return fmt.Errorf("close private key file %s: %w", privPath, err)
				}
				if err := os.Chmod(pubPath, 0o644); err != nil {
					return fmt.Errorf("chmod public key file %s: %w", pubPath, err)
				}
				if err := os.Chmod(privPath, 0o600); err != nil {
					return fmt.Errorf("chmod private key file %s: %w", privPath, err)
				}
			}

			return nil
		},
	}

	cmd.Flags().StringVar(&keyType, "type", "both", "Key type to generate (pq|classical|both)")
	cmd.Flags().StringVar(&outPrefix, "out-prefix", "", "Output file path prefix")
	cmd.Flags().BoolVar(&force, "force", false, "Overwrite existing files")

	return cmd
}