// Package health keeps the agent counters and serves the local /healthz and /metrics endpoints.
package health

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"internetpressure.io/probe-agent/internal/protocol"
)

// Gauges are read live from the rest of the agent.
type Gauges struct {
	Buffered      func() int
	SpoolBytes    func() int64
	ClockOffsetMs func() int64
	Targets       func() int
	Identity      func() *protocol.Identity
}

// State is the shared counter set.
type State struct {
	ProbeID      string
	Version      string
	Capabilities []string
	Started      time.Time
	Gauges       Gauges

	measurementsTotal atomic.Uint64
	errorsTotal       atomic.Uint64
	flushFailures     atomic.Uint64
	panics            atomic.Uint64
	lastFlush         atomic.Int64 // unix nano, 0 = never
	lastConfig        atomic.Int64
	lastConfigVersion atomic.Pointer[string]

	mu      sync.Mutex
	byKind  map[string]uint64
	byError map[string]uint64
}

// New creates a State.
func New(probeID, version string, caps []string, g Gauges) *State {
	s := &State{ProbeID: probeID, Version: version, Capabilities: caps, Started: time.Now(), Gauges: g,
		byKind: map[string]uint64{}, byError: map[string]uint64{}}
	empty := ""
	s.lastConfigVersion.Store(&empty)
	return s
}

// Record counts one measurement.
func (s *State) Record(m protocol.Measurement) {
	s.measurementsTotal.Add(1)
	s.mu.Lock()
	s.byKind[m.Kind]++
	if !m.OK {
		s.errorsTotal.Add(1)
		code := m.Error
		if code == "" {
			code = protocol.ErrOther
		}
		s.byError[code]++
	}
	s.mu.Unlock()
}

// RecordPanic counts a recovered panic in a check.
func (s *State) RecordPanic() { s.panics.Add(1) }

// RecordFlushFailure counts a failed batch POST.
func (s *State) RecordFlushFailure() { s.flushFailures.Add(1) }

// SetLastFlush marks a successful delivery.
func (s *State) SetLastFlush(t time.Time) { s.lastFlush.Store(t.UnixNano()) }

// SetLastConfig marks a successful config fetch.
func (s *State) SetLastConfig(t time.Time, version string) {
	s.lastConfig.Store(t.UnixNano())
	v := version
	s.lastConfigVersion.Store(&v)
}

// Uptime in seconds.
func (s *State) Uptime() int64 { return int64(time.Since(s.Started).Seconds()) }

// Totals returns (measurements, errors).
func (s *State) Totals() (uint64, uint64) { return s.measurementsTotal.Load(), s.errorsTotal.Load() }

// RSSMb returns the resident-set estimate (heap sys + stacks, in MiB) from runtime.MemStats.
func RSSMb() float64 {
	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	return float64(ms.Sys-ms.HeapReleased) / (1024 * 1024)
}

func fmtTime(unixNano int64) string {
	if unixNano == 0 {
		return ""
	}
	return protocol.FormatTime(time.Unix(0, unixNano))
}

// Snapshot builds the protocol health block.
func (s *State) Snapshot() *protocol.Health {
	meas, errs := s.Totals()
	h := &protocol.Health{
		TS:                protocol.FormatTime(time.Now()),
		AgentVersion:      s.Version,
		UptimeS:           s.Uptime(),
		MeasurementsTotal: meas,
		ErrorsTotal:       errs,
		RSSMb:             round1(RSSMb()),
		Goroutines:        runtime.NumGoroutine(),
		Capabilities:      s.Capabilities,
		OS:                runtime.GOOS,
		Arch:              runtime.GOARCH,
	}
	if g := s.Gauges; true {
		if g.Buffered != nil {
			h.Buffered = g.Buffered()
		}
		if g.SpoolBytes != nil {
			h.SpoolBytes = g.SpoolBytes()
		}
		if g.ClockOffsetMs != nil {
			h.ClockOffsetMs = g.ClockOffsetMs()
		}
		if g.Identity != nil {
			h.Identity = g.Identity()
		}
	}
	return h
}

func round1(v float64) float64 { return float64(int64(v*10+0.5)) / 10 }

// Handler serves /healthz and /metrics.
func (s *State) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", s.healthz)
	mux.HandleFunc("/metrics", s.metrics)
	return mux
}

