diamond-orm/query.go

package orm

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/fatih/structtag"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"log"
	"reflect"
	"strings"
	"time"
)

type Query struct {
	// the handle of the collection associated with this query
	Collection *mongo.Collection
	// the operation from which this query stems
	Op string
	// the model instance associated with this query
	Model  Model
	done   bool
	rawDoc any
	doc    any
}

const (
	OP_FIND_ONE   = "findOne"
	OP_FIND_PAGED = "findPaged"
	OP_FIND_ALL   = "findAll"
	OP_FIND       = "find"
)

func populate(r Reference, rcoll string, rawDoc interface{}, d string, src interface{}) any {
	rt := reflect.TypeOf(src)
	rv := reflect.ValueOf(src)
	srt := rt
	if srt.Kind() == reflect.Pointer {
		srt = rt.Elem()
	}
	if rv.Kind() != reflect.Pointer {
		rv = reflect.New(rt)
		rv.Elem().Set(reflect.ValueOf(src))
	}

	var fieldsMap [3]string
	type bsonWhat struct {
		What string
	}
	var w bsonWhat
	switch rawDoc.(type) {
	case bson.A:
		w.What = "A"
	case bson.D:
		w.What = "D"
	case bson.M:
		w.What = "M"
	default:
		w.What = "-"
	}
	var next string
	if len(strings.Split(d, ".")) > 1 {
		next = strings.Join(strings.Split(d, ".")[1:], ".")
		d = strings.Split(d, ".")[0]
	} else {
		next = d
	}
	var toReturn interface{}
	switch w.What {
	case "A":
		rvs := reflect.MakeSlice(rt, 0, 0)
		var rahh reflect.Value
		if rv.Kind() == reflect.Ptr {
			rahh = rv.Elem()
		} else {
			rahh = rv
		}
		for i, el := range rawDoc.(bson.A) {
			it := rahh.Index(i)
			popped := populate(r, rcoll, el, next, it.Interface())
			poppedVal := reflect.ValueOf(popped)
			if poppedVal.Kind() == reflect.Pointer {
				rvs = reflect.Append(rvs, poppedVal.Elem())
			} else {
				rvs = reflect.Append(rvs, poppedVal)
			}
		}
		if rv.CanSet() {
			rv.Set(rvs)
		} else if rv.Kind() == reflect.Pointer {
			rv.Elem().Set(rvs)
		} else {
			src = rvs.Interface()
			toReturn = src
		}
	case "D":
		loc := rawDoc.(bson.D)
		nrd := bson.M{}
		for _, el := range loc {
			nrd[el.Key] = el.Value
		}
		rawDoc = nrd
		fallthrough
	case "M":
		dd := rawDoc.(bson.M)
		var sf reflect.Value
		var rsf reflect.Value
		if rv.Kind() == reflect.Pointer {
			sf = rv.Elem().FieldByName(d)
		} else {
			sf = rv.FieldByName(d)
		}
		if rv.Kind() == reflect.Pointer {
			rsf = rv.Elem().FieldByName(d)
		} else {
			rsf = rv.FieldByName(d)
		}

		var ff reflect.StructField
		var ok bool
		if rt.Kind() == reflect.Pointer {
			ff, ok = rt.Elem().FieldByName(d)
		} else {
			ff, ok = rt.FieldByName(d)
		}

		if ok {
			tag, err := structtag.Parse(string(ff.Tag))
			if err == nil {
				val, err2 := tag.Get("bson")
				if err2 == nil {
					fttt := ff.Type
					if fttt.Kind() == reflect.Pointer || fttt.Kind() == reflect.Slice {
						fttt = fttt.Elem()
					}
					fieldsMap = [3]string{d, fttt.Name(), val.Name}
				}
			}
		}
		intermediate := populate(r, rcoll, dd[fieldsMap[2]], next, sf.Interface())
		if rsf.CanSet() {
			ival := reflect.ValueOf(intermediate)
			if ival.Kind() != reflect.Pointer && rsf.Kind() == reflect.Pointer {
				rsf.Set(ival.Elem())
			} else if ival.Kind() == reflect.Pointer && rsf.Kind() != reflect.Pointer {
				rsf.Set(ival.Elem())
			} else {
				rsf.Set(ival)
			}
		} else {
			src = intermediate
		}
	default:
		tto := r.HydratedType
		if tto.Kind() == reflect.Pointer || tto.Kind() == reflect.Slice {
			tto = tto.Elem()
		}

		rawt := ModelRegistry.new_(tto.Name())
		t := rawt.(IModel)
		q := bson.M{"_id": rawDoc}
		reso := DB.Collection(rcoll).FindOne(context.TODO(), q)
		reso.Decode(t)
		hatred := rv
		if hatred.Kind() == reflect.Pointer {
			hatred = hatred.Elem()
		}
		if hatred.CanSet() {
			if reflect.ValueOf(t).Kind() == reflect.Pointer {
				fmt.Println(reflect.ValueOf(t).Elem().Type().String())
				fmt.Println(rv.Type().String())
				hatred.Set(reflect.ValueOf(t).Elem())
			} else {
				hatred.Set(reflect.ValueOf(t))
			}
		} else {
			src = t
			toReturn = src
		}
	}

	if toReturn == nil {
		return rv.Interface()
	}
	return src
}

