package main

import (
	"bytes"
	"encoding/pem"
	"os"
	"path/filepath"
	"testing"

	"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 makeRegistry(t *testing.T) crypto.Registry {
	t.Helper()
	reg := crypto.NewRegistry()
	if err := reg.Register(0x0006, func() crypto.KEM { return mlkem768.New() }); err != nil {
		t.Fatalf("register mlkem768: %v", err)
	}
	if err := reg.Register(0x0007, func() crypto.KEM { return x25519.New() }); err != nil {
		t.Fatalf("register x25519: %v", err)
	}
	return reg
}

func TestKeygenCmdFlags(t *testing.T) {
	cmd := newKeygenCmd()

	// Verify flags exist and have correct defaults.
	if cmd.Flag("type") == nil {
		t.Fatal("missing --type flag")
	}
	if cmd.Flag("type").DefValue != "both" {
		t.Errorf("--type default = %q, want both", cmd.Flag("type").DefValue)
	}

	if cmd.Flag("out-prefix") == nil {
		t.Fatal("missing --out-prefix flag")
	}
	if cmd.Flag("out-prefix").DefValue != "" {
		t.Errorf("--out-prefix default = %q, want empty", cmd.Flag("out-prefix").DefValue)
	}

	if cmd.Flag("force") == nil {
		t.Fatal("missing --force flag")
	}
	if cmd.Flag("force").DefValue != "false" {
		t.Errorf("--force default = %q, want false", cmd.Flag("force").DefValue)
	}
}

func TestKeygenCmdInvalidType(t *testing.T) {
	cmd := newKeygenCmd()
	cmd.SetArgs([]string{"--type", "invalid", "--out-prefix", filepath.Join(t.TempDir(), "keys")})

	err := cmd.Execute()
	if err == nil {
		t.Fatal("expected error for invalid --type")
	}
}

func TestKeygenCmdBoth(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	// Verify all 4 files exist.
	files := []string{
		prefix + ".pq.pub.pem",
		prefix + ".pq.priv.pem",
		prefix + ".classical.pub.pem",
		prefix + ".classical.priv.pem",
	}
	for _, f := range files {
		if _, err := os.Stat(f); err != nil {
			t.Errorf("expected file %s to exist: %v", f, err)
		}
	}
}

func TestKeygenCmdPEMTypes(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	tests := []struct {
		path       string
		wantType   string
		wantKeyLen int
	}{
		{prefix + ".pq.pub.pem", "ML-KEM-768 PUBLIC KEY", 1184},
		{prefix + ".pq.priv.pem", "ML-KEM-768 PRIVATE KEY", 64},
		{prefix + ".classical.pub.pem", "X25519 PUBLIC KEY", 32},
		{prefix + ".classical.priv.pem", "X25519 PRIVATE KEY", 32},
	}

	for _, tt := range tests {
		data, err := os.ReadFile(tt.path)
		if err != nil {
			t.Fatalf("read %s: %v", tt.path, err)
		}
		block, _ := pem.Decode(data)
		if block == nil {
			t.Fatalf("failed to decode PEM from %s", tt.path)
		}
		if block.Type != tt.wantType {
			t.Errorf("%s PEM type = %q, want %q", tt.path, block.Type, tt.wantType)
		}
		if len(block.Bytes) != tt.wantKeyLen {
			t.Errorf("%s key len = %d, want %d", tt.path, len(block.Bytes), tt.wantKeyLen)
		}
	}
}

