package main

import (
    "context"
    "fmt"
    "github.com/gorilla/mux"
    "go.mongodb.org/mongo-driver/bson"
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
    "log"
    "net/http"
    "time"
)

const (
    handlerTimeout = 3 * time.Second
    mongoUri       = "mongodb://localhost:27017"
    mongoDbName    = "mydb"
)

func main() {
    r := mux.NewRouter()

    mgo, _ := mongo.NewClient(options.Client().ApplyURI(mongoUri))
    err := mgo.Connect(context.Background())
    if err != nil {
        panic(err.Error())
    }

    s := service{
        mgo.Database(mongoDbName),
    }

    withMiddleware := func(f http.HandlerFunc) http.HandlerFunc {
        var f2 http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
            key := r.URL.Query().Get("key")
            userId, err := s.authorize(r.Context(), key)
            if err != nil {
                http.Error(w, "error during authorization: "+err.Error(), 500)
                return
            }

            // context.WithValue() "wraps" the existing request context and adds a single
            // piece of data associated to a key, in our case, a user ID.
            newCtx := context.WithValue(r.Context(), "userId", userId)

            f(w, r.WithContext(newCtx))
        }

        // http.TimeoutHandler, internally, calls context.WithTimeout(), wrapping the context
        // again and therefore adding the timeout functionality we want.
        return http.TimeoutHandler(f2, handlerTimeout, "").ServeHTTP
    }

    r.HandleFunc("/", withMiddleware(s.doStuff))

    srv := &http.Server{Handler: r}

    srv.ListenAndServe()
}

type service struct {
    db *mongo.Database
}

// authorize takes a string key, runs a DB query, and returns the corresponding user ID.
func (s *service) authorize(ctx context.Context, key string) (string, error) {
    var user struct {
        Key    string
        UserId string
    }

    filter := bson.M{"Key": key}

    err := s.db.Collection("users").FindOne(ctx, filter).Decode(&user)
    if err != nil {
        log.Printf("authorize(): error during query: %v", err)
    }

    return user.UserId, err
}

// doStuff outputs every piece of data found in the DB for the current user.
func (s *service) doStuff(w http.ResponseWriter, r *http.Request) {
    ctx := r.Context()
    userId := ctx.Value("userId")

    filter := bson.M{"UserId": userId}

    cur, err := s.db.Collection("data").Find(ctx, filter)
    if err != nil {
        log.Printf("doStuff(): error during query: %v", err)
        return
    }

    for cur.Next(ctx) {
        var data struct {
            Data string
        }

        err := cur.Decode(&data)
        if err != nil {
            log.Printf("doStuff(): error during decoding: %v", err.Error())
            continue
        }

        w.Write([]byte(fmt.Sprintf("data record: %v\n", data.Data)))
    }
}