// Package health tracks, for every host, whether it answered its last poll, which models it has
// loaded, and when it last answered. The proxy reads the table to choose a host and to record a
// request failure against a host.
package health

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sort"
	"sync"
	"time"
)

// RecoveryPolls is the number of good polls in a row a host needs after a failure before it is
// trusted again; a router loading a model flaps, so one good poll is not enough.
const RecoveryPolls = 2

// MaxModelsBody bounds how many bytes we read from either /health or /v1/models.
const MaxModelsBody = 1 << 20

// ModelCtx is what /props?model=X taught us about one loaded model.
type ModelCtx struct {
	NCtx  int `json:"n_ctx"`
	Slots int `json:"slots"`
}

// Status is a snapshot of one host's health, safe to copy.
type Status struct {
	Healthy     bool                `json:"healthy"`
	Loaded      []string            `json:"loaded"`      // sorted, unique model ids from the last good poll
	LastOK      time.Time           `json:"last_ok"`     // zero if never
	LastErr     string              `json:"last_err"`    // "" after a good poll
	Consecutive int                 `json:"consecutive"` // good polls in a row
	NCtx        int                 `json:"n_ctx"`       // total context from /props; 0 = unknown (a router's own /props carries none)
	Slots       int                 `json:"slots"`       // total_slots from /props; 0 = unknown
	Models      map[string]ModelCtx `json:"models"`      // per loaded model; never nil after a poll
}

// PerSlotCtx is the context one request may use: NCtx divided by Slots, or the
// whole NCtx when Slots is unknown (0). It is 0 when NCtx is unknown.
func (s Status) PerSlotCtx() int {
	if s.NCtx == 0 || s.Slots == 0 {
		return s.NCtx
	}
	return s.NCtx / s.Slots
}

// PerSlotCtxFor is the per-slot context for one model on this host: Models[model]
// when present (NCtx/Slots, 0 when either is 0); else, when model is in Loaded,
// the host-level PerSlotCtx(); else 0 ("unknown" / not resident).
func (s Status) PerSlotCtxFor(model string) int {
	if mc, ok := s.Models[model]; ok {
		if mc.NCtx == 0 || mc.Slots == 0 {
			return 0
		}
		return mc.NCtx / mc.Slots
	}
	if contains(s.Loaded, model) {
		return s.PerSlotCtx()
	}
	return 0
}

type entry struct {
	status     Status
	everFailed bool
}

type pollResult struct {
	ok        bool
	cancelled bool
	reason    string
	loaded    []string
	nctx      int
	slots     int
	models    map[string]ModelCtx
}

// Table maps a host name to its health status. All methods are safe for concurrent use.
type Table struct {
	mu       sync.Mutex
	baseURL  map[string]string
	entries  map[string]entry
	client   *http.Client
	interval time.Duration
}

// New builds a table from hosts (name -> base URL, no trailing slash). A nil client becomes a
// five-second-timeout client. Every host starts untrusted with an empty Loaded slice.
func New(hosts map[string]string, interval time.Duration, client *http.Client) *Table {
	if client == nil {
		client = &http.Client{Timeout: 5 * time.Second}
	}
	t := &Table{
		baseURL:  hosts,
		entries:  make(map[string]entry, len(hosts)),
		client:   client,
		interval: interval,
	}
	for name := range hosts {
		t.entries[name] = entry{status: Status{Loaded: []string{}}}
	}
	return t
}

// Run polls once immediately, then every interval, and returns when ctx is done.
func (t *Table) Run(ctx context.Context) {
	t.PollOnce(ctx)
	ticker := time.NewTicker(t.interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			t.PollOnce(ctx)
		}
	}
}

// PollOnce polls every host concurrently and returns when all have finished.
func (t *Table) PollOnce(ctx context.Context) {
	t.mu.Lock()
	names := make([]string, 0, len(t.baseURL))
	for name := range t.baseURL {
		names = append(names, name)
	}
	t.mu.Unlock()

	var wg sync.WaitGroup
	for _, name := range names {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			t.pollHost(ctx, name)
		}(name)
	}
	wg.Wait()
}

// Get returns a copy of one host's status; ok is false for an unknown name.
func (t *Table) Get(name string) (Status, bool) {
	t.mu.Lock()
	defer t.mu.Unlock()
	e, ok := t.entries[name]
	if !ok {
		return Status{}, false
	}
	return copyStatus(e.status), true
}

// All returns copies of every host's status.
func (t *Table) All() map[string]Status {
	t.mu.Lock()
	defer t.mu.Unlock()
	out := make(map[string]Status, len(t.entries))
	for name, e := range t.entries {
		out[name] = copyStatus(e.status)
	}
	return out
}

