package main

import (
	"bufio"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
)

// Body-size histogram bucket upper bounds in bytes.
const promNumBodyBounds = 7

var promBodyBounds = [promNumBodyBounds]int64{256, 1024, 4096, 16384, 65536, 262144, 1048576}

// Request-time histogram bucket upper bounds in seconds (standard Prometheus defaults).
const promNumTimeBounds = 11

var promTimeBounds = [promNumTimeBounds]float64{.005, .01, .025, .05, .1, .25, .5, 1, 2.5, 5, 10}

const promCounterCap = 100_000 // safety cap on {host,method,status} counter entries

// promCounterKey is the label set for per-request counters.
type promCounterKey struct {
	Host   string
	Method string
	Status string
}

// promBodyEntry holds the body_bytes_sent histogram for one host.
type promBodyEntry struct {
	buckets [promNumBodyBounds + 1]int64 // indices 0..N-1: le=bound[i]; index N: le=+Inf
	sum     int64
}

// promTimeEntry holds the request_time histogram for one host.
type promTimeEntry struct {
	buckets [promNumTimeBounds + 1]int64
	sum     float64
}

// PromStore accumulates Prometheus metrics ingested from log records.
//
// Ingest must be called from exactly one goroutine (the store's Run goroutine).
// ServeHTTP may be called from any number of goroutines concurrently.
type PromStore struct {
	mu       sync.Mutex
	counters map[promCounterKey]int64
	body     map[string]*promBodyEntry // keyed by host
	reqTime  map[string]*promTimeEntry // keyed by host
}

// NewPromStore returns an empty PromStore ready for use.
func NewPromStore() *PromStore {
	return &PromStore{
		counters: make(map[promCounterKey]int64, 1024),
		body:     make(map[string]*promBodyEntry, 64),
		reqTime:  make(map[string]*promTimeEntry, 64),
	}
}

// Ingest records one log record into the Prometheus metrics.
// Must be called from a single goroutine.
func (p *PromStore) Ingest(r LogRecord) {
	p.mu.Lock()

	// --- per-{host,method,status} request counter ---
	ck := promCounterKey{Host: r.Website, Method: r.Method, Status: r.Status}
	if _, ok := p.counters[ck]; ok {
		p.counters[ck]++
	} else if len(p.counters) < promCounterCap {
		p.counters[ck] = 1
	}

	// --- body_bytes_sent histogram (keyed by host only) ---
	be, ok := p.body[r.Website]
	if !ok {
		be = &promBodyEntry{}
		p.body[r.Website] = be
	}
	for i, bound := range promBodyBounds {
		if r.BodyBytesSent <= bound {
			be.buckets[i]++
		}
	}
	be.buckets[promNumBodyBounds]++ // +Inf
	be.sum += r.BodyBytesSent

	// --- request_time histogram (keyed by host only) ---
	te, ok := p.reqTime[r.Website]
	if !ok {
		te = &promTimeEntry{}
		p.reqTime[r.Website] = te
	}
	for i, bound := range promTimeBounds {
		if r.RequestTime <= bound {
			te.buckets[i]++
		}
	}
	te.buckets[promNumTimeBounds]++ // +Inf
	te.sum += r.RequestTime

	p.mu.Unlock()
}

