package main

import (
    "compress/gzip"
    "context"
    "fmt"
    "io"
    "log"
    "net/http"
    "os"
    "path/filepath"
    "strings"
    "sync/atomic"
    "time"

    "github.com/joho/godotenv"
)

type ctxKey string

const ctxReqID ctxKey = "reqID"

// atomic counter for simple request IDs
var reqIDCounter uint64

// log levels
const (
    levelDebug = iota
    levelInfo
    levelWarn
    levelError
)

var currentLogLevel = levelInfo

func parseLogLevel(s string) int {
    switch strings.ToLower(s) {
    case "debug":
        return levelDebug
    case "info", "":
        return levelInfo
    case "warn", "warning":
        return levelWarn
    case "error", "err":
        return levelError
    default:
        return levelInfo
    }
}

// proxyToValhallaPOST forwards a POST with JSON body to a specific Valhalla endpoint (e.g., /isochrone)
func proxyToValhallaPOST(w http.ResponseWriter, r *http.Request, target string) {
    if r.Method != http.MethodPost {
        http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
        return
    }
    req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, target, r.Body)
    if err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
    }
    req.Header = r.Header.Clone()
    if req.Header.Get("Content-Type") == "" {
        req.Header.Set("Content-Type", "application/json")
    }

    client := http.Client{Timeout: 20 * time.Second}
    start := time.Now()
    resp, err := client.Do(req)
    if err != nil {
        if rid, ok := r.Context().Value(ctxReqID).(string); ok {
            logf(levelError, "rid=%s upstream=valhalla target=%s error=%v", rid, target, err)
        } else {
            logf(levelError, "upstream=valhalla target=%s error=%v", target, err)
        }
        http.Error(w, err.Error(), http.StatusBadGateway)
        return
    }
    defer resp.Body.Close()

    // Forward common headers and status
    for key, values := range resp.Header {
        if key == "Content-Type" || key == "Content-Encoding" || key == "Cache-Control" || key == "Vary" {
            for _, v := range values {
                w.Header().Add(key, v)
            }
        }
    }
    w.WriteHeader(resp.StatusCode)
    _, _ = io.Copy(w, resp.Body)

    if rid, ok := r.Context().Value(ctxReqID).(string); ok {
        logf(levelDebug, "rid=%s upstream=valhalla target=%s status=%d dur=%s", rid, target, resp.StatusCode, time.Since(start))
    } else {
        logf(levelDebug, "upstream=valhalla target=%s status=%d dur=%s", target, resp.StatusCode, time.Since(start))
    }
}

func logf(level int, format string, args ...any) {
    if level < currentLogLevel {
        return
    }
    prefix := "INFO"
    switch level {
    case levelDebug:
        prefix = "DEBUG"
    case levelInfo:
        prefix = "INFO"
    case levelWarn:
        prefix = "WARN"
    case levelError:
        prefix = "ERROR"
    }
    // Prepend the prefix as the first argument
    log.Printf("[%s] "+format, append([]any{prefix}, args...)...)
}

