package store

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"sort"
	"time"

	"git.smesh.lol/torly/pkg/nostr"
)

// This file owns the social graph indexes and the _graph traversal.
//
// Every edge is stored twice, once in each direction:
//
//	epg event -> pubkey       peg pubkey -> event
//	eeg event -> event        gee event -> event
//	ppg pubkey -> pubkey      gpp pubkey -> pubkey
//
// The forward direction is what a traversal scans; the reverse direction is
// what makes deletion and inbound traversal one prefix scan instead of a full
// table scan.
//
// Pubkeys are interned in pks/spk first, so an edge carries two five-byte
// serials rather than two 32-byte keys.

// ErrGraph is returned for a malformed graph query.
var ErrGraph = errors.New("store: invalid graph query")

// lookupPubkeySerial resolves a pubkey without allocating a serial.
func (e *Engine) lookupPubkeySerial(pub []byte) (uint64, bool, error) {
	if len(pub) < 32 {
		return 0, false, errors.New("store: short pubkey")
	}
	h := Hash8(pub)
	prefix := makePksPrefix(h)
	// The common case is a direct point lookup. Hash collisions are possible
	// but vanishingly rare, so only the fallback scans every row for the hash.
	if rec, ok, err := e.pks.Get(prefix); err != nil {
		return 0, false, err
	} else if ok {
		ser := GetSerial(rec[3+HashLen:])
		if row, exists, err := e.spk.Get(makeSpk(ser)); err != nil {
			return 0, false, err
		} else if exists && len(row) >= spkRecLen+spkValLen &&
			bytes.Equal(row[spkRecLen:spkRecLen+spkValLen], pub[:32]) {
			return ser, true, nil
		}
	}

	var found uint64
	var ok bool
	var scanErr error
	err := e.pks.Scan(prefix, prefixEnd(prefix), func(rec []byte) bool {
		ser := GetSerial(rec[3+HashLen:])
		row, exists, err := e.spk.Get(makeSpk(ser))
		if err != nil {
			scanErr = err
			return false
		}
		if exists && len(row) >= spkRecLen+spkValLen &&
			bytes.Equal(row[spkRecLen:spkRecLen+spkValLen], pub[:32]) {
			found = ser
			ok = true
			return false
		}
		return true
	})
	if err != nil {
		return 0, false, err
	}
	if scanErr != nil {
		return 0, false, scanErr
	}
	return found, ok, nil
}

// pubkeySerial interns a pubkey, allocating a serial the first time it is seen.
func (e *Engine) pubkeySerial(pub []byte) (uint64, error) {
	if e.bulkPub != nil {
		if ser, ok := e.bulkPub[string(pub[:32])]; ok {
			return ser, nil
		}
	}
	if ser, ok, err := e.lookupPubkeySerial(pub); err != nil {
		return 0, err
	} else if ok {
		return ser, nil
	}
	ser := e.nextPubSer
	if ser == 0 || ser > MaxSerial {
		return 0, errors.New("store: pubkey serial space exhausted")
	}
	e.nextPubSer++
	// The value is written before the lookup row: a crash between the two
	// leaves an unreachable pubkey value, never a lookup row that resolves to
	// a missing key.
	if err := e.spk.Put(makeSpkRec(ser, pub)); err != nil {
		return 0, err
	}
	if err := e.pks.Put(makePks(Hash8(pub), ser)); err != nil {
		return 0, err
	}
	if e.bulkPub != nil {
		e.bulkPub[string(pub[:32])] = ser
	}
	return ser, nil
}

func (e *Engine) pubkeyBySerial(ser uint64) ([32]byte, bool) {
	var out [32]byte
	row, ok, err := e.spk.Get(makeSpk(ser))
	if err != nil || !ok || len(row) < spkRecLen+spkValLen {
		return out, false
	}
	copy(out[:], row[spkRecLen:spkRecLen+spkValLen])
	return out, true
}

func (e *Engine) eventIDBySerial(ser uint64) ([32]byte, bool) {
	var out [32]byte
	row, ok, err := e.sei.Get(makeSei(ser))
	if err != nil || !ok || len(row) < seiRecLen+seiValLen {
		return out, false
	}
	copy(out[:], row[seiRecLen:seiRecLen+seiValLen])
	return out, true
}