// ServeHTTP renders all metrics in the Prometheus text exposition format (0.0.4).
func (p *PromStore) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
	// Snapshot everything under the lock, then render without holding it.
	p.mu.Lock()

	type counterSnap struct {
		k promCounterKey
		v int64
	}
	counters := make([]counterSnap, 0, len(p.counters))
	for k, v := range p.counters {
		counters = append(counters, counterSnap{k, v})
	}

	type bodySnap struct {
		host string
		e    promBodyEntry
	}
	bodySnaps := make([]bodySnap, 0, len(p.body))
	for h, e := range p.body {
		bodySnaps = append(bodySnaps, bodySnap{h, *e})
	}

	type timeSnap struct {
		host string
		e    promTimeEntry
	}
	timeSnaps := make([]timeSnap, 0, len(p.reqTime))
	for h, e := range p.reqTime {
		timeSnaps = append(timeSnaps, timeSnap{h, *e})
	}

	p.mu.Unlock()

	// Sort for stable, human-readable output.
	sort.Slice(counters, func(i, j int) bool {
		a, b := counters[i].k, counters[j].k
		if a.Host != b.Host {
			return a.Host < b.Host
		}
		if a.Method != b.Method {
			return a.Method < b.Method
		}
		return a.Status < b.Status
	})
	sort.Slice(bodySnaps, func(i, j int) bool { return bodySnaps[i].host < bodySnaps[j].host })
	sort.Slice(timeSnaps, func(i, j int) bool { return timeSnaps[i].host < timeSnaps[j].host })

	w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
	bw := bufio.NewWriterSize(w, 256*1024)

	// nginx_http_requests_total
	fmt.Fprintln(bw, "# HELP nginx_http_requests_total Total number of HTTP requests processed.")
	fmt.Fprintln(bw, "# TYPE nginx_http_requests_total counter")
	for _, c := range counters {
		fmt.Fprintf(bw, "nginx_http_requests_total{host=%q,method=%q,status=%q} %d\n",
			c.k.Host, c.k.Method, c.k.Status, c.v)
	}

	// nginx_http_response_body_bytes (histogram, labeled by host)
	fmt.Fprintln(bw, "# HELP nginx_http_response_body_bytes HTTP response body size distribution in bytes.")
	fmt.Fprintln(bw, "# TYPE nginx_http_response_body_bytes histogram")
	for _, s := range bodySnaps {
		for i, bound := range promBodyBounds {
			fmt.Fprintf(bw, "nginx_http_response_body_bytes_bucket{host=%q,le=%q} %d\n",
				s.host, fmt.Sprintf("%d", bound), s.e.buckets[i])
		}
		fmt.Fprintf(bw, "nginx_http_response_body_bytes_bucket{host=%q,le=\"+Inf\"} %d\n",
			s.host, s.e.buckets[promNumBodyBounds])
		fmt.Fprintf(bw, "nginx_http_response_body_bytes_count{host=%q} %d\n",
			s.host, s.e.buckets[promNumBodyBounds])
		fmt.Fprintf(bw, "nginx_http_response_body_bytes_sum{host=%q} %d\n",
			s.host, s.e.sum)
	}

	// nginx_http_request_duration_seconds (histogram, labeled by host)
	fmt.Fprintln(bw, "# HELP nginx_http_request_duration_seconds HTTP request processing time in seconds.")
	fmt.Fprintln(bw, "# TYPE nginx_http_request_duration_seconds histogram")
	for _, s := range timeSnaps {
		for i, bound := range promTimeBounds {
			fmt.Fprintf(bw, "nginx_http_request_duration_seconds_bucket{host=%q,le=%q} %d\n",
				s.host, formatFloat(bound), s.e.buckets[i])
		}
		fmt.Fprintf(bw, "nginx_http_request_duration_seconds_bucket{host=%q,le=\"+Inf\"} %d\n",
			s.host, s.e.buckets[promNumTimeBounds])
		fmt.Fprintf(bw, "nginx_http_request_duration_seconds_count{host=%q} %d\n",
			s.host, s.e.buckets[promNumTimeBounds])
		fmt.Fprintf(bw, "nginx_http_request_duration_seconds_sum{host=%q} %g\n",
			s.host, s.e.sum)
	}

	bw.Flush()
}

// formatFloat renders a float64 bucket bound without trailing zeros but always
// with at least one decimal place, matching Prometheus convention (e.g. "0.5", "10").
func formatFloat(f float64) string {
	s := fmt.Sprintf("%g", f)
	if !strings.Contains(s, ".") && !strings.Contains(s, "e") {
		s += ".0" // ensure it looks like a float, not an integer
	}
	return s
}