func TestKeygenCmdRoundTrip(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	km := keymanager.NewKeyManager(reg)

	schemes := []struct {
		pubPath  string
		privPath string
		schemeID uint16
		pubLen   int
		privLen  int
	}{
		{prefix + ".pq.pub.pem", prefix + ".pq.priv.pem", 0x0006, 1184, 64},
		{prefix + ".classical.pub.pem", prefix + ".classical.priv.pem", 0x0007, 32, 32},
	}

	for _, s := range schemes {
		pub, err := km.LoadPub(s.pubPath, s.schemeID)
		if err != nil {
			t.Fatalf("LoadPub %s: %v", s.pubPath, err)
		}
		priv, err := km.LoadPriv(s.privPath, s.schemeID)
		if err != nil {
			t.Fatalf("LoadPriv %s: %v", s.privPath, err)
		}

		if pub.SchemeID() != s.schemeID {
			t.Errorf("pub.SchemeID() = 0x%04x, want 0x%04x", pub.SchemeID(), s.schemeID)
		}
		if priv.SchemeID() != s.schemeID {
			t.Errorf("priv.SchemeID() = 0x%04x, want 0x%04x", priv.SchemeID(), s.schemeID)
		}
		if len(pub.Raw()) != s.pubLen {
			t.Errorf("pub.Raw() len = %d, want %d", len(pub.Raw()), s.pubLen)
		}
		if len(priv.Raw()) != s.privLen {
			t.Errorf("priv.Raw() len = %d, want %d", len(priv.Raw()), s.privLen)
		}
		if len(pub.KeyID()) != 8 {
			t.Errorf("pub.KeyID() len = %d, want 8", len(pub.KeyID()))
		}
		if len(priv.KeyID()) != 8 {
			t.Errorf("priv.KeyID() len = %d, want 8", len(priv.KeyID()))
		}

		// Verify the raw bytes round-trip correctly by checking the PEM
		// contents match what LoadPub/LoadPriv return.
		pubPEM, err := os.ReadFile(s.pubPath)
		if err != nil {
			t.Fatalf("read pub PEM: %v", err)
		}
		pubBlock, _ := pem.Decode(pubPEM)
		if pubBlock == nil {
			t.Fatal("failed to decode pub PEM")
		}
		if !bytes.Equal(pub.Raw(), pubBlock.Bytes) {
			t.Errorf("pub.Raw() does not match PEM bytes")
		}

		privPEM, err := os.ReadFile(s.privPath)
		if err != nil {
			t.Fatalf("read priv PEM: %v", err)
		}
		privBlock, _ := pem.Decode(privPEM)
		if privBlock == nil {
			t.Fatal("failed to decode priv PEM")
		}
		if !bytes.Equal(priv.Raw(), privBlock.Bytes) {
			t.Errorf("priv.Raw() does not match PEM bytes")
		}
	}
}

func TestKeygenCmdPQOnly(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "pq", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, err := os.Stat(prefix + ".pq.pub.pem"); err != nil {
		t.Errorf("expected pq.pub.pem to exist: %v", err)
	}
	if _, err := os.Stat(prefix + ".pq.priv.pem"); err != nil {
		t.Errorf("expected pq.priv.pem to exist: %v", err)
	}
	if _, err := os.Stat(prefix + ".classical.pub.pem"); !os.IsNotExist(err) {
		t.Errorf("expected classical.pub.pem to NOT exist")
	}
	if _, err := os.Stat(prefix + ".classical.priv.pem"); !os.IsNotExist(err) {
		t.Errorf("expected classical.priv.pem to NOT exist")
	}
}

func TestKeygenCmdClassicalOnly(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "classical", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, err := os.Stat(prefix + ".classical.pub.pem"); err != nil {
		t.Errorf("expected classical.pub.pem to exist: %v", err)
	}
	if _, err := os.Stat(prefix + ".classical.priv.pem"); err != nil {
		t.Errorf("expected classical.priv.pem to exist: %v", err)
	}
	if _, err := os.Stat(prefix + ".pq.pub.pem"); !os.IsNotExist(err) {
		t.Errorf("expected pq.pub.pem to NOT exist")
	}
	if _, err := os.Stat(prefix + ".pq.priv.pem"); !os.IsNotExist(err) {
		t.Errorf("expected pq.priv.pem to NOT exist")
	}
}

func TestKeygenCmdPermissions(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	pubFiles := []string{prefix + ".pq.pub.pem", prefix + ".classical.pub.pem"}
	for _, f := range pubFiles {
		info, err := os.Stat(f)
		if err != nil {
			t.Fatalf("stat %s: %v", f, err)
		}
		mode := info.Mode().Perm()
		if mode != 0o644 {
			t.Errorf("%s permissions = 0%o, want 0644", f, mode)
		}
	}

	privFiles := []string{prefix + ".pq.priv.pem", prefix + ".classical.priv.pem"}
	for _, f := range privFiles {
		info, err := os.Stat(f)
		if err != nil {
			t.Fatalf("stat %s: %v", f, err)
		}
		mode := info.Mode().Perm()
		if mode != 0o600 {
			t.Errorf("%s permissions = 0%o, want 0600", f, mode)
		}
	}
}

