package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"slices"
	"strings"
	"unicode"
)

// Models lists the models the server offers, by name.
func (c *Client) Models(ctx context.Context) ([]string, error) {
	var list struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := c.getJSON(ctx, c.BaseURL+"/models", &list); err != nil {
		return nil, fmt.Errorf("listing models: %w", err)
	}
	var names []string
	for _, m := range list.Data {
		// A name is drawn on screen and saved as the model to use: one
		// with terminal controls in it is not a name to offer.
		if m.ID != "" && !strings.ContainsFunc(m.ID, func(r rune) bool { return unicode.IsControl(r) || unicode.Is(unicode.Cf, r) }) {
			names = append(names, m.ID)
		}
	}
	slices.Sort(names)
	return names, nil
}

// maxContext bounds the context sizes believed from a server.
const maxContext = 100_000_000

// ContextSize asks the server how many tokens the model can take in,
// prompt and reply together: vLLM and SGLang list it with the model, and
// llama.cpp gives it in /props, beside the API. It returns 0 if the server
// does not say.
func (c *Client) ContextSize(ctx context.Context) (int, error) {
	var list struct {
		Data []struct {
			ID          string `json:"id"`
			MaxModelLen int    `json:"max_model_len"`
		} `json:"data"`
	}
	if err := c.getJSON(ctx, c.BaseURL+"/models", &list); err != nil {
		return 0, err
	}
	for _, m := range list.Data {
		if m.ID == c.Model && m.MaxModelLen > 0 && m.MaxModelLen <= maxContext {
			return m.MaxModelLen, nil
		}
	}
	var props struct {
		Settings struct {
			NCtx int `json:"n_ctx"`
		} `json:"default_generation_settings"`
	}
	if err := c.getJSON(ctx, strings.TrimSuffix(c.BaseURL, "/v1")+"/props", &props); err != nil {
		return 0, nil // not llama.cpp
	}
	if n := props.Settings.NCtx; n > 0 && n <= maxContext {
		return n, nil
	}
	return 0, nil
}

// getJSON reads a JSON document from the server.
func (c *Client) getJSON(ctx context.Context, url string, v any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	if c.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 10<<20))
	if err != nil {
		return err
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("%s: status %d: %.200s", url, resp.StatusCode, body)
	}
	if err := json.Unmarshal(body, v); err != nil {
		return fmt.Errorf("%s is not the JSON expected: %w", url, err)
	}
	return nil
}