// Populate populates document references via reflection
func (q *Query) Populate(fields ...string) *Query {
	_, cm, _ := ModelRegistry.HasByName(q.Model.typeName)

	if cm != nil {
		rawDoc := q.rawDoc
		for _, field := range fields {
			// 0 = fieldname, 1 = typename, 2 = bson name

			var r Reference

			for k2, v := range cm.References {
				if strings.HasPrefix(k2, field) {
					r = v
					break
				}
			}

			if r.Exists {
				// get self
				// get ptr
				// find
				// unmarshal...
				htt := r.HydratedType
				if htt.Kind() == reflect.Pointer || htt.Kind() == reflect.Slice {
					htt = htt.Elem()
				}
				if strings.HasSuffix(field, ".") || strings.HasPrefix(field, ".") {
					log.Printf("WARN: invalid field name passed to Populate(). skipping...\n")
					continue
				}

				tto := r.HydratedType
				if tto.Kind() == reflect.Pointer || tto.Kind() == reflect.Slice {
					tto = tto.Elem()
				}
				_, refColl, _ := ModelRegistry.HasByName(tto.Name())
				var tmp1 interface{}
				if arr, ok := rawDoc.(bson.A); ok {
					typ := cm.Type
					slic := reflect.New(reflect.SliceOf(typ))
					for i, val2 := range arr {
						ref := reflect.ValueOf(q.doc)
						if ref.Kind() == reflect.Pointer {
							ref = ref.Elem()
						}
						src := ref.Index(i).Interface()
						inter := populate(r, refColl.Collection, val2, field, src)
						if reflect.ValueOf(inter).Kind() == reflect.Pointer {
							slic.Elem().Set(reflect.Append(slic.Elem(), reflect.ValueOf(inter).Elem()))
						} else {
							slic.Elem().Set(reflect.Append(slic, reflect.ValueOf(inter)))
						}
					}
					tmp1 = slic.Interface()
				} else {
					tmp1 = populate(r, refColl.Collection, rawDoc, field, reflect.ValueOf(q.doc).Interface())
				}
				q.doc = tmp1
			}
		}
	}
	return q
}
func (q *Query) reOrganize() {
	var trvo reflect.Value
	if arr, ok := q.rawDoc.(bson.A); ok {
		typ := ModelRegistry[q.Model.typeName].Type
		slic := reflect.New(reflect.SliceOf(typ))
		for _, v2 := range arr {
			inter := reflect.ValueOf(rerere(v2, typ))
			if inter.Kind() == reflect.Pointer {
				inter = inter.Elem()
			}
			slic.Elem().Set(reflect.Append(slic.Elem(), inter))
		}
		trvo = slic.Elem()
	} else {
		trvo = reflect.ValueOf(rerere(q.rawDoc, reflect.TypeOf(q.doc)))
		for {
			if trvo.Kind() == reflect.Pointer {
				trvo = trvo.Elem()
			} else {
				break
			}
		}
	}

	resV := reflect.ValueOf(q.doc)
	for {
		if resV.Kind() == reflect.Pointer {
			resV = resV.Elem()
		} else {
			break
		}
	}
	resV.Set(trvo)
}

