package authz

import (
	"bytes"
	"encoding/json"
	"errors"
	"io/fs"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// MemoryNonceStore is a non-durable NonceStore for tests. Replay protection does
// NOT survive process restart — never use it on a real host.
type MemoryNonceStore struct {
	mu   sync.Mutex
	seen map[string]time.Time
}

// NewMemoryNonceStore builds an empty in-memory store.
func NewMemoryNonceStore() *MemoryNonceStore {
	return &MemoryNonceStore{seen: make(map[string]time.Time)}
}

// SeenOrRecord reports whether nonce was already recorded, recording it if not.
func (m *MemoryNonceStore) SeenOrRecord(nonce string, exp time.Time) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.seen[nonce]; ok {
		return true
	}
	m.seen[nonce] = exp
	return false
}

// FileNonceStore is the durable, crash-safe NonceStore for the host. Mechanism:
// an fsync'd append-only JSONL log with an in-memory index, periodic compaction,
// and expiry-only pruning.
//
// Durability guarantee: a nonce is on disk AND fsync'd before SeenOrRecord returns
// false, so the caller acting on a verified op always does so AFTER the durable
// record. A crash between verify and execute therefore drops the op (fail-safe
// direction) and never enables a replay. Replay protection survives restarts: the
// log is replayed into the index on Open.
//
// Pruning: a nonce is dropped only after its exp (compaction), never before —
// pruning before expiry would reopen the replay window. (An expired nonce can't be
// replayed anyway: the time-window check rejects an expired op before the nonce
// check, so pruning is housekeeping, not an authz hole.)
//
// Concurrency: a single mutex guards the file handle and index (single-process; the
// agent is concurrent — 03 §10).
type FileNonceStore struct {
	mu           sync.Mutex
	path         string
	f            *os.File
	idx          map[string]time.Time
	sinceCompact int
	now          func() time.Time

	// CompactEvery is the append count that triggers a compaction (default 1000).
	CompactEvery int
}

type nonceRecord struct {
	Nonce string    `json:"n"`
	Exp   time.Time `json:"e"`
}

// OpenFileNonceStore opens (or creates) the durable store at path, replaying any
// existing log into the index.
func OpenFileNonceStore(path string) (*FileNonceStore, error) {
	s := &FileNonceStore{
		path:         path,
		idx:          make(map[string]time.Time),
		now:          func() time.Time { return time.Now().UTC() },
		CompactEvery: 1000,
	}
	if err := s.load(); err != nil {
		return nil, err
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		return nil, err
	}
	s.f = f
	syncDir(filepath.Dir(path)) // make a freshly-created file's dir entry durable
	return s, nil
}

func (s *FileNonceStore) load() error {
	b, err := os.ReadFile(s.path)
	if errors.Is(err, fs.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	for _, line := range bytes.Split(b, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		var r nonceRecord
		if json.Unmarshal(line, &r) != nil {
			continue // skip a torn trailing line from a crash mid-append
		}
		s.idx[r.Nonce] = r.Exp
	}
	return nil
}

// SeenOrRecord durably records an unseen nonce before returning false. On any I/O
// failure it returns true (fail-safe: the op is NOT executed rather than risk an
// unrecorded nonce enabling a later replay).
func (s *FileNonceStore) SeenOrRecord(nonce string, exp time.Time) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.idx[nonce]; ok {
		return true
	}
	rec, _ := json.Marshal(nonceRecord{Nonce: nonce, Exp: exp})
	rec = append(rec, '\n')
	if _, err := s.f.Write(rec); err != nil {
		return true
	}
	if err := s.f.Sync(); err != nil {
		return true
	}
	s.idx[nonce] = exp
	s.sinceCompact++
	s.maybeCompact()
	return false
}

// Close releases the file handle.
func (s *FileNonceStore) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.f != nil {
		return s.f.Close()
	}
	return nil
}

// maybeCompact rewrites the log keeping only non-expired entries once enough
// appends have accrued. Caller holds the mutex. Compaction is housekeeping: the
// recorded nonce is already durable, so a compaction failure never fails the op.
func (s *FileNonceStore) maybeCompact() {
	if s.CompactEvery <= 0 || s.sinceCompact < s.CompactEvery {
		return
	}
	s.sinceCompact = 0
	now := s.now()

	live := make(map[string]time.Time, len(s.idx))
	var buf bytes.Buffer
	for n, e := range s.idx {
		if e.Before(now) {
			continue // prune AFTER expiry only — safe
		}
		live[n] = e
		rec, _ := json.Marshal(nonceRecord{Nonce: n, Exp: e})
		buf.Write(rec)
		buf.WriteByte('\n')
	}

	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil {
		return // keep using the existing handle; nonce already durable
	}
	if tf, err := os.OpenFile(tmp, os.O_WRONLY, 0o600); err == nil {
		_ = tf.Sync()
		_ = tf.Close()
	}
	if s.f != nil {
		_ = s.f.Close()
	}
	if err := os.Rename(tmp, s.path); err != nil {
		s.f, _ = os.OpenFile(s.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
		return
	}
	syncDir(filepath.Dir(s.path))
	s.f, _ = os.OpenFile(s.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	s.idx = live
}

// syncDir best-effort fsyncs a directory so a create/rename is durable.
func syncDir(dir string) {
	if d, err := os.Open(dir); err == nil {
		_ = d.Sync()
		_ = d.Close()
	}
}