// A single event can carry an unbounded tag list. Such events are not graph
// indexed at all: partial truncation would silently change the meaning of a
// follow list or thread, while indexing them makes one event dominate a graph
// rebuild. The event itself is still stored and available to NIP-01 queries.
const maxGraphTags = 256

func graphTagCount(ev *nostr.Event) int {
	for i := range ev.Tags {
		if i >= maxGraphTags {
			return i + 1
		}
	}
	return len(ev.Tags)
}

// indexGraph writes the author, p-tag and e-tag edges for one event. An event
// with more than maxGraphTags total tags contributes no graph edges at all.
func (e *Engine) indexGraph(ev *nostr.Event, ser uint64) error {
	if graphTagCount(ev) > maxGraphTags {
		return nil
	}
	authorSer, err := e.pubkeySerial(ev.PubKey[:])
	if err != nil {
		return err
	}
	if err := e.epg.Put(makeEpg(ser, authorSer, ev.Kind, DirAuthor)); err != nil {
		return err
	}
	if err := e.peg.Put(makePeg(authorSer, ev.Kind, DirAuthor, ser)); err != nil {
		return err
	}

	seenPK := map[uint64]struct{}{authorSer: {}}
	for _, tag := range ev.Tags {
		if tag.Name() != "p" || len(tag) < 2 {
			continue
		}
		var pub [32]byte
		if nostr.Decode32(tag[1], pub[:]) != nil {
			continue
		}
		pkSer, err := e.pubkeySerial(pub[:])
		if err != nil {
			return err
		}
		if _, ok := seenPK[pkSer]; ok {
			continue
		}
		seenPK[pkSer] = struct{}{}
		if err := e.epg.Put(makeEpg(ser, pkSer, ev.Kind, DirPTagOut)); err != nil {
			return err
		}
		if err := e.peg.Put(makePeg(pkSer, ev.Kind, DirPTagIn, ser)); err != nil {
			return err
		}
		if err := e.ppg.Put(makePpg(authorSer, pkSer, ev.Kind, DirPKOut, ser)); err != nil {
			return err
		}
		if err := e.gpp.Put(makeGpp(pkSer, ev.Kind, DirPKIn, authorSer, ser)); err != nil {
			return err
		}
	}

	// A single event can carry an unbounded tag list. The graph is an
	// acceleration structure, so cap the work one event can cause while the
	// event itself remains fully stored and queryable.
	seenEV := make(map[uint64]struct{})
	for _, tag := range ev.Tags {
		if tag.Name() != "e" || len(tag) < 2 {
			continue
		}
		var id [32]byte
		if nostr.Decode32(tag[1], id[:]) != nil {
			continue
		}
		tgtSer, ok, err := e.serialByID(id[:])
		if err != nil || !ok {
			// The target is not stored yet. backfillInboundETags creates this
			// edge when it does arrive.
			continue
		}
		if _, ok := seenEV[tgtSer]; ok {
			continue
		}
		seenEV[tgtSer] = struct{}{}
		if err := e.eeg.Put(makeEeg(ser, tgtSer, ev.Kind, DirETagOut)); err != nil {
			return err
		}
		if err := e.gee.Put(makeGee(tgtSer, ev.Kind, DirETagIn, ser)); err != nil {
			return err
		}
	}
	return nil
}

// backfillInboundETags creates the inbound edges for events that were stored
// before the event they reference. Without it, publishing a reply before its
// root would leave the thread graph incomplete forever. The tag index already
// holds the reference, so the backfill costs one prefix scan per new event.
func (e *Engine) backfillInboundETags(ev *nostr.Event, ser uint64) error {
	targetHex := nostr.Hex(ev.ID[:])
	vh := Hash8([]byte(targetHex))
	start := makeTC('e', vh, 0, 0)
	end := prefixEnd(start[:3+1+HashLen])
	var scanErr error
	err := e.tc.Scan(start, end, func(rec []byte) bool {
		srcSer := GetSerial(rec[len(rec)-SerialLen:])
		if srcSer == ser {
			return true
		}
		src, err := e.getBySerial(srcSer)
		if err != nil || !src.Tags.Has("e", targetHex) {
			return true
		}
		if graphTagCount(src) > maxGraphTags {
			return true
		}
		row, ok, err := e.eeg.Get(makeEeg(srcSer, ser, src.Kind, DirETagOut))
		if err != nil {
			scanErr = err
			return false
		}
		if ok && len(row) > 0 {
			return true
		}
		if err := e.eeg.Put(makeEeg(srcSer, ser, src.Kind, DirETagOut)); err != nil {
			scanErr = err
			return false
		}
		if err := e.gee.Put(makeGee(ser, src.Kind, DirETagIn, srcSer)); err != nil {
			scanErr = err
			return false
		}
		return true
	})
	if err != nil {
		return err
	}
	return scanErr
}

