package health

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/url"
)

// modelProps is the part of a /props body that carries a model's context size.
type modelProps struct {
	Generation struct {
		NCtx *int `json:"n_ctx"`
	} `json:"default_generation_settings"`
	TotalSlots *int `json:"total_slots"`
}

// props reads <base>/props best-effort. A request that fails because ctx is
// done yields a cancelled result so the caller records nothing; any other
// outcome (status, body, or missing fields) leaves context unknown without
// failing the poll. When the body says "role":"router" the host carries no
// context of its own, so NCtx and Slots stay 0 whatever else it reports.
func (t *Table) props(ctx context.Context, base string) (int, int, pollResult) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props", nil)
	if err != nil {
		if ctx.Err() != nil {
			return 0, 0, pollResult{cancelled: true}
		}
		return 0, 0, pollResult{}
	}
	resp, err := t.client.Do(req)
	if err != nil {
		if ctx.Err() != nil {
			return 0, 0, pollResult{cancelled: true}
		}
		return 0, 0, pollResult{}
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return 0, 0, pollResult{}
	}
	var p struct {
		Role       string `json:"role"`
		Generation struct {
			NCtx *int `json:"n_ctx"`
		} `json:"default_generation_settings"`
		TotalSlots *int `json:"total_slots"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil {
		return 0, 0, pollResult{}
	}
	if p.Role == "router" {
		return 0, 0, pollResult{}
	}
	nctx, slots := 0, 0
	if p.Generation.NCtx != nil {
		nctx = *p.Generation.NCtx
	}
	if p.TotalSlots != nil {
		slots = *p.TotalSlots
	}
	if nctx < 0 {
		nctx = 0
	}
	if slots < 0 {
		slots = 0
	}
	return nctx, slots, pollResult{}
}

// propsModels asks /props?model= for each loaded model and returns a fresh,
// non-nil map of what each answered. A failed or malformed answer leaves that
// id absent; a request that fails because ctx is done cancels the whole poll.
// No model is ever asked that is not in loaded.
func (t *Table) propsModels(ctx context.Context, base string, loaded []string) (map[string]ModelCtx, pollResult) {
	models := make(map[string]ModelCtx, len(loaded))
	for _, id := range loaded {
		mc, ok, r := t.propsModel(ctx, base, id)
		if r.cancelled || r.reason != "" {
			return nil, r
		}
		if !ok {
			continue
		}
		models[id] = mc
	}
	return models, pollResult{}
}

// propsModel reads <base>/props?model=<id> best-effort. ok is false when the
// answer failed or was malformed (the model stays unknown); a request that
// fails because ctx is done yields a cancelled result so the caller records
// nothing.
func (t *Table) propsModel(ctx context.Context, base, id string) (ModelCtx, bool, pollResult) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/props?model="+url.QueryEscape(id), nil)
	if err != nil {
		if ctx.Err() != nil {
			return ModelCtx{}, false, pollResult{cancelled: true}
		}
		return ModelCtx{}, false, pollResult{}
	}
	resp, err := t.client.Do(req)
	if err != nil {
		if ctx.Err() != nil {
			return ModelCtx{}, false, pollResult{cancelled: true}
		}
		return ModelCtx{}, false, pollResult{}
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return ModelCtx{}, false, pollResult{}
	}
	var p modelProps
	if err := json.NewDecoder(io.LimitReader(resp.Body, MaxModelsBody)).Decode(&p); err != nil {
		return ModelCtx{}, false, pollResult{}
	}
	mc, r := clampCtx(p.Generation.NCtx, p.TotalSlots)
	return mc, true, r
}

// clampCtx turns the two optional fields into a non-negative context and slots.
func clampCtx(nctx, slots *int) (ModelCtx, pollResult) {
	c := ModelCtx{}
	if nctx != nil {
		c.NCtx = *nctx
	}
	if slots != nil {
		c.Slots = *slots
	}
	if c.NCtx < 0 {
		c.NCtx = 0
	}
	if c.Slots < 0 {
		c.Slots = 0
	}
	return c, pollResult{}
}