func TestKeygenCmdNoOverwrite(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	// Create an existing file.
	existing := prefix + ".pq.pub.pem"
	if err := os.WriteFile(existing, []byte("existing"), 0o644); err != nil {
		t.Fatalf("write existing file: %v", err)
	}

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	err := cmd.Execute()
	if err == nil {
		t.Fatal("expected error when file exists without --force")
	}

	// Verify existing file was not overwritten.
	data, err := os.ReadFile(existing)
	if err != nil {
		t.Fatalf("read existing file: %v", err)
	}
	if !bytes.Equal(data, []byte("existing")) {
		t.Error("existing file was overwritten without --force")
	}
}

func TestKeygenCmdForceOverwrite(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	// Create an existing file.
	existing := prefix + ".pq.pub.pem"
	if err := os.WriteFile(existing, []byte("existing"), 0o644); err != nil {
		t.Fatalf("write existing file: %v", err)
	}

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix, "--force"})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	// Verify file was overwritten with valid PEM.
	data, err := os.ReadFile(existing)
	if err != nil {
		t.Fatalf("read overwritten file: %v", err)
	}
	block, _ := pem.Decode(data)
	if block == nil {
		t.Fatal("overwritten file is not valid PEM")
	}
	if block.Type != "ML-KEM-768 PUBLIC KEY" {
		t.Errorf("overwritten PEM type = %q, want ML-KEM-768 PUBLIC KEY", block.Type)
	}
}

func TestKeygenCmdKeyIDConsistency(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()
	prefix := filepath.Join(dir, "keys")

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	km := keymanager.NewKeyManager(reg)

	schemes := []struct {
		pubPath  string
		privPath string
		schemeID uint16
		pubLen   int
		privLen  int
	}{
		{prefix + ".pq.pub.pem", prefix + ".pq.priv.pem", 0x0006, 1184, 64},
		{prefix + ".classical.pub.pem", prefix + ".classical.priv.pem", 0x0007, 32, 32},
	}

	for _, s := range schemes {
		pub, err := km.LoadPub(s.pubPath, s.schemeID)
		if err != nil {
			t.Fatalf("LoadPub %s: %v", s.pubPath, err)
		}
		priv, err := km.LoadPriv(s.privPath, s.schemeID)
		if err != nil {
			t.Fatalf("LoadPriv %s: %v", s.privPath, err)
		}

		if len(pub.Raw()) != s.pubLen {
			t.Errorf("pub %s raw len = %d, want %d", s.pubPath, len(pub.Raw()), s.pubLen)
		}
		if len(priv.Raw()) != s.privLen {
			t.Errorf("priv %s raw len = %d, want %d", s.privPath, len(priv.Raw()), s.privLen)
		}

		// KeyID must be present and 8 bytes for both pub and priv.
		if len(pub.KeyID()) != 8 {
			t.Errorf("pub %s KeyID len = %d, want 8", s.pubPath, len(pub.KeyID()))
		}
		if len(priv.KeyID()) != 8 {
			t.Errorf("priv %s KeyID len = %d, want 8", s.privPath, len(priv.KeyID()))
		}
	}
}

func TestKeygenCmdEmptyPrefix(t *testing.T) {
	reg := makeRegistry(t)
	dir := t.TempDir()

	cmd := newKeygenCmdWithDeps(reg)
	cmd.SetArgs([]string{"--type", "pq"})
	cmd.SetOut(nil)
	cmd.SetErr(nil)
	// Change working directory to temp dir so empty prefix creates files there.
	origWd, _ := os.Getwd()
	if err := os.Chdir(dir); err != nil {
		t.Fatalf("Chdir(%q): %v", dir, err)
	}
	defer func() {
		if err := os.Chdir(origWd); err != nil {
			t.Fatalf("Chdir(%q): %v", origWd, err)
		}
	}()

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, err := os.Stat(filepath.Join(dir, ".pq.pub.pem")); err != nil {
		t.Errorf("expected .pq.pub.pem to exist: %v", err)
	}
	if _, err := os.Stat(filepath.Join(dir, ".pq.priv.pem")); err != nil {
		t.Errorf("expected .pq.priv.pem to exist: %v", err)
	}
}