package model
|
|
import (
|
"errors"
|
"fmt"
|
"io"
|
"net/http"
|
"os"
|
"os/exec"
|
"path/filepath"
|
"strconv"
|
"strings"
|
"time"
|
|
"voicesnap/internal/logger"
|
)
|
|
// ProgressCallback is called with download progress (percent, downloaded bytes, total bytes).
|
type ProgressCallback func(percent float64, downloaded, total int64)
|
|
// Download downloads and extracts the ASR model.
|
// It tries the primary URL first, then falls back to the fallback URL.
|
func Download(primaryURL, fallbackURL, modelsDir string, progress ProgressCallback) error {
|
profile := DefaultModelProfile()
|
urls := []string{primaryURL, fallbackURL}
|
return DownloadProfile(profile, urls, modelsDir, progress)
|
}
|
|
func DownloadProfile(profile ModelProfile, urls []string, modelsDir string, progress ProgressCallback) error {
|
if err := os.MkdirAll(modelsDir, 0755); err != nil {
|
return fmt.Errorf("failed to create models dir: %w", err)
|
}
|
cleanupStaleDownloadDirs(profile, modelsDir)
|
|
var lastErr error
|
for i, url := range urls {
|
if url == "" {
|
continue
|
}
|
if err := downloadAndInstallFromURL(profile, url, modelsDir, i+1, progress); err != nil {
|
lastErr = err
|
logger.Info("Model install from URL %d failed: %v", i+1, err)
|
continue
|
}
|
|
if err := UpdateInstalledModelState(modelsDir, profile, profile.InstallDirName, "new"); err != nil {
|
logger.Error("Failed to update model install state: %v", err)
|
}
|
logger.Info("Model download and extraction complete")
|
return nil
|
}
|
|
if lastErr == nil {
|
lastErr = fmt.Errorf("no download URL configured")
|
}
|
return fmt.Errorf("download failed: %w", lastErr)
|
}
|
|
func cleanupStaleDownloadDirs(profile ModelProfile, modelsDir string) {
|
downloadRoot := filepath.Join(modelsDir, ".downloads", profile.ID)
|
if err := os.RemoveAll(downloadRoot); err != nil {
|
logger.Info("Failed to remove stale download dir %s: %v", downloadRoot, err)
|
}
|
}
|
|
func downloadAndInstallFromURL(profile ModelProfile, url, modelsDir string, attempt int, progress ProgressCallback) error {
|
runID := fmt.Sprintf("%d-%d", time.Now().UnixNano(), attempt)
|
downloadDir := filepath.Join(modelsDir, ".downloads", profile.ID, runID)
|
extractDir := filepath.Join(downloadDir, "extract")
|
stagingDir := filepath.Join(modelsDir, ".staging", profile.ID+"-"+runID)
|
archivePath := filepath.Join(downloadDir, "model_package")
|
|
if err := os.MkdirAll(extractDir, 0755); err != nil {
|
return fmt.Errorf("failed to create download dir: %w", err)
|
}
|
defer os.RemoveAll(downloadDir)
|
defer os.RemoveAll(stagingDir)
|
|
logger.Info("Downloading model from URL %d: %s", attempt, url)
|
if err := downloadFile(url, archivePath, progress); err != nil {
|
return err
|
}
|
|
logger.Info("Extracting model archive...")
|
if err := extractArchive(archivePath, extractDir); err != nil {
|
return fmt.Errorf("extraction failed: %w", err)
|
}
|
|
sourceDir, err := findInstallSource(profile, extractDir)
|
if err != nil {
|
return err
|
}
|
|
if err := moveOrCopyDir(sourceDir, stagingDir); err != nil {
|
return fmt.Errorf("failed to stage model: %w", err)
|
}
|
|
validation := ValidateModelDir(profile, stagingDir)
|
if !validation.Valid {
|
return fmt.Errorf("downloaded model is incomplete: missing=%v problems=%v", validation.Missing, validation.Problems)
|
}
|
|
finalDir := filepath.Join(modelsDir, profile.InstallDirName)
|
return installStagedModel(profile, stagingDir, finalDir, runID)
|
}
|
|
func downloadFile(url, destPath string, progress ProgressCallback) error {
|
const maxAttempts = 6
|
var lastErr error
|
var lastPercent = -1.0
|
|
for attempt := 1; attempt <= maxAttempts; attempt++ {
|
downloaded := existingFileSize(destPath)
|
err := downloadFileAttempt(url, destPath, downloaded, &lastPercent, progress)
|
if err == nil {
|
return nil
|
}
|
if isNonRetryableDownloadError(err) {
|
return err
|
}
|
lastErr = err
|
logger.Info("Download attempt %d/%d failed: %v", attempt, maxAttempts, err)
|
if attempt < maxAttempts {
|
time.Sleep(time.Duration(attempt) * time.Second)
|
}
|
}
|
|
return fmt.Errorf("download failed after %d attempts: %w", maxAttempts, lastErr)
|
}
|
|
func downloadFileAttempt(url, destPath string, resumeFrom int64, lastPercent *float64, progress ProgressCallback) error {
|
req, err := http.NewRequest(http.MethodGet, url, nil)
|
if err != nil {
|
return nonRetryableDownloadError{err: err}
|
}
|
if resumeFrom > 0 {
|
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", resumeFrom))
|
}
|
|
resp, err := http.DefaultClient.Do(req)
|
if err != nil {
|
return err
|
}
|
defer resp.Body.Close()
|
|
appendMode := resumeFrom > 0 && resp.StatusCode == http.StatusPartialContent
|
if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
|
err := fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status)
|
if resp.StatusCode < http.StatusInternalServerError {
|
return nonRetryableDownloadError{err: err}
|
}
|
return err
|
}
|
if resumeFrom > 0 && resp.StatusCode == http.StatusOK {
|
logger.Info("Download server ignored Range request; restarting download")
|
resumeFrom = 0
|
appendMode = false
|
}
|
|
totalBytes := resp.ContentLength
|
if appendMode {
|
if total := parseContentRangeTotal(resp.Header.Get("Content-Range")); total > 0 {
|
totalBytes = total
|
} else if resp.ContentLength > 0 {
|
totalBytes = resumeFrom + resp.ContentLength
|
}
|
}
|
|
flag := os.O_CREATE | os.O_WRONLY
|
if appendMode {
|
flag |= os.O_APPEND
|
} else {
|
flag |= os.O_TRUNC
|
}
|
out, err := os.OpenFile(destPath, flag, 0644)
|
if err != nil {
|
return err
|
}
|
defer out.Close()
|
|
downloaded := resumeFrom
|
if !appendMode {
|
downloaded = 0
|
}
|
buf := make([]byte, 32*1024)
|
for {
|
n, readErr := resp.Body.Read(buf)
|
if n > 0 {
|
if _, writeErr := out.Write(buf[:n]); writeErr != nil {
|
return writeErr
|
}
|
downloaded += int64(n)
|
|
if totalBytes > 0 && progress != nil {
|
pct := float64(downloaded) / float64(totalBytes) * 100
|
if pct-*lastPercent >= 0.5 {
|
*lastPercent = pct
|
progress(pct, downloaded, totalBytes)
|
}
|
}
|
}
|
if readErr != nil {
|
if readErr == io.EOF {
|
if totalBytes > 0 && downloaded < totalBytes {
|
return fmt.Errorf("short download: got %d of %d bytes", downloaded, totalBytes)
|
}
|
return nil
|
}
|
return readErr
|
}
|
}
|
}
|
|
func existingFileSize(path string) int64 {
|
info, err := os.Stat(path)
|
if err != nil || info.IsDir() {
|
return 0
|
}
|
return info.Size()
|
}
|
|
func parseContentRangeTotal(value string) int64 {
|
if value == "" {
|
return 0
|
}
|
slash := strings.LastIndex(value, "/")
|
if slash < 0 || slash == len(value)-1 {
|
return 0
|
}
|
total, err := strconv.ParseInt(value[slash+1:], 10, 64)
|
if err != nil {
|
return 0
|
}
|
return total
|
}
|
|
type nonRetryableDownloadError struct {
|
err error
|
}
|
|
func (e nonRetryableDownloadError) Error() string {
|
return e.err.Error()
|
}
|
|
func (e nonRetryableDownloadError) Unwrap() error {
|
return e.err
|
}
|
|
func isNonRetryableDownloadError(err error) bool {
|
var target nonRetryableDownloadError
|
return errors.As(err, &target)
|
}
|
|
func extractArchive(archivePath, destDir string) error {
|
cmd := exec.Command("tar", "-xf", archivePath, "-C", destDir)
|
output, err := cmd.CombinedOutput()
|
if err != nil {
|
return fmt.Errorf("tar extraction failed: %v, output: %s", err, string(output))
|
}
|
return nil
|
}
|
|
func findInstallSource(profile ModelProfile, extractDir string) (string, error) {
|
if validation := ValidateModelDir(profile, extractDir); validation.Valid {
|
return extractDir, nil
|
}
|
|
var found string
|
err := filepath.WalkDir(extractDir, func(path string, entry os.DirEntry, walkErr error) error {
|
if walkErr != nil {
|
return walkErr
|
}
|
if found != "" || !entry.IsDir() {
|
return nil
|
}
|
if path == extractDir {
|
return nil
|
}
|
if validation := ValidateModelDir(profile, path); validation.Valid {
|
found = path
|
return filepath.SkipDir
|
}
|
return nil
|
})
|
if err != nil {
|
return "", err
|
}
|
if found == "" {
|
return "", fmt.Errorf("no complete %s model found in archive", profile.ID)
|
}
|
return found, nil
|
}
|
|
func installStagedModel(profile ModelProfile, stagingDir, finalDir, runID string) error {
|
var backupDir string
|
if _, err := os.Stat(finalDir); err == nil {
|
backupDir = finalDir + ".backup-" + runID
|
if err := os.Rename(finalDir, backupDir); err != nil {
|
return fmt.Errorf("failed to move existing model dir: %w", err)
|
}
|
}
|
|
if err := os.Rename(stagingDir, finalDir); err != nil {
|
if copyErr := copyDir(stagingDir, finalDir); copyErr != nil {
|
restoreBackup(finalDir, backupDir)
|
return fmt.Errorf("failed to install staged model: rename=%v copy=%w", err, copyErr)
|
}
|
}
|
|
if validation := ValidateModelDir(profile, finalDir); !validation.Valid {
|
restoreBackup(finalDir, backupDir)
|
return fmt.Errorf("installed model is incomplete: missing=%v problems=%v", validation.Missing, validation.Problems)
|
}
|
|
if backupDir != "" {
|
os.RemoveAll(backupDir)
|
}
|
return nil
|
}
|
|
func restoreBackup(finalDir, backupDir string) {
|
if backupDir == "" {
|
return
|
}
|
os.RemoveAll(finalDir)
|
os.Rename(backupDir, finalDir)
|
}
|
|
func moveOrCopyDir(src, dst string) error {
|
os.RemoveAll(dst)
|
if err := os.Rename(src, dst); err == nil {
|
return nil
|
}
|
return copyDir(src, dst)
|
}
|
|
func copyDir(src, dst string) error {
|
return filepath.WalkDir(src, func(path string, entry os.DirEntry, walkErr error) error {
|
if walkErr != nil {
|
return walkErr
|
}
|
rel, err := filepath.Rel(src, path)
|
if err != nil {
|
return err
|
}
|
target := filepath.Join(dst, rel)
|
info, err := entry.Info()
|
if err != nil {
|
return err
|
}
|
if entry.IsDir() {
|
return os.MkdirAll(target, info.Mode())
|
}
|
in, err := os.Open(path)
|
if err != nil {
|
return err
|
}
|
defer in.Close()
|
|
out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, info.Mode())
|
if err != nil {
|
return err
|
}
|
if _, err := io.Copy(out, in); err != nil {
|
out.Close()
|
return err
|
}
|
if err := out.Close(); err != nil {
|
return err
|
}
|
return nil
|
})
|
}
|
