//go:build ai

package main

import (
	"context"
	"io"
	"net/http"
	"os"
	"path"
	"testing"
	"time"
)

func TestAILocal(t *testing.T) {
	ctx, can := context.WithTimeout(context.Background(), time.Minute)
	defer can()

	d := os.TempDir()
	for k, u := range map[string]string{
		"checkpoints": "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin",
		"tokenizer":   "https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin",
	} {
		func() {
			if _, err := os.Stat(path.Join(d, k)); os.IsNotExist(err) {
				t.Logf("downloading %s from %s", u, k)

				resp, err := http.Get(u)
				if err != nil {
					t.Fatal(err)
				}
				defer resp.Body.Close()

				f, err := os.Create(path.Join(d, k))
				if err != nil {
					t.Fatal(err)
				}
				defer f.Close()

				if _, err := io.Copy(f, resp.Body); err != nil {
					f.Close()
					os.Remove(path.Join(d, k))
					t.Fatal(err)
				}
			}
		}()
	}

	ai := NewAILocal(
		path.Join(d, "checkpoints"),
		path.Join(d, "tokenizer"),
		0.9,
		256,
		0.9,
	)

	if result, err := ai.Do(ctx, "hello world"); err != nil {
		t.Fatal(err)
	} else {
		t.Logf("%s", result)
	}
}