package model
|
|
import (
|
"fmt"
|
"io"
|
"net/http"
|
"os"
|
"os/exec"
|
"path/filepath"
|
"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 {
|
resp, err := http.Get(url)
|
if err != nil {
|
return err
|
}
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
return fmt.Errorf("HTTP %d: %s", resp.StatusCode, resp.Status)
|
}
|
|
totalBytes := resp.ContentLength
|
out, err := os.Create(destPath)
|
if err != nil {
|
return err
|
}
|
defer out.Close()
|
|
buf := make([]byte, 32*1024)
|
var downloaded int64
|
lastPercent := -1.0
|
|
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 {
|
break
|
}
|
return readErr
|
}
|
}
|
|
return nil
|
}
|
|
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
|
})
|
}
|
