package main

import (
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/composite"
	"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"
)

var (
	newRestoreKeyManager = keymanager.NewKeyManager
	newRestoreDecryptor  = func(registry crypto.Registry) crypto.Decryptor {
		return composite.NewDecryptor(registry)
	}
	restoreOutput io.Writer = os.Stdout
	restoreOsOpen           = os.Open
)

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Restore a backup from a .pqenc file",
	PreRunE: func(cmd *cobra.Command, args []string) error {
		privkeyPq, _ := cmd.Flags().GetString("privkey-pq")
		privkeyClassical, _ := cmd.Flags().GetString("privkey-classical")

		var pqMissing bool
		if _, err := os.Stat(privkeyPq); err != nil {
			pqMissing = true
		}

		var classicalMissing bool
		if _, err := os.Stat(privkeyClassical); err != nil {
			classicalMissing = true
		}

		if pqMissing || classicalMissing {
			return fmt.Errorf(
				"--privkey-pq and --privkey-classical are both required (AND model)",
			)
		}
		return nil
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		inPath, _ := cmd.Flags().GetString("in")
		outPath, _ := cmd.Flags().GetString("out")
		privkeyPq, _ := cmd.Flags().GetString("privkey-pq")
		privkeyClassical, _ := cmd.Flags().GetString("privkey-classical")

		inFile, err := restoreOsOpen(inPath)
		if err != nil {
			return fmt.Errorf("open input file: %w", err)
		}
		defer inFile.Close()

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

		keyManager := newRestoreKeyManager(registry)

		pqPriv, err := keyManager.LoadPriv(privkeyPq, 0x0006)
		if err != nil {
			return fmt.Errorf("load PQ private key: %w", err)
		}

		classicalPriv, err := keyManager.LoadPriv(privkeyClassical, 0x0007)
		if err != nil {
			return fmt.Errorf("load classical private key: %w", err)
		}

		var out io.Writer
		if outPath != "" {
			outFile, err := os.Create(outPath)
			if err != nil {
				return fmt.Errorf("create output file: %w", err)
			}
			defer outFile.Close()
			out = outFile
		} else {
			out = restoreOutput
		}

		decryptor := newRestoreDecryptor(registry)
		if err := decryptor.Decrypt(
			inFile,
			[]crypto.RecipientPriv{pqPriv, classicalPriv},
			out,
		); err != nil {
			return err
		}

		return nil
	},
}

func init() {
	restoreCmd.Flags().String("in", "", "Input .pqenc file path")
	restoreCmd.Flags().String("privkey-pq", "", "Path to PQ private key PEM")
	restoreCmd.Flags().String(
		"privkey-classical",
		"",
		"Path to classical private key PEM",
	)
	restoreCmd.Flags().String("out", "", "Output file path (empty = stdout)")

	_ = restoreCmd.MarkFlagRequired("in")
	_ = restoreCmd.MarkFlagRequired("privkey-pq")
	_ = restoreCmd.MarkFlagRequired("privkey-classical")
}