// removeGraph deletes every graph edge that names the event. Scans are used
// instead of key reconstruction because a referenced event may already have
// been removed and its serial can no longer be looked up.
func (e *Engine) removeGraph(ev *nostr.Event, ser uint64) {
	// event -> pubkey, plus the matching reverse pubkey -> event rows.
	_ = e.scanSerial(e.epg, pEpg, ser, func(rec []byte) bool {
		pkSer := GetSerial(rec[8:13])
		kind := GetKind(rec[13:15])
		e.epg.Delete(rec)
		switch rec[15] {
		case DirAuthor:
			e.peg.Delete(makePeg(pkSer, kind, DirAuthor, ser))
		case DirPTagOut:
			e.peg.Delete(makePeg(pkSer, kind, DirPTagIn, ser))
		}
		return true
	})

	// event -> event, outbound e-tags.
	_ = e.scanSerial(e.eeg, pEeg, ser, func(rec []byte) bool {
		tgtSer := GetSerial(rec[8:13])
		kind := GetKind(rec[13:15])
		e.eeg.Delete(rec)
		e.gee.Delete(makeGee(tgtSer, kind, DirETagIn, ser))
		return true
	})

	// event -> event, inbound e-tags where this event is the target.
	_ = e.scanSerial(e.gee, pGee, ser, func(rec []byte) bool {
		kind := GetKind(rec[8:10])
		srcSer := GetSerial(rec[11:16])
		e.gee.Delete(rec)
		e.eeg.Delete(makeEeg(srcSer, ser, kind, DirETagOut))
		return true
	})

	authorSer, ok, err := e.lookupPubkeySerial(ev.PubKey[:])
	if err != nil || !ok {
		return
	}

	// pubkey -> pubkey, edges sourced by this event's author.
	_ = e.scanSerial(e.ppg, pPpg, authorSer, func(rec []byte) bool {
		if GetSerial(rec[16:21]) != ser {
			return true
		}
		tgtPK := GetSerial(rec[8:13])
		kind := GetKind(rec[13:15])
		e.ppg.Delete(rec)
		e.gpp.Delete(makeGpp(tgtPK, kind, DirPKIn, authorSer, ser))
		return true
	})
}

func (e *Engine) scanSerial(f *sortedFile, p prefix, ser uint64, fn func(rec []byte) bool) error {
	start := graphPrefix(p, ser)
	return f.Scan(start, prefixEnd(start), fn)
}

// TraverseGraph runs a _graph query against the materialized indexes.
func (e *Engine) TraverseGraph(q *nostr.GraphQuery) (*nostr.GraphResult, error) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.closed {
		return nil, ErrClosed
	}
	if err := q.Validate(); err != nil {
		return nil, err
	}
	if err := e.flushPending(); err != nil {
		return nil, err
	}
	start := time.Now()
	var (
		res *nostr.GraphResult
		err error
	)
	switch q.Edge {
	case nostr.GraphEdgePubkeyPubkey:
		res, err = e.traversePubkeyPubkey(q)
	case nostr.GraphEdgePubkeyEvent:
		res, err = e.traversePubkeyEvent(q)
	case nostr.GraphEdgeEventEvent:
		res, err = e.traverseEventEvent(q)
	default:
		return nil, ErrGraph
	}
	if err != nil {
		return nil, err
	}
	res.Sort()
	res.Elapsed = time.Since(start).String()
	return res, nil
}