func main() {
    _ = godotenv.Load(".env") // Load .env file

    valhallaURL := os.Getenv("VALHALLA_URL")
    if valhallaURL == "" {
        valhallaURL = "http://valhalla:8002/route"
    }
    // Derive Valhalla base for other endpoints like /isochrone
    valhallaBase := strings.TrimSuffix(valhallaURL, "/route")

    nominatimBase := os.Getenv("NOMINATIM_URL")
    if nominatimBase == "" {
        nominatimBase = "https://nominatim.openstreetmap.org"
    }
    nominatimUA := os.Getenv("NOMINATIM_USER_AGENT")
    if nominatimUA == "" {
        nominatimUA = "FreeMoto/1.0 (+https://fm.ztsw.de/)"
    }

    currentLogLevel = parseLogLevel(os.Getenv("LOG_LEVEL"))

    port := os.Getenv("PORT")
    if port == "" {
        port = "8080"
    }

    // Static files under /static with long-lived caching and gzip for text assets
    http.Handle("/static/", withLogging(withCacheGzip(http.StripPrefix("/static/", http.FileServer(http.Dir("./static"))), true)))
    http.Handle("/healthz", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.WriteHeader(http.StatusOK)
        _, _ = w.Write([]byte("ok"))
    })))
    http.Handle("/route", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        proxyToValhalla(w, r, valhallaURL)
    })))
    http.Handle("/isochrone", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // Forward to Valhalla isochrone endpoint (POST JSON)
        if r.Method != http.MethodPost {
            http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
            return
        }
        target := valhallaBase + "/isochrone"
        proxyToValhallaPOST(w, r, target)
    })))
    // Nominatim proxy endpoints
    http.Handle("/geocode", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        proxyToNominatimGET(w, r, nominatimBase+"/search", nominatimUA)
    })))
    http.Handle("/reverse", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        proxyToNominatimGET(w, r, nominatimBase+"/reverse", nominatimUA)
    })))
    http.Handle("/", withLogging(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        // Serve index
        if r.URL.Path == "/" {
            // Short cache for HTML entrypoint
            setShortCache(w)
            maybeGzipAndServeFile(w, r, "./static/index.html")
            return
        }
        // Serve other assets from ./static with long cache for versioned files
        // Decide cache and compression based on extension
        ext := strings.ToLower(filepath.Ext(r.URL.Path))
        if isTextLikeExt(ext) {
            setLongCache(w)
            w.Header().Add("Vary", "Accept-Encoding")
            // Use gzip when supported
            if acceptsGzip(r) {
                gz := gzip.NewWriter(w)
                defer gz.Close()
                w.Header().Set("Content-Encoding", "gzip")
                // Avoid Content-Length when gzipping
                w.Header().Del("Content-Length")
                gzw := &gzipResponseWriter{ResponseWriter: w, gw: gz}
                http.ServeFile(gzw, r, "./static/"+r.URL.Path[1:])
                return
            }
            http.ServeFile(w, r, "./static/"+r.URL.Path[1:])
            return
        }
        // Non-text assets: just set long cache
        setLongCache(w)
        http.ServeFile(w, r, "./static/"+r.URL.Path[1:])
    })))
    logf(levelInfo, "Listening on :%s", port)
    http.ListenAndServe(":"+port, nil)
}

func proxyToValhalla(w http.ResponseWriter, r *http.Request, valhallaURL string) {
    // Create outbound request with caller context and preserve headers
    req, err := http.NewRequestWithContext(r.Context(), http.MethodPost, valhallaURL, r.Body)
    if err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
    }
    req.Header = r.Header.Clone()

    client := http.Client{Timeout: 15 * time.Second}
    start := time.Now()
    resp, err := client.Do(req)
    if err != nil {
        if rid, ok := r.Context().Value(ctxReqID).(string); ok {
            logf(levelError, "rid=%s upstream=valhalla error=%v", rid, err)
        } else {
            logf(levelError, "upstream=valhalla error=%v", err)
        }
        http.Error(w, err.Error(), http.StatusBadGateway)
        return
    }
    defer resp.Body.Close()
    // Propagate relevant headers and status code
    for key, values := range resp.Header {
        // Optionally filter headers; here we forward common ones
        if key == "Content-Type" || key == "Content-Encoding" || key == "Cache-Control" || key == "Vary" {
            for _, v := range values {
                w.Header().Add(key, v)
            }
        }
    }
    w.WriteHeader(resp.StatusCode)
    _, _ = io.Copy(w, resp.Body)

    if rid, ok := r.Context().Value(ctxReqID).(string); ok {
        logf(levelDebug, "rid=%s upstream=valhalla status=%d dur=%s", rid, resp.StatusCode, time.Since(start))
    } else {
        logf(levelDebug, "upstream=valhalla status=%d dur=%s", resp.StatusCode, time.Since(start))
    }
}

// proxyToNominatimGET proxies GET requests to Nominatim endpoints (search/reverse),
// sets a descriptive User-Agent, and enforces a timeout.
func proxyToNominatimGET(w http.ResponseWriter, r *http.Request, target string, userAgent string) {
    if r.Method != http.MethodGet {
        http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
        return
    }
    // Preserve original query parameters
    url := target
    if rq := r.URL.RawQuery; rq != "" {
        url = url + "?" + rq
    }

    req, err := http.NewRequestWithContext(r.Context(), http.MethodGet, url, nil)
    if err != nil {
        http.Error(w, err.Error(), http.StatusBadRequest)
        return
    }
    // Set a proper User-Agent per Nominatim usage policy
    req.Header.Set("User-Agent", userAgent)
    req.Header.Set("Accept", "application/json")

    client := http.Client{Timeout: 10 * time.Second}
    start := time.Now()
    resp, err := client.Do(req)
    if err != nil {
        if rid, ok := r.Context().Value(ctxReqID).(string); ok {
            logf(levelError, "rid=%s upstream=nominatim url=%s error=%v", rid, url, err)
        } else {
            logf(levelError, "upstream=nominatim url=%s error=%v", url, err)
        }
        http.Error(w, err.Error(), http.StatusBadGateway)
        return
    }
    defer resp.Body.Close()

    // Forward selected headers and status code
    if ct := resp.Header.Get("Content-Type"); ct != "" {
        w.Header().Set("Content-Type", ct)
    }
    w.WriteHeader(resp.StatusCode)
    _, _ = io.Copy(w, resp.Body)

    if rid, ok := r.Context().Value(ctxReqID).(string); ok {
        logf(levelDebug, "rid=%s upstream=nominatim status=%d dur=%s", rid, resp.StatusCode, time.Since(start))
    } else {
        logf(levelDebug, "upstream=nominatim status=%d dur=%s", resp.StatusCode, time.Since(start))
    }
}

