package modelplan

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

// FetchModel reads a model's architecture from its repository.
//
// Read rather than asked for, because every field here is one somebody would
// otherwise have to look up and retype: the layer count, the KV head count, the
// head dimension and the parameter total. Those are exactly the numbers a
// person gets subtly wrong — attention heads instead of KV heads being the
// classic, and worth a factor of five on a model with grouped-query attention.
//
// Two documents: config.json for the architecture, and the repository index for
// the parameter count, which is the one number config.json does not carry.
func FetchModel(ctx context.Context, repo string) (Model, error) {
	base := "https://huggingface.co"
	cfg, err := fetchConfig(ctx, base, repo)
	if err != nil {
		return Model{}, err
	}
	m, err := cfg.toModel(repo)
	if err != nil {
		return Model{}, err
	}
	// Best effort. A repository that does not publish its parameter count is
	// still usable — the caller can supply one — and failing the whole lookup
	// over it would be worse than saying which number is missing.
	if n, err := fetchParamCount(ctx, base, repo); err == nil && n > 0 {
		m.Params = n
	}
	return m, nil
}

type hfConfig struct {
	NumHiddenLayers   int    `json:"num_hidden_layers"`
	NumAttentionHeads int    `json:"num_attention_heads"`
	NumKeyValueHeads  int    `json:"num_key_value_heads"`
	HiddenSize        int    `json:"hidden_size"`
	HeadDim           int    `json:"head_dim"`
	MaxPositions      int    `json:"max_position_embeddings"`
	ModelType         string `json:"model_type"`
}

func (c hfConfig) toModel(repo string) (Model, error) {
	if c.NumHiddenLayers == 0 {
		return Model{}, fmt.Errorf("%s: config.json has no num_hidden_layers; is it a text model?", repo)
	}
	m := Model{
		Name:       repo[strings.LastIndex(repo, "/")+1:],
		Layers:     c.NumHiddenLayers,
		MaxContext: c.MaxPositions,
	}

	// Grouped-query attention: when the KV head count is absent the model has
	// as many as it has attention heads, which is multi-head attention and the
	// older, more expensive arrangement.
	m.KVHeads = c.NumKeyValueHeads
	if m.KVHeads == 0 {
		m.KVHeads = c.NumAttentionHeads
	}

	// head_dim is explicit in newer configs and derived in older ones. Derived
	// from the attention heads, never the KV heads: the hidden size is split
	// across the query heads.
	m.HeadDim = c.HeadDim
	if m.HeadDim == 0 {
		if c.NumAttentionHeads == 0 {
			return Model{}, fmt.Errorf("%s: config.json gives neither head_dim nor num_attention_heads", repo)
		}
		m.HeadDim = c.HiddenSize / c.NumAttentionHeads
	}
	return m, nil
}

func fetchConfig(ctx context.Context, base, repo string) (hfConfig, error) {
	var c hfConfig
	body, err := get(ctx, base+"/"+repo+"/resolve/main/config.json")
	if err != nil {
		return c, fmt.Errorf("reading %s: %w", repo, err)
	}
	if err := json.Unmarshal(body, &c); err != nil {
		return c, fmt.Errorf("%s: config.json did not parse: %w", repo, err)
	}
	return c, nil
}

// fetchParamCount reads the parameter total from the repository's own index.
//
// Taken from the safetensors index rather than estimated from the architecture,
// because an estimate is wrong by whatever the embeddings and the tied weights
// come to, and that error lands directly on the weights figure — the largest
// term in the whole calculation.
func fetchParamCount(ctx context.Context, base, repo string) (int64, error) {
	body, err := get(ctx, base+"/api/models/"+repo)
	if err != nil {
		return 0, err
	}
	var info struct {
		Safetensors struct {
			Total int64 `json:"total"`
		} `json:"safetensors"`
	}
	if err := json.Unmarshal(body, &info); err != nil {
		return 0, err
	}
	return info.Safetensors.Total, nil
}

func get(ctx context.Context, url string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		// A private or gated repository is the common case here, and "404" on
		// its own sends people looking for a typo they did not make.
		return nil, fmt.Errorf("%s returned %d; a private or gated repository needs to be downloaded by hand",
			url, resp.StatusCode)
	}
	return io.ReadAll(io.LimitReader(resp.Body, 4<<20))
}