func (e *Engine) seedPubkey(q *nostr.GraphQuery) (uint64, bool, error) {
	var raw [32]byte
	if err := nostr.Decode32(q.Seed, raw[:]); err != nil {
		return 0, false, errGraphSeed()
	}
	return e.lookupPubkeySerial(raw[:])
}

func errGraphSeed() error { return errors.New("store: graph seed is not a pubkey") }

// traversePubkeyPubkey is a breadth-first walk of kind-3 p-tag edges. Kind 3
// is the contact list: a p tag there is a follow.
func (e *Engine) traversePubkeyPubkey(q *nostr.GraphQuery) (*nostr.GraphResult, error) {
	res := nostr.NewGraphResult()
	seed, ok, err := e.seedPubkey(q)
	if err != nil || !ok {
		return res, err
	}
	visited := map[uint64]struct{}{seed: {}}
	frontier := []uint64{seed}
	var scanErr error

	for depth := 1; depth <= q.Depth; depth++ {
		var next []uint64
		for _, ser := range frontier {
			if res.TotalPubkeys >= e.maxResults {
				break
			}
			visit := func(neighbor uint64) bool {
				if res.TotalPubkeys >= e.maxResults {
					return false
				}
				if _, seen := visited[neighbor]; seen {
					return true
				}
				visited[neighbor] = struct{}{}
				next = append(next, neighbor)
				pub, ok := e.pubkeyBySerial(neighbor)
				if ok {
					res.AddPubkeyAtDepth(nostr.Hex(pub[:]), depth)
				}
				return res.TotalPubkeys < e.maxResults
			}
			if q.IsOutbound() {
				scanErr = e.scanSerial(e.ppg, pPpg, ser, func(rec []byte) bool {
					if GetKind(rec[13:15]) != 3 {
						return true
					}
					return visit(GetSerial(rec[8:13]))
				})
			}
			if scanErr == nil && q.IsInbound() {
				scanErr = e.scanSerial(e.gpp, pGpp, ser, func(rec []byte) bool {
					if GetKind(rec[8:10]) != 3 {
						return true
					}
					return visit(GetSerial(rec[11:16]))
				})
			}
			if scanErr != nil {
				return nil, scanErr
			}
		}
		if len(next) == 0 {
			break
		}
		frontier = next
	}
	return res, nil
}

// traversePubkeyEvent returns the events directly connected to the seed
// pubkey: authored events for "out", p-tag mentions for "in". The pubkey↔event
// graph is bipartite, so a second hop would change the result type; orly's
// extension also defines pe as a one-hop query.
func (e *Engine) traversePubkeyEvent(q *nostr.GraphQuery) (*nostr.GraphResult, error) {
	res := nostr.NewGraphResult()
	seed, ok, err := e.seedPubkey(q)
	if err != nil || !ok {
		return res, err
	}
	visit := func(ser uint64) bool {
		if res.TotalEvents >= e.maxResults {
			return false
		}
		id, ok := e.eventIDBySerial(ser)
		if ok {
			res.AddEventAtDepth(nostr.Hex(id[:]), 1)
		}
		return res.TotalEvents < e.maxResults
	}
	if q.IsOutbound() {
		if err := e.scanSerial(e.peg, pPeg, seed, func(rec []byte) bool {
			if rec[10] != DirAuthor {
				return true
			}
			return visit(GetSerial(rec[11:16]))
		}); err != nil {
			return nil, err
		}
	}
	if q.IsInbound() {
		if err := e.scanSerial(e.peg, pPeg, seed, func(rec []byte) bool {
			if rec[10] != DirPTagIn {
				return true
			}
			return visit(GetSerial(rec[11:16]))
		}); err != nil {
			return nil, err
		}
	}
	return res, nil
}