func rerere(input interface{}, resType reflect.Type) interface{} {
	t := reflect.TypeOf(input)
	v := reflect.ValueOf(input)
	if input == nil {
		return nil
	}
	if v.Type().Kind() == reflect.Pointer {
		v = v.Elem()
	}
	if t.Kind() == reflect.Pointer {
		t = t.Elem()
	}

	if resType.Kind() == reflect.Pointer {
		resType = resType.Elem()
	}
	resV := reflect.New(resType)
	var rve reflect.Value = resV
	if rve.Kind() == reflect.Pointer {
		rve = resV.Elem()
	}

	if d, isD := v.Interface().(bson.D); isD {
		m := bson.M{}
		for _, el := range d {
			m[el.Key] = el.Value
		}
		input = m
		v = reflect.ValueOf(input)
	}

	switch resType.Kind() {
	case reflect.Struct:
		shouldBreak := false
		mipmap, ok := v.Interface().(bson.M)
		if ok {
			for i := 0; i < resType.NumField(); i++ {
				ft := resType.Field(i)
				fv := rve.Field(i)
				if ft.Anonymous {
					fv.Set(handleAnon(input, ft.Type, fv))
					continue
				}
				tags, err := structtag.Parse(string(ft.Tag))
				panik(err)
				btag, err := tags.Get("bson")
				if err != nil {
					continue
				}
				if btag.Name == "-" {
					continue
				}
				intermediate := mipmap[btag.Name]
				_, err = tags.Get("ref")
				if err != nil {
					tmp := rerere(intermediate, ft.Type)
					fuck := reflect.ValueOf(tmp)
					if tmp != nil {
						if fuck.Type().Kind() == reflect.Pointer {
							fuck = fuck.Elem()
						}
						fv.Set(fuck)
					} else {
						fv.Set(reflect.Zero(ft.Type))
					}
					shouldBreak = true
				} else {
					tt := ft.Type
					if tt.Kind() == reflect.Pointer {
						tt = tt.Elem()
					}
					tmp := rerere(intermediate, ft.Type)
					if tmp != nil {
						if reflect.ValueOf(tmp).Kind() == reflect.Pointer && fv.Kind() != reflect.Pointer {
							fv.Set(reflect.ValueOf(tmp).Elem())
						} else {
							fv.Set(reflect.ValueOf(tmp))
						}
					} else {
						fv.Set(reflect.Zero(ft.Type))
					}
				}
			}
			if shouldBreak {
				// break
			}
		} else {
			nunu := ModelRegistry.new_(resType.Name())
			ider, ok := nunu.(HasID)
			if ok {
				ider.SetId(v.Interface())
				if reflect.ValueOf(ider).Kind() == reflect.Pointer {
					nunu = reflect.ValueOf(ider).Elem().Interface()
				}
				rve.Set(reflect.ValueOf(nunu))
			}
		}
	case reflect.Slice:
		arr := v.Interface().(bson.A)
		for _, it := range arr {
			if it != nil {
				tmp := reflect.ValueOf(rerere(it, rve.Type().Elem()))
				if tmp.Kind() == reflect.Pointer {
					tmp = tmp.Elem()
				}
				rve.Set(reflect.Append(rve, tmp))
			}
		}
	default:
		if resType.AssignableTo(v.Type()) {
			rve.Set(reflect.ValueOf(input))
		} else {
			switch rve.Interface().(type) {
			case int, int32, int64, uint, uint32, uint64:
				rve.Set(reflect.ValueOf(coerceInt(v, rve)))
			}
		}
	}

	return resV.Interface()
}

func handleAnon(raw interface{}, rtype reflect.Type, rval reflect.Value) reflect.Value {
	f := rval
	g := rtype
	if rtype.Kind() == reflect.Pointer {
		g = rtype.Elem()
	}
	if !f.CanSet() {
		f = reflect.New(g)
		f.Elem().Set(rval)
	}
	if rtype.Kind() != reflect.Struct {
		return rval
	}

	for i := 0; i < rtype.NumField(); i++ {
		typeField := rtype.Field(i)
		valueField := f.Field(i)
		tags, err := structtag.Parse(string(typeField.Tag))
		if !typeField.IsExported() {
			continue
		}
		if err == nil {
			var btag *structtag.Tag
			btag, err = tags.Get("bson")
			if err != nil {
				continue
			}
			if btag.Name == "-" {
				continue
			}
			amap, ok := raw.(bson.M)
			if ok {
				fval := amap[btag.Name]
				if reflect.TypeOf(fval) == reflect.TypeFor[primitive.DateTime]() {
					fval = time.UnixMilli(int64(fval.(primitive.DateTime)))
				}
				if valueField.Kind() == reflect.Pointer {
					valueField.Elem().Set(reflect.ValueOf(fval))
				} else {
					valueField.Set(reflect.ValueOf(fval))
				}
			}
		}
	}
	return f
}

// JSON - marshals this Query's results into json format
func (q *Query) JSON() (string, error) {
	res, err := json.MarshalIndent(q.doc, "\n", "\t")
	if err != nil {
		return "", err
	}
	return string(res[:]), nil
}

// Exec - executes the query and puts its results into the
// provided argument.
//
// Will panic if called more than once on the same Query instance.
func (q *Query) Exec(result interface{}) {
	if q.done {
		panic("Exec() has already been called!")
	}
	reflect.ValueOf(result).Elem().Set(reflect.ValueOf(q.doc).Elem())
	q.Model.self = q.doc
	q.done = true
}