func (s *State) healthz(w http.ResponseWriter, _ *http.Request) {
	targets := 0
	if s.Gauges.Targets != nil {
		targets = s.Gauges.Targets()
	}
	var buffered int
	var spool, offset int64
	if s.Gauges.Buffered != nil {
		buffered = s.Gauges.Buffered()
	}
	if s.Gauges.SpoolBytes != nil {
		spool = s.Gauges.SpoolBytes()
	}
	if s.Gauges.ClockOffsetMs != nil {
		offset = s.Gauges.ClockOffsetMs()
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]any{
		"ok":              true,
		"probe_id":        s.ProbeID,
		"version":         s.Version,
		"buffered":        buffered,
		"spool_bytes":     spool,
		"last_flush":      fmtTime(s.lastFlush.Load()),
		"last_config":     fmtTime(s.lastConfig.Load()),
		"config_version":  *s.lastConfigVersion.Load(),
		"targets":         targets,
		"clock_offset_ms": offset,
		"uptime_s":        s.Uptime(),
		"capabilities":    s.Capabilities,
	})
}

func (s *State) metrics(w http.ResponseWriter, _ *http.Request) {
	w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
	s.mu.Lock()
	kinds := make([]string, 0, len(s.byKind))
	for k := range s.byKind {
		kinds = append(kinds, k)
	}
	sort.Strings(kinds)
	codes := make([]string, 0, len(s.byError))
	for k := range s.byError {
		codes = append(codes, k)
	}
	sort.Strings(codes)
	byKind := make(map[string]uint64, len(kinds))
	byErr := make(map[string]uint64, len(codes))
	for k, v := range s.byKind {
		byKind[k] = v
	}
	for k, v := range s.byError {
		byErr[k] = v
	}
	s.mu.Unlock()

	fmt.Fprintln(w, "# TYPE ip_probe_measurements_total counter")
	for _, k := range kinds {
		fmt.Fprintf(w, "ip_probe_measurements_total{kind=%q} %d\n", k, byKind[k])
	}
	fmt.Fprintln(w, "# TYPE ip_probe_errors_total counter")
	for _, k := range codes {
		fmt.Fprintf(w, "ip_probe_errors_total{code=%q} %d\n", k, byErr[k])
	}
	fmt.Fprintln(w, "# TYPE ip_probe_flush_failures_total counter")
	fmt.Fprintf(w, "ip_probe_flush_failures_total %d\n", s.flushFailures.Load())
	fmt.Fprintln(w, "# TYPE ip_probe_check_panics_total counter")
	fmt.Fprintf(w, "ip_probe_check_panics_total %d\n", s.panics.Load())
	fmt.Fprintln(w, "# TYPE ip_probe_buffered gauge")
	if s.Gauges.Buffered != nil {
		fmt.Fprintf(w, "ip_probe_buffered %d\n", s.Gauges.Buffered())
	}
	fmt.Fprintln(w, "# TYPE ip_probe_spool_bytes gauge")
	if s.Gauges.SpoolBytes != nil {
		fmt.Fprintf(w, "ip_probe_spool_bytes %d\n", s.Gauges.SpoolBytes())
	}
	fmt.Fprintln(w, "# TYPE ip_probe_clock_offset_ms gauge")
	if s.Gauges.ClockOffsetMs != nil {
		fmt.Fprintf(w, "ip_probe_clock_offset_ms %d\n", s.Gauges.ClockOffsetMs())
	}
	fmt.Fprintln(w, "# TYPE ip_probe_targets gauge")
	if s.Gauges.Targets != nil {
		fmt.Fprintf(w, "ip_probe_targets %d\n", s.Gauges.Targets())
	}
	fmt.Fprintln(w, "# TYPE ip_probe_uptime_seconds gauge")
	fmt.Fprintf(w, "ip_probe_uptime_seconds %d\n", s.Uptime())
	fmt.Fprintln(w, "# TYPE ip_probe_rss_mb gauge")
	fmt.Fprintf(w, "ip_probe_rss_mb %.1f\n", RSSMb())
	fmt.Fprintln(w, "# TYPE ip_probe_goroutines gauge")
	fmt.Fprintf(w, "ip_probe_goroutines %d\n", runtime.NumGoroutine())
	fmt.Fprintln(w, "# TYPE ip_probe_info gauge")
	fmt.Fprintf(w, "ip_probe_info{probe_id=%q,version=%q,os=%q,arch=%q} 1\n", s.ProbeID, s.Version, runtime.GOOS, runtime.GOARCH)
}

// Serve runs the local HTTP server on listen until ctx is cancelled.
func (s *State) Serve(ctx context.Context, listen string) error {
	ln, err := net.Listen("tcp", listen)
	if err != nil {
		return err
	}
	srv := &http.Server{Handler: s.Handler(), ReadHeaderTimeout: 5 * time.Second}
	go func() {
		<-ctx.Done()
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		srv.Shutdown(shutdownCtx)
	}()
	if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
		return err
	}
	return nil
}