// Copyright 2025 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build vm_histogram

package promutil

import (
	"compress/gzip"
	"io"
	"net/http"
	"strings"

	vmetrics "github.com/VictoriaMetrics/metrics"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var promHandler = promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{
	DisableCompression: true,
})

type combinedMetricsHandler struct{}

func (h combinedMetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")

	var writer io.Writer = w
	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		gz := gzip.NewWriter(w)
		defer gz.Close()
		writer = gz
	}

	metricsWriter := &metricsResponseWriter{Writer: writer}
	promHandler.ServeHTTP(metricsWriter, r)
	if err := metricsWriter.ensureTrailingNewline(); err != nil {
		return
	}
	vmetrics.WritePrometheus(writer, true)
}

type metricsResponseWriter struct {
	io.Writer
	written         bool
	endsWithNewline bool
}

func (w *metricsResponseWriter) Header() http.Header {
	return http.Header{}
}

func (w *metricsResponseWriter) WriteHeader(int) {}

func (w *metricsResponseWriter) Write(b []byte) (int, error) {
	n, err := w.Writer.Write(b)
	if n > 0 {
		w.written = true
		w.endsWithNewline = b[n-1] == '\n'
	}
	return n, err
}

func (w *metricsResponseWriter) ensureTrailingNewline() error {
	if !w.written || w.endsWithNewline {
		return nil
	}
	if _, err := w.Writer.Write([]byte{'\n'}); err != nil {
		return err
	}
	w.endsWithNewline = true
	return nil
}

// MetricsHandler returns an HTTP handler that exposes metrics for both
// Prometheus and VictoriaMetrics when the vm_histogram build tag is enabled.
// It first serves Prometheus metrics via promhttp, then appends VM metrics,
// ensuring the output formatting remains valid for scrapers.
func MetricsHandler() http.Handler {
	return combinedMetricsHandler{}
}