package handoff

import (
	"bytes"
	"context"
	"crypto/rand"
	"database/sql"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"amadan.net/rastrillo/rastrillo/assertion"
	"amadan.net/rastrillo/rastrillo/crypto"
	_ "modernc.org/sqlite"
)

type testSealer struct{ key []byte }

func (s testSealer) Seal(p string, b []byte) ([]byte, error) {
	return crypto.SealSym(crypto.Derive(s.key, p), b)
}
func (s testSealer) Open(p string, b []byte) ([]byte, error) {
	return crypto.OpenSym(crypto.Derive(s.key, p), b)
}

type fixture struct {
	s      *Store
	k      *crypto.Keypair
	now    time.Time
	dbpath string
	db     *sql.DB
}

func fresh(t *testing.T) fixture {
	t.Helper()
	path := filepath.Join(t.TempDir(), "handoff.db")
	db, e := sql.Open("sqlite", path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)")
	if e != nil {
		t.Fatal(e)
	}
	db.SetMaxOpenConns(4)
	t.Cleanup(func() { db.Close() })
	schema, e := Migrations.ReadFile("migrations/0001_exchange.sql")
	if e != nil {
		t.Fatal(e)
	}
	if _, e = db.Exec(string(schema)); e != nil {
		t.Fatal(e)
	}
	key := make([]byte, 32)
	rand.Read(key)
	k, e := crypto.Generate()
	if e != nil {
		t.Fatal(e)
	}
	s, e := NewStore(db, testSealer{key}, "https://id.example.com", "https://app.example.com", map[string][]byte{"k1": k.SignPub()})
	if e != nil {
		t.Fatal(e)
	}
	return fixture{s: s, k: k, now: time.Unix(1700000000, 0), dbpath: path, db: db}
}
func (f fixture) token(t *testing.T, nonce string) string {
	t.Helper()
	id, e := assertion.RandomID()
	if e != nil {
		t.Fatal(e)
	}
	s, e := assertion.Sign(f.k, assertion.Claims{Version: 1, Issuer: "https://id.example.com", Audience: "https://app.example.com", Subject: "person_123", KeyID: "k1", ID: id, Nonce: nonce, IssuedAt: f.now.Unix(), ExpiresAt: f.now.Unix() + 120, AuthTime: f.now.Unix() - 900})
	if e != nil {
		t.Fatal(e)
	}
	return s
}
func (f fixture) pending(t *testing.T) (string, string) {
	t.Helper()
	secret, nonce, e := f.s.Start(context.Background(), "/documents/private-canary?team=blue", "", f.now)
	if e != nil {
		t.Fatal(e)
	}
	return secret, nonce
}
func (f fixture) received(t *testing.T, nonce string) (string, string) {
	t.Helper()
	token := f.token(t, nonce)
	handle, e := f.s.Receive(context.Background(), token, f.now)
	if e != nil {
		t.Fatal(e)
	}
	return handle, token
}
func count(t *testing.T, db *sql.DB, table string) int {
	t.Helper()
	var n int
	if e := db.QueryRow("SELECT COUNT(*) FROM " + table).Scan(&n); e != nil {
		t.Fatal(e)
	}
	return n
}
func failFinish(t *testing.T, f fixture, secret, handle, subject string, now time.Time, want error) {
	t.Helper()
	c, p, e := f.s.Finish(context.Background(), secret, handle, subject, now)
	if e != want || p != "" || !reflect.DeepEqual(c, assertion.Claims{}) {
		t.Fatalf("finish refused: %v wanted %v zero=%v", e, want, reflect.DeepEqual(c, assertion.Claims{}))
	}
}

