package api

import (
    "context"
    "crypto/rand"
    "encoding/hex"
    "encoding/json"
    "net/http"
)

type ctxKey string

var traceKey ctxKey = "trace_id"
var sqlKey ctxKey = "sql"
var metaKey ctxKey = "req_meta"
var payloadKey ctxKey = "payload"

func withTrace(h http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        tid := r.Header.Get("X-Request-ID")
        if tid == "" {
            buf := make([]byte, 16)
            _, _ = rand.Read(buf)
            tid = hex.EncodeToString(buf)
        }
        w.Header().Set("X-Request-ID", tid)
        m := ReqMeta{Method: r.Method, Path: r.URL.Path, Query: r.URL.RawQuery, Remote: r.RemoteAddr}
        ctx := context.WithValue(r.Context(), traceKey, tid)
        ctx = context.WithValue(ctx, metaKey, m)
        h.ServeHTTP(w, r.WithContext(ctx))
    })
}

func TraceIDFrom(r *http.Request) string {
	v := r.Context().Value(traceKey)
	if v == nil {
		return ""
	}
	s, _ := v.(string)
	return s
}

func WithSQL(r *http.Request, sql string) *http.Request {
    return r.WithContext(context.WithValue(r.Context(), sqlKey, sql))
}

func SQLFrom(r *http.Request) string {
    v := r.Context().Value(sqlKey)
    if v == nil {
        return ""
    }
    s, _ := v.(string)
    return s
}

type ReqMeta struct {
    Method string
    Path   string
    Query  string
    Remote string
}

func MetaFrom(r *http.Request) ReqMeta {
    v := r.Context().Value(metaKey)
    if v == nil {
        return ReqMeta{}
    }
    m, _ := v.(ReqMeta)
    return m
}

func WithPayload(r *http.Request, v interface{}) *http.Request {
    b, _ := json.Marshal(v)
    return r.WithContext(context.WithValue(r.Context(), payloadKey, string(b)))
}

func PayloadFrom(r *http.Request) string {
    v := r.Context().Value(payloadKey)
    if v == nil {
        return ""
    }
    s, _ := v.(string)
    return s
}