package web

import (
	"context"
	"errors"
	"slices"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
	"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
)

// The mirror page's Strata tab: the models Strata runs, each size with what
// it downloads and keeps in RAM and where it would run here, and a button
// that fetches a size's files into the mirror, where the version's serve
// form then serves it on Strata.

type strataView struct {
	// Nodes are the nodes that have Strata installed.
	Nodes    []string
	Families []strataFamilyRow
}

type strataFamilyRow struct {
	Name, Title, Repo string
	// Fetchable is a family whose files llmgw fetches into the mirror;
	// FetchesItself one whose files Strata fetches onto the node, served
	// from here at once.
	Fetchable, FetchesItself bool
	Sizes                    []strataSizeRow
}

type strataSizeRow struct {
	Size              string
	DownloadGB, RAMGB string
	Runs              string
	Fits, NVIDIAOnly  bool
}

func (p *AccountPages) strataView(l *Lang) *strataView {
	v := &strataView{}
	var pool []modelplan.PoolNode
	if fit := p.fitOf(); fit != nil {
		pool = fit.pool
	}
	for _, n := range pool {
		if slices.Contains(n.Engines, "strata") {
			v.Nodes = append(v.Nodes, n.ID)
		}
	}
	for _, f := range strata.Families {
		row := strataFamilyRow{Name: f.Name, Title: f.Title, Repo: f.Repo, Fetchable: f.Fetchable(),
			FetchesItself: f.FetchesItself()}
		for _, s := range strata.SizesOf(f.Name) {
			sr := strataSizeRow{Size: s.Name, DownloadGB: l.Num(s.DownloadGB, 0), RAMGB: l.Int(int64(s.MemoryGB())),
				NVIDIAOnly: s.NVIDIAOnly}
			opts := modelplan.StrataOptions(f.Name, []modelplan.Artifact{{Name: s.Name + ".gguf", Format: "gguf",
				Bytes: int64(s.DownloadGB * 1e9)}}, pool)
			switch {
			case len(opts) > 0:
				sr.Fits, sr.Runs = true, l.T("explore.strata", opts[0].Class, l.Num(opts[0].NeedGB, 0))
			case len(v.Nodes) == 0:
				sr.Runs = l.T("strata.no_node")
			case f.MinCC > 0 && cardsOlderThan(pool, f.MinCC):
				sr.Runs = l.T("strata.old_card", l.Num(f.MinCC, 1))
			default:
				sr.Runs = l.T("strata.no_room", s.MemoryGB())
			}
			row.Sizes = append(row.Sizes, sr)
		}
		v.Families = append(v.Families, row)
	}
	return v
}

// cardsOlderThan says whether every node with Strata has a card of a
// compute capability below cc, as far as they say.
func cardsOlderThan(pool []modelplan.PoolNode, cc float64) bool {
	older := false
	for _, n := range pool {
		if !slices.Contains(n.Engines, "strata") {
			continue
		}
		if n.CC == 0 || n.CC >= cc {
			return false
		}
		older = true
	}
	return older
}

// serveStrata adds a model of a family Strata fetches itself: its source
// the repository, its replica on a card of the class the plan found.
func (p *AccountPages) serveStrata(ctx context.Context, l *Lang, by, family, size string) (string, error) {
	fam, ok := strata.FamilyNamed(family)
	s, err := strata.SizeOf(family, size)
	if !ok || err != nil {
		return "", errors.New(l.T("strata.err.unknown", family, size))
	}
	if !fam.FetchesItself() || p.catalog == nil {
		return "", errors.New(l.T("strata.err.not_servable", fam.Title))
	}
	var pool []modelplan.PoolNode
	if fit := p.fitOf(); fit != nil {
		pool = fit.pool
	}
	opts := modelplan.StrataOptions(fam.Name, []modelplan.Artifact{{Name: s.Name + ".gguf", Format: "gguf"}}, pool)
	if len(opts) == 0 {
		return "", errors.New(l.T("strata.no_room", s.MemoryGB()))
	}
	name := fam.Model + "-" + strings.ToLower(s.Name)
	disk := s.DownloadGB
	if s.DiskGB > disk {
		disk = s.DiskGB
	}
	m := catalog.Model{Name: name, Source: strata.RepoOf(fam, s), RoutingProfile: "chat", ContextLength: strataContext,
		Scaling: catalog.Scaling{Min: 1, Max: 1},
		Profiles: []catalog.Profile{{Name: "strata-" + opts[0].Class, Engine: "strata", GPUClass: opts[0].Class,
			GPUsPerNode: 1, Weight: 1, Priority: 100, MinMemoryGB: s.MemoryGB(), MinDiskGB: int(disk) + 10,
			Args: []string{"--family", fam.Name, "--size", s.Name, "--context", strconv.Itoa(strataContext), "--vision", "no"},
			// Its files come onto the node on its first start, tens of
			// gigabytes from the internet, before it is prepared.
			StartupTimeout: 4 * time.Hour}}}
	if err := p.catalog.Add(ctx, m, by); err != nil {
		return "", err
	}
	return l.T("strata.msg.served", name, l.Num(disk, 0)), nil
}

// strataContext is the context a Strata model is served with from here.
const strataContext = 32768

// fetchStrata fetches a Strata size's files from its family's repository
// into the mirror, each checked against the sha256 the registry published.
func (p *AccountPages) fetchStrata(ctx context.Context, l *Lang, m MirrorAdmin, family, size string) (string, error) {
	fam, ok := strata.FamilyNamed(family)
	if _, err := strata.SizeOf(family, size); !ok || err != nil {
		return "", errors.New(l.T("strata.err.unknown", family, size))
	}
	if !fam.Fetchable() || p.registry == nil {
		return "", errors.New(l.T("strata.err.not_fetchable", fam.Title))
	}
	repo, err := p.registry.Repo(ctx, fam.Repo)
	if err != nil {
		return "", err
	}
	var paths []string
	sums := map[string]string{}
	for _, f := range repo.Files {
		paths = append(paths, f.Path)
		sums[f.Path] = f.SHA256
	}
	files := strata.FilesOf(size, paths)
	if len(files) == 0 {
		return "", errors.New(l.T("strata.err.no_files", fam.Repo, size))
	}
	model := fam.Model + "-" + strings.ToLower(size)
	version := repo.ShortRevision()
	for i, path := range files {
		req := mirror.FetchRequest{Source: repo.DownloadURL("", path), Model: model, Version: version,
			File: mirror.LastPathElement(path), SHA256: sums[path]}
		if _, err := m.Fetch(ctx, req); err != nil && !strings.Contains(err.Error(), "already") {
			return "", errors.New(l.T("mirror.err.partly_started", i, len(files), req.File, l.Err(err)))
		}
	}
	return l.T("mirror.msg.started", model, version, l.N("mirror.files", int64(len(files)))), nil
}