package main

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
)

// sha256File is the hash of a whole file.
func sha256File(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

// downloadFile fetches url into dest, resuming a .part file, and checks the size at the end, and the sha256 when `sum` is
// given (a file of the right size can still be wrong: that is how a damaged model made noise instead of pictures). onBytes
// is called with the number of bytes just written. A file that is already complete (and matches `sum`) is left alone; one
// that does not match is removed and fetched again.
func downloadFile(ctx context.Context, url, dest string, size int64, sum string, onBytes func(int64)) error {
	if fi, err := os.Stat(dest); err == nil && fi.Size() == size {
		if sum == "" {
			return nil
		}
		if got, err := sha256File(dest); err == nil && got == sum {
			return nil
		}
		logf("%s is damaged (its sha256 is not the published one): fetching it again", filepath.Base(dest))
		if err := os.Remove(dest); err != nil {
			return err
		}
	}
	if err := os.MkdirAll(filepath.Dir(dest), 0755); err != nil {
		return err
	}
	part := dest + ".part"
	var have int64
	if fi, err := os.Stat(part); err == nil {
		have = fi.Size()
	}
	if have > size { // a leftover from another file
		os.Remove(part)
		have = 0
	}
	if have < size {
		r, err := hfRequest(ctx, url, have)
		if err != nil {
			return err
		}
		defer r.Body.Close()
		switch r.StatusCode {
		case 200, 206:
		case 401, 403:
			return fmt.Errorf("Hugging Face refused access to %s (put a token in /etc/wxch/hf_token)", filepath.Base(dest))
		default:
			return fmt.Errorf("download of %s failed: HTTP %d", filepath.Base(dest), r.StatusCode)
		}
		flag := os.O_CREATE | os.O_WRONLY | os.O_APPEND
		if r.StatusCode == 200 { // the server ignored Range: start over
			flag = os.O_CREATE | os.O_WRONLY | os.O_TRUNC
			have = 0
		}
		out, err := os.OpenFile(part, flag, 0644)
		if err != nil {
			return err
		}
		if have > 0 {
			onBytes(have)
		}
		buf := make([]byte, 1<<20)
		for {
			n, rerr := r.Body.Read(buf)
			if n > 0 {
				if _, werr := out.Write(buf[:n]); werr != nil {
					out.Close()
					return werr
				}
				onBytes(int64(n))
			}
			if rerr == io.EOF {
				break
			}
			if rerr != nil {
				out.Close()
				return fmt.Errorf("download of %s interrupted: %v (install again to resume)", filepath.Base(dest), rerr)
			}
		}
		if err := out.Close(); err != nil {
			return err
		}
	}
	if fi, err := os.Stat(part); err != nil || fi.Size() != size {
		os.Remove(part)
		return fmt.Errorf("%s has the wrong size after download (install again)", filepath.Base(dest))
	}
	if sum != "" {
		got, err := sha256File(part)
		if err != nil {
			return err
		}
		if got != sum {
			os.Remove(part)
			return fmt.Errorf("%s is damaged after download (its sha256 is not the published one); install again", filepath.Base(dest))
		}
	}
	return os.Rename(part, dest)
}