package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"io"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/kami/maven/internal/store"
	"github.com/kami/maven/internal/webauthn"
)

func randBytes(t *testing.T, n int) []byte {
	t.Helper()
	b := make([]byte, n)
	if _, err := io.ReadFull(rand.Reader, b); err != nil {
		t.Fatalf("rand: %v", err)
	}
	b[0] |= 1
	return b
}

func TestDaemonLockStartsLockedAndFlips(t *testing.T) {
	dl := newDaemonLock(true)
	if !dl.isLocked() {
		t.Fatal("newDaemonLock(true) is not locked")
	}
	dl.unlock(nil)
	if dl.isLocked() {
		t.Fatal("still locked after unlock")
	}
	if newDaemonLock(false).isLocked() {
		t.Fatal("newDaemonLock(false) reports locked")
	}
}

// closeStore must be safe on a daemon that never unlocked and safe twice —
// shutdown runs it unconditionally.
func TestDaemonLockCloseStoreIsSafeWhenNeverUnlocked(t *testing.T) {
	dl := newDaemonLock(true)
	if err := dl.closeStore(); err != nil {
		t.Fatalf("closeStore with no store: %v", err)
	}
	if err := dl.closeStore(); err != nil {
		t.Fatalf("second closeStore: %v", err)
	}
}

// The data-loss bug: in locked mode the store is opened on an IPC goroutine
// inside UnlockFn, and shutdown runs on main. Without the handoff nothing
// calls Close, and Close is what re-encrypts the tmpfs working copy back over
// the ciphertext file — so every write of a cold-started session vanished.
func TestDaemonLockSealsTheStoreOpenedAfterUnlock(t *testing.T) {
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "maven.db")
	tmpfs := filepath.Join(dir, "work")
	key := randBytes(t, 32)
	// Store.Close zeroes the key slice it was handed (encState.key is the
	// caller's backing array), so the next boot needs its own copy — exactly
	// as mavend keeps envKeyBytes separate from the config's key.
	nextBoot := bytes.Clone(key)
	ctx := context.Background()

	// Cold start: locked, no store.
	dl := newDaemonLock(true)

	// ... unlock arrives, opens the store and hands it over.
	st, err := store.OpenEncrypted(ctx, dbPath, tmpfs, key)
	if err != nil {
		t.Fatalf("OpenEncrypted: %v", err)
	}
	dl.unlock(st)
	if _, err := st.WriteNote(ctx, time.Now(), "заметка после холодного старта", nil, "test"); err != nil {
		t.Fatalf("WriteNote: %v", err)
	}

	// Shutdown.
	if err := dl.closeStore(); err != nil {
		t.Fatalf("closeStore: %v", err)
	}
	if err := dl.closeStore(); err != nil {
		t.Fatalf("second closeStore after a real store: %v", err)
	}

	// Next boot with the same key must see the write.
	st2, err := store.OpenEncrypted(ctx, dbPath, tmpfs, nextBoot)
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	defer st2.Close()
	notes, err := st2.RecentNotes(ctx, 10)
	if err != nil {
		t.Fatalf("RecentNotes: %v", err)
	}
	if len(notes) != 1 {
		t.Fatalf("got %d notes after a cold-started session, want 1 — the session was lost", len(notes))
	}
}

// The whole point of the wrapped blob: what sits in the state dir must not let
// anyone open the database. Nothing written there may contain the key, and the
// ciphertext must not be readable with a wrong one.
func TestColdStartLeavesNoPlaintextKeyOnDisk(t *testing.T) {
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "maven.db")
	tmpfs := filepath.Join(dir, "work")
	wrappedPath := filepath.Join(dir, "db_key.wrapped")
	key := randBytes(t, 32)
	secret := randBytes(t, 32)
	ctx := context.Background()

	blob, err := webauthn.WrapKey(key, secret)
	if err != nil {
		t.Fatalf("WrapKey: %v", err)
	}
	if err := os.WriteFile(wrappedPath, blob, 0o600); err != nil {
		t.Fatalf("write wrapped key: %v", err)
	}

	st, err := store.OpenEncrypted(ctx, dbPath, tmpfs, key)
	if err != nil {
		t.Fatalf("OpenEncrypted: %v", err)
	}
	if _, err := st.WriteNote(ctx, time.Now(), "секрет", nil, "test"); err != nil {
		t.Fatalf("WriteNote: %v", err)
	}
	if err := st.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	// Walk everything in the state dir; none of it may contain the key.
	err = filepath.Walk(dir, func(p string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		b, rerr := os.ReadFile(p)
		if rerr != nil {
			return nil // unreadable is not a leak
		}
		if bytes.Contains(b, key) {
			t.Errorf("%s contains the plaintext encryption key", p)
		}
		return nil
	})
	if err != nil {
		t.Fatalf("walk: %v", err)
	}

	// The wrapped file must have owner-only permissions.
	fi, err := os.Stat(wrappedPath)
	if err != nil {
		t.Fatalf("stat: %v", err)
	}
	if perm := fi.Mode().Perm(); perm != 0o600 {
		t.Errorf("wrapped key file mode = %o, want 600", perm)
	}

	// A wrong passkey must not open the store.
	if _, _, err := webauthn.UnwrapKey(blob, randBytes(t, 32)); err == nil {
		t.Fatal("a wrong PRF secret unwrapped the key")
	}
	if _, err := store.OpenEncrypted(ctx, dbPath, filepath.Join(dir, "work2"), randBytes(t, 32)); err == nil {
		t.Fatal("the encrypted store opened under a wrong key")
	}

	// And the right one round-trips back to a readable database.
	got, version, err := webauthn.UnwrapKey(blob, secret)
	if err != nil {
		t.Fatalf("UnwrapKey: %v", err)
	}
	if version != webauthn.BlobV2 {
		t.Errorf("blob version = %v, want v2", version)
	}
	st2, err := store.OpenEncrypted(ctx, dbPath, tmpfs, got)
	if err != nil {
		t.Fatalf("reopen with the unwrapped key: %v", err)
	}
	defer st2.Close()
	notes, err := st2.RecentNotes(ctx, 10)
	if err != nil {
		t.Fatalf("RecentNotes: %v", err)
	}
	if len(notes) != 1 {
		t.Fatalf("got %d notes, want 1", len(notes))
	}
}