func TestExchangeAndReplay(t *testing.T) {
	f := fresh(t)
	secret, nonce := f.pending(t)
	handle, token := f.received(t, nonce)
	c, path, e := f.s.Finish(context.Background(), secret, handle, "", f.now)
	if e != nil || c.Subject != "person_123" || c.AuthTime != f.now.Unix()-900 || path != "/documents/private-canary?team=blue" {
		t.Fatalf("finish: %v", e)
	}
	if count(t, f.db, "oficina_handoff_pending") != 0 || count(t, f.db, "oficina_handoff_inbox") != 0 || count(t, f.db, "oficina_handoff_replay") != 1 {
		t.Fatal("consumption not atomic")
	}
	failFinish(t, f, secret, handle, "", f.now, ErrRefused)
	if h, e := f.s.Receive(context.Background(), token, f.now); e != ErrRefused || h != "" {
		t.Fatal("replayed callback accepted")
	}
	if count(t, f.db, "oficina_handoff_inbox") != 0 {
		t.Fatal("replay allocated inbox")
	}
}
func TestBrowserAndSubjectBinding(t *testing.T) {
	f := fresh(t)
	alice, an := f.pending(t)
	bob, bn := f.pending(t)
	ah, _ := f.received(t, an)
	bh, _ := f.received(t, bn)
	failFinish(t, f, alice, bh, "", f.now, ErrRefused)
	failFinish(t, f, bob, ah, "", f.now, ErrRefused)
	failFinish(t, f, alice, ah, "someone_else", f.now, ErrRefused)
	if _, _, e := f.s.Finish(context.Background(), alice, ah, "person_123", f.now); e != nil {
		t.Fatal("failed attempts consumed legitimate request", e)
	}
}
func TestConcurrentFinishAndCallback(t *testing.T) {
	f := fresh(t)
	secret, nonce := f.pending(t)
	token := f.token(t, nonce)
	var mu sync.Mutex
	var handles []string
	var wg sync.WaitGroup
	for range 12 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			h, e := f.s.Receive(context.Background(), token, f.now)
			if e == nil {
				mu.Lock()
				handles = append(handles, h)
				mu.Unlock()
			} else if e != ErrRefused {
				t.Error(e)
			}
		}()
	}
	wg.Wait()
	if len(handles) != 1 {
		t.Fatalf("callback winners %d", len(handles))
	}
	wins := 0
	for range 12 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, _, e := f.s.Finish(context.Background(), secret, handles[0], "", f.now)
			if e == nil {
				mu.Lock()
				wins++
				mu.Unlock()
			} else if e != ErrRefused {
				t.Error(e)
			}
		}()
	}
	wg.Wait()
	if wins != 1 {
		t.Fatalf("finish winners %d", wins)
	}
}
func TestReplacementAndExpiry(t *testing.T) {
	f := fresh(t)
	old, nonce := f.pending(t)
	handle, _ := f.received(t, nonce)
	if _, _, e := f.s.Start(context.Background(), "/", old, f.now); e != nil {
		t.Fatal(e)
	}
	failFinish(t, f, old, handle, "", f.now, ErrRefused)
	secret, nonce := f.pending(t)
	handle, _ = f.received(t, nonce)
	failFinish(t, f, secret, handle, "", f.now.Add(120*time.Second), ErrRefused)
	if e := f.s.Sweep(context.Background(), f.now.Add(601*time.Second)); e != nil {
		t.Fatal(e)
	}
	if count(t, f.db, "oficina_handoff_pending") != 0 || count(t, f.db, "oficina_handoff_inbox") != 0 {
		t.Fatal("expired rows survive")
	}
}
func TestClockFenceAndRestart(t *testing.T) {
	f := fresh(t)
	secret, nonce := f.pending(t)
	handle, token := f.received(t, nonce)
	if _, _, e := f.s.Finish(context.Background(), secret, handle, "", f.now); e != nil {
		t.Fatal(e)
	}
	if e := f.s.Sweep(context.Background(), f.now.Add(150*time.Second)); e != nil {
		t.Fatal(e)
	}
	if count(t, f.db, "oficina_handoff_replay") != 1 {
		t.Fatal("replay pruned at retention boundary")
	}
	if e := f.s.Sweep(context.Background(), f.now.Add(151*time.Second)); e != nil {
		t.Fatal(e)
	}
	if count(t, f.db, "oficina_handoff_replay") != 0 {
		t.Fatal("old replay not pruned")
	}
	if e := f.db.Close(); e != nil {
		t.Fatal(e)
	}
	db, e := sql.Open("sqlite", f.dbpath)
	if e != nil {
		t.Fatal(e)
	}
	defer db.Close()
	f.s.db = db
	if e = f.s.Sweep(context.Background(), f.now.Add(121*time.Second)); e != nil {
		t.Fatal("tolerated rollback refused", e)
	}
	if _, e = f.s.Receive(context.Background(), token, f.now.Add(121*time.Second)); e != ErrRefused {
		t.Fatal("pruned token reappeared", e)
	}
	if e = f.s.Sweep(context.Background(), f.now.Add(120*time.Second)); e != ErrClock {
		t.Fatal("large rollback not fenced", e)
	}
	if replayRetention < clockTolerance {
		t.Fatal("clock rollback exceeds retained history")
	}
}
func TestStorageFailureDoesNotReturnIdentity(t *testing.T) {
	f := fresh(t)
	secret, nonce := f.pending(t)
	handle, _ := f.received(t, nonce)
	if _, e := f.db.Exec(`CREATE TRIGGER break_delete BEFORE DELETE ON oficina_handoff_pending BEGIN SELECT RAISE(ABORT,'synthetic failure'); END`); e != nil {
		t.Fatal(e)
	}
	failFinish(t, f, secret, handle, "", f.now, ErrStorage)
	if count(t, f.db, "oficina_handoff_replay") != 0 || count(t, f.db, "oficina_handoff_pending") != 1 || count(t, f.db, "oficina_handoff_inbox") != 1 {
		t.Fatal("partial transaction survived")
	}
}
func TestPrivateValuesNeverReachDBOrWAL(t *testing.T) {
	f := fresh(t)
	secret, nonce := f.pending(t)
	handle, token := f.received(t, nonce)
	for _, path := range []string{f.dbpath, f.dbpath + "-wal"} {
		b, e := os.ReadFile(path)
		if e != nil {
			t.Fatal(e)
		}
		for _, value := range []string{secret, nonce, handle, token, "private-canary", "?team=blue", "person_123"} {
			if bytes.Contains(b, []byte(value)) {
				t.Fatalf("plaintext leaked in %s", filepath.Base(path))
			}
		}
	}
}
func TestBounds(t *testing.T) {
	f := fresh(t)
	tx, e := f.db.Begin()
	if e != nil {
		t.Fatal(e)
	}
	for i := 0; i < rowLimit; i++ {
		if _, e = tx.Exec(`INSERT INTO oficina_handoff_pending(secret_hash,nonce_hash,sealed_path,created_at,expires_at) VALUES(printf('%d',?),printf('%d',?),x'00',?,?)`, i, i, f.now.Unix()-1, f.now.Unix()+600); e != nil {
			t.Fatal(e)
		}
	}
	if e = tx.Commit(); e != nil {
		t.Fatal(e)
	}
	_, nonce := f.pending(t)
	if count(t, f.db, "oficina_handoff_pending") != rowLimit {
		t.Fatal("pending cap not bounded")
	}
	tx, e = f.db.Begin()
	if e != nil {
		t.Fatal(e)
	}
	for i := 0; i < rowLimit; i++ {
		if _, e = tx.Exec(`INSERT INTO oficina_handoff_inbox(handle_hash,jti_hash,sealed_token,created_at,expires_at) VALUES(printf('%d',?),printf('%d',?),x'00',?,?)`, i, i, f.now.Unix(), f.now.Unix()+120); e != nil {
			t.Fatal(e)
		}
	}
	if e = tx.Commit(); e != nil {
		t.Fatal(e)
	}
	if h, e := f.s.Receive(context.Background(), f.token(t, nonce), f.now); e != ErrFull || h != "" {
		t.Fatal("full inbox admitted", e)
	}
}
func TestSafePathsAndNonce(t *testing.T) {
	for _, s := range []string{"/", "/documents/a?team=1", "/sheets?q=hello%20world"} {
		if !SafePath(s) {
			t.Fatal(s)
		}
	}
	for _, s := range []string{"", "//evil.example/", "https://evil.example/", "/\\evil", "/%2f%2fevil", "/%2e%2e/auth/start", "/auth/start", "/%61uth/start", "/signin", "/authorize/resume", "/a\n", "/a#fragment", strings.Repeat("a", 2049)} {
		if SafePath(s) {
			t.Fatal(s)
		}
	}
	for _, s := range []string{"", "A", "AAAAAAAAAAAAAAAAAAAAAA\n", "AAAAAAAAAAAAAAAAAAAAAB"} {
		if ValidNonce(s) {
			t.Fatal("invalid nonce accepted")
		}
	}
}