// traverseEventEvent starts from every event authored by the seed pubkey and
// follows e-tag edges. The authored events are roots, not results.
func (e *Engine) traverseEventEvent(q *nostr.GraphQuery) (*nostr.GraphResult, error) {
	res := nostr.NewGraphResult()
	seed, ok, err := e.seedPubkey(q)
	if err != nil || !ok {
		return res, err
	}
	roots, err := e.authoredEventSerials(seed)
	if err != nil {
		return nil, err
	}
	sort.Slice(roots, func(i, j int) bool { return roots[i] < roots[j] })

	seen := make(map[uint64]struct{}, len(roots))
	for _, root := range roots {
		seen[root] = struct{}{}
	}
	for _, root := range roots {
		local := map[uint64]struct{}{root: {}}
		frontier := []uint64{root}
		for depth := 1; depth <= q.Depth; depth++ {
			var next []uint64
			for _, ser := range frontier {
				if res.TotalEvents >= e.maxResults {
					break
				}
				visit := func(neighbor uint64) bool {
					if _, ok := local[neighbor]; ok {
						return true
					}
					local[neighbor] = struct{}{}
					next = append(next, neighbor)
					if _, ok := seen[neighbor]; !ok {
						seen[neighbor] = struct{}{}
						if id, ok := e.eventIDBySerial(neighbor); ok {
							res.AddEventAtDepth(nostr.Hex(id[:]), depth)
						}
					}
					return res.TotalEvents < e.maxResults
				}
				if q.IsOutbound() {
					err = e.scanSerial(e.eeg, pEeg, ser, func(rec []byte) bool {
						return visit(GetSerial(rec[8:13]))
					})
				}
				if err == nil && q.IsInbound() {
					err = e.scanSerial(e.gee, pGee, ser, func(rec []byte) bool {
						return visit(GetSerial(rec[11:16]))
					})
				}
				if err != nil {
					return nil, err
				}
			}
			if len(next) == 0 {
				break
			}
			frontier = next
		}
	}
	return res, nil
}

func (e *Engine) authoredEventSerials(pubSer uint64) ([]uint64, error) {
	var out []uint64
	seen := make(map[uint64]struct{})
	err := e.scanSerial(e.peg, pPeg, pubSer, func(rec []byte) bool {
		if rec[10] != DirAuthor {
			return true
		}
		ser := GetSerial(rec[11:16])
		if _, ok := seen[ser]; ok {
			return true
		}
		seen[ser] = struct{}{}
		out = append(out, ser)
		return true
	})
	return out, err
}

// RebuildGraph clears the pubkey and graph indexes and materializes them from
// the events that survived import. Bulk imports skip graph maintenance while
// writing; by the time this runs every e-tag target that exists is already
// stored, so no per-event backfill scan is needed.
func (e *Engine) RebuildGraph() error {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.closed {
		return ErrClosed
	}
	for _, f := range []*sortedFile{e.pks, e.spk, e.epg, e.peg, e.eeg, e.gee, e.ppg, e.gpp} {
		if err := f.Clear(); err != nil {
			return err
		}
	}
	e.nextPubSer = 1
	e.bulkPub = make(map[string]uint64)
	if err := e.spk.Scan(nil, nil, func(rec []byte) bool {
		var pub [32]byte
		copy(pub[:], rec[spkRecLen:spkRecLen+spkValLen])
		e.bulkPub[string(pub[:])] = GetSerial(rec[3:])
		return true
	}); err != nil {
		return err
	}

	live := e.bulkIDs
	if live == nil {
		live = make(map[[32]byte]uint64)
		if err := e.sei.Scan(nil, nil, func(rec []byte) bool {
			var id [32]byte
			copy(id[:], rec[seiRecLen:])
			live[id] = GetSerial(rec[3:])
			return true
		}); err != nil {
			return err
		}
	}
	total := len(live)
	done := 0
	var graphErr error
	err := e.w.ForEachFrom(0, func(ser uint64, data []byte) bool {
		ev, derr := decodeEvent(data)
		if derr != nil {
			return true
		}
		if _, ok := live[ev.ID]; !ok {
			return true
		}
		if ierr := e.indexGraph(ev, ser); ierr != nil {
			graphErr = ierr
			return false
		}
		done++
		if done%100000 == 0 {
			fmt.Fprintf(os.Stderr, "graph: %d/%d live events\n", done, total)
		}
		return true
	})
	if err != nil {
		return err
	}
	if graphErr != nil {
		return graphErr
	}

	e.bulk = false
	e.bulkRep = nil
	e.bulkIDs = nil
	e.bulkPub = nil
	return e.flushPending()
}