package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func sumOf(b []byte) string { h := sha256.Sum256(b); return hex.EncodeToString(h[:]) }

func fileServer(t *testing.T, body []byte) string {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write(body) }))
	t.Cleanup(srv.Close)
	return srv.URL
}

func TestDownloadChecksTheHash(t *testing.T) {
	good := []byte(strings.Repeat("a good model ", 1000))
	url := fileServer(t, good)
	dest := filepath.Join(t.TempDir(), "m", "x.safetensors")
	if err := downloadFile(context.Background(), url, dest, int64(len(good)), sumOf(good), func(int64) {}); err != nil {
		t.Fatal(err)
	}
	if b, _ := os.ReadFile(dest); string(b) != string(good) {
		t.Fatal("the file is not what was served")
	}

	// the right size but other content: refused, and nothing is left behind as if it were a model
	other := []byte(strings.Repeat("a BAD model!! ", 1000))[:len(good)]
	url = fileServer(t, other)
	dest2 := filepath.Join(t.TempDir(), "x.safetensors")
	err := downloadFile(context.Background(), url, dest2, int64(len(good)), sumOf(good), func(int64) {})
	if err == nil || !strings.Contains(err.Error(), "damaged") {
		t.Fatalf("a damaged download must be refused: %v", err)
	}
	if _, e := os.Stat(dest2); e == nil {
		t.Error("the damaged file was kept")
	}
	if _, e := os.Stat(dest2 + ".part"); e == nil {
		t.Error("the damaged .part was kept")
	}
}

func TestDownloadReplacesAnExistingDamagedFile(t *testing.T) {
	good := []byte(strings.Repeat("a good model ", 1000))
	url := fileServer(t, good)
	dest := filepath.Join(t.TempDir(), "x.safetensors")
	damaged := []byte(strings.Repeat("a BAD model!! ", 1000))[:len(good)] // same size, other content
	os.WriteFile(dest, damaged, 0o644)
	if err := downloadFile(context.Background(), url, dest, int64(len(good)), sumOf(good), func(int64) {}); err != nil {
		t.Fatal(err)
	}
	if b, _ := os.ReadFile(dest); string(b) != string(good) {
		t.Fatal("the damaged file was not replaced")
	}
}

func TestDownloadLeavesAMatchingFileAlone(t *testing.T) {
	good := []byte(strings.Repeat("a good model ", 1000))
	called := false
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called = true }))
	defer srv.Close()
	dest := filepath.Join(t.TempDir(), "x.safetensors")
	os.WriteFile(dest, good, 0o644)
	if err := downloadFile(context.Background(), srv.URL, dest, int64(len(good)), sumOf(good), func(int64) {}); err != nil || called {
		t.Fatalf("a good file must not be fetched again: err=%v called=%v", err, called)
	}
	// without a published hash only the size counts, as before
	if err := downloadFile(context.Background(), srv.URL, dest, int64(len(good)), "", func(int64) {}); err != nil || called {
		t.Fatalf("size only: err=%v called=%v", err, called)
	}
}