// MarkDown records a failure seen by the proxy. Loaded is left as last seen.
func (t *Table) MarkDown(name, reason string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	e, ok := t.entries[name]
	if !ok {
		return
	}
	e.everFailed = true
	e.status.Healthy = false
	e.status.Consecutive = 0
	e.status.LastErr = "marked down: " + reason
	t.entries[name] = e
}

func (t *Table) pollHost(ctx context.Context, name string) {
	base := t.baseURL[name]
	r := t.poll(ctx, base)

	t.mu.Lock()
	defer t.mu.Unlock()
	e := t.entries[name]
	if r.cancelled {
		return
	}
	if r.ok {
		e.status.Consecutive++
		e.status.LastOK = time.Now()
		e.status.LastErr = ""
		e.status.Loaded = r.loaded
		e.status.NCtx = r.nctx
		e.status.Slots = r.slots
		e.status.Models = r.models
		e.status.Healthy = !e.everFailed || e.status.Consecutive >= RecoveryPolls
	} else {
		e.everFailed = true
		e.status.Healthy = false
		e.status.Consecutive = 0
		e.status.LastErr = r.reason
	}
	t.entries[name] = e
}

// poll runs the two requests for one host. A request that fails because ctx is done yields a
// cancelled result so the caller records nothing; any other failure yields a reason.
func (t *Table) poll(ctx context.Context, base string) pollResult {
	loaded, r := t.check(ctx, base+"/health", "health")
	if r.cancelled || r.reason != "" {
		return r
	}
	loaded, r = t.check(ctx, base+"/v1/models", "models")
	if r.cancelled || r.reason != "" {
		return r
	}
	nctx, slots, r := t.props(ctx, base)
	if r.cancelled || r.reason != "" {
		return r
	}
	models, r := t.propsModels(ctx, base, loaded)
	if r.cancelled || r.reason != "" {
		return r
	}
	return pollResult{ok: true, loaded: loaded, nctx: nctx, slots: slots, models: models}
}

// check performs one GET and, on success, returns the decoded model ids. Health checks use the
// status code only; the models check also decodes the resident model ids.
func (t *Table) check(ctx context.Context, url, prefix string) ([]string, pollResult) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, t.fail(ctx, prefix, err)
	}
	resp, err := t.client.Do(req)
	if err != nil {
		return nil, t.fail(ctx, prefix, err)
	}

	if prefix == "health" {
		// We only need the status code; bound and drain the body.
		func() {
			defer resp.Body.Close()
			_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, MaxModelsBody))
		}()
		if resp.StatusCode != http.StatusOK {
			return nil, pollResult{reason: fmt.Sprintf("health: HTTP %d", resp.StatusCode)}
		}
		return nil, pollResult{}
	}

	// models: bound and decode the body.
	if resp.StatusCode != http.StatusOK {
		defer resp.Body.Close()
		return nil, pollResult{reason: fmt.Sprintf("models: HTTP %d", resp.StatusCode)}
	}
	var m struct {
		Data []struct {
			ID     string `json:"id"`
			Status struct {
				Value string `json:"value"`
			} `json:"status"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&m); err != nil {
		defer resp.Body.Close()
		return nil, t.fail(ctx, "models", err)
	}
	defer resp.Body.Close()

	// A model is loaded when it has no status or status.value == "loaded"; any other
	// value ("unloaded", "loading", …) is not loaded and must never be asked /props?model=.
	loaded := make([]string, 0, len(m.Data))
	seen := make(map[string]struct{}, len(m.Data))
	for _, d := range m.Data {
		if d.ID == "" {
			continue
		}
		if _, ok := seen[d.ID]; ok {
			continue
		}
		if d.Status.Value != "" && d.Status.Value != "loaded" {
			continue
		}
		seen[d.ID] = struct{}{}
		loaded = append(loaded, d.ID)
	}
	sort.Strings(loaded)
	return loaded, pollResult{}
}

// fail builds a failure result, treating a done ctx as a cancellation rather than a failure.
func (t *Table) fail(ctx context.Context, prefix string, err error) pollResult {
	if ctx.Err() != nil {
		return pollResult{cancelled: true}
	}
	return pollResult{reason: prefix + ": " + err.Error()}
}

func copyStatus(s Status) Status {
	out := s
	if s.Loaded != nil {
		out.Loaded = make([]string, len(s.Loaded))
		copy(out.Loaded, s.Loaded)
	}
	if s.Models != nil {
		out.Models = make(map[string]ModelCtx, len(s.Models))
		for k, v := range s.Models {
			out.Models[k] = v
		}
	}
	return out
}

func contains(list []string, v string) bool {
	for _, s := range list {
		if s == v {
			return true
		}
	}
	return false
}