// withLogging wraps handlers to add request ID, status, duration logging
func withLogging(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        rid := fmt.Sprintf("%x", atomic.AddUint64(&reqIDCounter, 1))
        ctx := context.WithValue(r.Context(), ctxReqID, rid)
        r = r.WithContext(ctx)

        lrw := &loggingResponseWriter{ResponseWriter: w, status: 200}
        start := time.Now()
        next.ServeHTTP(lrw, r)
        dur := time.Since(start)

        logf(levelInfo, "rid=%s method=%s path=%s status=%d bytes=%d dur=%s ua=\"%s\" ip=%s",
            rid, r.Method, r.URL.Path, lrw.status, lrw.bytes, dur, r.UserAgent(), clientIP(r))
    })
}

type loggingResponseWriter struct {
    http.ResponseWriter
    status int
    bytes  int
}

func (lrw *loggingResponseWriter) WriteHeader(code int) {
    lrw.status = code
    lrw.ResponseWriter.WriteHeader(code)
}

func (lrw *loggingResponseWriter) Write(b []byte) (int, error) {
    n, err := lrw.ResponseWriter.Write(b)
    lrw.bytes += n
    return n, err
}

func clientIP(r *http.Request) string {
    // honor X-Forwarded-For if present (first IP)
    if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
        return xff
    }
    return r.RemoteAddr
}

// ---- Performance helpers: caching and gzip compression ----

func setShortCache(w http.ResponseWriter) {
    // Short cache e.g., for HTML entrypoint
    w.Header().Set("Cache-Control", "public, max-age=300")
}

func setLongCache(w http.ResponseWriter) {
    // Long-lived cache for versioned static assets
    w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
}

func isTextLikeExt(ext string) bool {
    switch ext {
    case ".html", ".css", ".js", ".mjs", ".json", ".svg", ".xml", ".txt", ".map", ".webmanifest":
        return true
    default:
        return false
    }
}

func acceptsGzip(r *http.Request) bool {
    ae := r.Header.Get("Accept-Encoding")
    return strings.Contains(ae, "gzip")
}

type gzipResponseWriter struct {
    http.ResponseWriter
    gw *gzip.Writer
}

func (g *gzipResponseWriter) Write(b []byte) (int, error) {
    return g.gw.Write(b)
}

// withCacheGzip wraps a handler to set cache headers and gzip text-like responses
func withCacheGzip(next http.Handler, longCache bool) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if longCache {
            setLongCache(w)
        } else {
            setShortCache(w)
        }
        ext := strings.ToLower(filepath.Ext(r.URL.Path))
        if isTextLikeExt(ext) && acceptsGzip(r) {
            w.Header().Add("Vary", "Accept-Encoding")
            w.Header().Set("Content-Encoding", "gzip")
            w.Header().Del("Content-Length")
            gz := gzip.NewWriter(w)
            defer gz.Close()
            gzw := &gzipResponseWriter{ResponseWriter: w, gw: gz}
            next.ServeHTTP(gzw, r)
            return
        }
        next.ServeHTTP(w, r)
    })
}

func maybeGzipAndServeFile(w http.ResponseWriter, r *http.Request, path string) {
    // For text-like files, gzip when supported
    ext := strings.ToLower(filepath.Ext(path))
    if isTextLikeExt(ext) && acceptsGzip(r) {
        w.Header().Add("Vary", "Accept-Encoding")
        w.Header().Set("Content-Encoding", "gzip")
        w.Header().Del("Content-Length")
        gz := gzip.NewWriter(w)
        defer gz.Close()
        gzw := &gzipResponseWriter{ResponseWriter: w, gw: gz}
        http.ServeFile(gzw, r, path)
        return
    }
    http.ServeFile(w, r, path)
}