main.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

var (
	dataDir   = "/data"
	binDir    = "/app/bin"
	apiKey    = ""
	port      = "7860" // Hugging Face Spaces default port
)

func init() {
	if val, ok := os.LookupEnv("DATA_DIR"); ok {
		dataDir = val
	}
	if val, ok := os.LookupEnv("BIN_DIR"); ok {
		binDir = val
	}
	if val, ok := os.LookupEnv("HF_SAVE_API_KEY"); ok {
		apiKey = val
	} else {
		log.Println("WARNING: HF_SAVE_API_KEY environment variable is not set. Write access will not be authenticated.")
	}
	if val, ok := os.LookupEnv("PORT"); ok {
		port = val
	}

	// Ensure data directory exists
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		log.Fatalf("Failed to create data directory %s: %v", dataDir, err)
	}
}

func main() {
	// API key middleware for write operations
	authRequired := func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			if apiKey == "" {
				next(w, r)
				return
			}
			authHeader := r.Header.Get("Authorization")
			if !strings.HasPrefix(authHeader, "Bearer ") {
				http.Error(w, "Unauthorized: Missing or invalid Authorization header", http.StatusUnauthorized)
				return
			}
			token := strings.TrimPrefix(authHeader, "Bearer ")
			if token != apiKey {
				http.Error(w, "Unauthorized: Invalid API key", http.StatusUnauthorized)
				return
			}
			next(w, r)
		}
	}

	http.HandleFunc("/init", handleInit)
	http.HandleFunc("/bin/", handleServeBinaries)
	http.HandleFunc("/upload", authRequired(handleUpload))
	http.HandleFunc("/list", handleList)
	http.HandleFunc("/download", handleDownload)
	http.HandleFunc("/health", handleHealth)

	log.Printf("Server starting on port %s...", port)
	log.Printf("Data directory: %s", dataDir)
	log.Printf("Binaries directory: %s", binDir)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("OK"))
}

// Serves the install script for various platforms
func handleInit(w http.ResponseWriter, r *http.Request) {
	platform := r.URL.Query().Get("platform")
	host := r.Host
	scheme := "https"
	if r.TLS == nil && !strings.Contains(host, "hf.space") {
		scheme = "http"
	}
	serverURL := fmt.Sprintf("%s://%s", scheme, host)

	// If token isn't passed, CLI will prompt the user to input the pre-shared key during setup
	token := r.URL.Query().Get("token")

	w.Header().Set("Content-Type", "text/plain")

	switch strings.ToLower(platform) {
	case "windows":
		// PowerShell script
		script := fmt.Sprintf(`$ServerUrl = "%s"
$Token = "%s"
$BinDir = "$HOME\.local\bin"
if (!(Test-Path $BinDir)) { New-Item -ItemType Directory -Force -Path $BinDir }
$ExePath = "$BinDir\hf-save.exe"
Write-Host "Downloading hf-save CLI..." -ForegroundColor Cyan
Invoke-WebRequest -Uri "$ServerUrl/bin/windows-amd64" -OutFile $ExePath
if (!(($env:Path -split ';') -contains $BinDir)) {
    [System.Environment]::SetEnvironmentVariable("Path", $env:Path + ";$BinDir", "User")
    $env:Path += ";$BinDir"
    Write-Host "Added $BinDir to PATH. You may need to restart your terminal." -ForegroundColor Yellow
}
$ConfigDir = "$HOME\.config\hf-save"
if (!(Test-Path $ConfigDir)) { New-Item -ItemType Directory -Force -Path $ConfigDir }
if ($Token -eq "") {
    $Token = Read-Host "Enter your HF_SAVE_API_KEY (Pre-shared API Key)"
}
$Config = @{
    server_url = $ServerUrl
    api_key = $Token
} | ConvertTo-Json
$Config | Out-File -FilePath "$ConfigDir\config.json" -Encoding utf8
Write-Host "Installation successful! Try running: hf-save --help" -ForegroundColor Green
`, serverURL, token)
		w.Write([]byte(script))

	case "mac", "darwin":
		// Bash/Zsh script for macOS
		script := fmt.Sprintf(`#!/bin/bash
set -e
SERVER_URL="%s"
TOKEN="%s"
BIN_DIR="$HOME/.local/bin"
mkdir -p "$BIN_DIR"
# Determine CPU architecture
ARCH="amd64"
if [[ "$(uname -m)" == "arm64" ]]; then
    ARCH="arm64"
fi
echo "Downloading hf-save CLI for macOS ($ARCH)..."
curl -fsSL "$SERVER_URL/bin/darwin-$ARCH" -o "$BIN_DIR/hf-save"
chmod +x "$BIN_DIR/hf-save"

# Add to path helper
if [[ ":$PATH:" != *":$BIN_DIR:"* ]]; then
    echo "Adding $BIN_DIR to PATH in shell profile..."
    if [[ "$SHELL" == */zsh ]]; then
        echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$HOME/.zshrc"
    else
        echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$HOME/.bashrc"
    fi
    export PATH="$BIN_DIR:$PATH"
fi

CONFIG_DIR="$HOME/.config/hf-save"
mkdir -p "$CONFIG_DIR"
if [ -z "$TOKEN" ]; then
    read -p "Enter your HF_SAVE_API_KEY (Pre-shared API Key): " TOKEN
fi
cat << EOF > "$CONFIG_DIR/config.json"
{
  "server_url": "$SERVER_URL",
  "api_key": "$TOKEN"
}
EOF
echo "Installation successful! Try running: hf-save --help (Or reload your shell profile)"
`, serverURL, token)
		w.Write([]byte(script))

	default: // Default to linux
		script := fmt.Sprintf(`#!/bin/bash
set -e
SERVER_URL="%s"
TOKEN="%s"
BIN_DIR="$HOME/.local/bin"
mkdir -p "$BIN_DIR"
# Determine CPU architecture
ARCH="amd64"
if [[ "$(uname -m)" == "aarch64" ]]; then
    ARCH="arm64"
fi
echo "Downloading hf-save CLI for Linux ($ARCH)..."
curl -fsSL "$SERVER_URL/bin/linux-$ARCH" -o "$BIN_DIR/hf-save"
chmod +x "$BIN_DIR/hf-save"

# Add to path helper
if [[ ":$PATH:" != *":$BIN_DIR:"* ]]; then
    echo "Adding $BIN_DIR to PATH in shell profile..."
    echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$HOME/.bashrc"
    export PATH="$BIN_DIR:$PATH"
fi

CONFIG_DIR="$HOME/.config/hf-save"
mkdir -p "$CONFIG_DIR"
if [ -z "$TOKEN" ]; then
    read -p "Enter your HF_SAVE_API_KEY (Pre-shared API Key): " TOKEN
fi
cat << EOF > "$CONFIG_DIR/config.json"
{
  "server_url": "$SERVER_URL",
  "api_key": "$TOKEN"
}
EOF
echo "Installation successful! Try running: hf-save --help (Or run: export PATH=\$HOME/.local/bin:\$PATH)"
`, serverURL, token)
		w.Write([]byte(script))
	}
}

// Serves compiled CLI binaries
func handleServeBinaries(w http.ResponseWriter, r *http.Request) {
	filename := strings.TrimPrefix(r.URL.Path, "/bin/")
	if filename == "" {
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	// Sanitize filename
	filename = filepath.Base(filename)
	binaryPath := filepath.Join(binDir, filename)

	if _, err := os.Stat(binaryPath); os.IsNotExist(err) {
		// Fallback for names containing extensions like .exe
		if strings.HasSuffix(filename, ".exe") {
			binaryPath = filepath.Join(binDir, "windows-amd64")
		} else if strings.Contains(filename, "linux") {
			binaryPath = filepath.Join(binDir, "linux-amd64")
		} else if strings.Contains(filename, "mac") || strings.Contains(filename, "darwin") {
			binaryPath = filepath.Join(binDir, "darwin-amd64")
		}
	}

	if _, err := os.Stat(binaryPath); os.IsNotExist(err) {
		http.Error(w, "Binary not found on server", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filename))
	w.Header().Set("Content-Type", "application/octet-stream")
	http.ServeFile(w, r, binaryPath)
}

// Handles multipart file uploads (Streams files straight to directory structure)
func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method Not Allowed", http.StatusMethodNotAllowed)
		return
	}

	mr, err := r.MultipartReader()
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to read multipart data: %v", err), http.StatusBadRequest)
		return
	}

	var currentPrefix string
	fileCount := 0

	for {
		part, err := mr.NextPart()
		if err == io.EOF {
			break
		}
		if err != nil {
			http.Error(w, fmt.Sprintf("Error parsing part: %v", err), http.StatusInternalServerError)
			return
		}

		formName := part.FormName()

		if formName == "prefix" {
			prefixVal, err := io.ReadAll(part)
			if err != nil {
				http.Error(w, "Failed to read prefix", http.StatusBadRequest)
				return
			}
			currentPrefix = string(prefixVal)
			part.Close()
			continue
		}

		if formName == "files" {
			if currentPrefix == "" {
				// Fallback to today's date if prefix wasn't provided first
				currentPrefix = time.Now().Format("2006-01-02/15-04-05_upload")
			}

			// Clean prefix
			cleanPrefix := filepath.Clean(currentPrefix)
			// Prevent directory traversal attacks
			if strings.Contains(cleanPrefix, "..") || strings.HasPrefix(cleanPrefix, "/") {
				http.Error(w, "Invalid prefix path", http.StatusBadRequest)
				return
			}

			// The file's relative path is stored in the header's filename
			relPath := part.FileName()
			if relPath == "" {
				part.Close()
				continue
			}

			// Clean and resolve local target path
			cleanRelPath := filepath.Clean(relPath)
			if strings.Contains(cleanRelPath, "..") || strings.HasPrefix(cleanRelPath, "/") {
				part.Close()
				continue
			}

			targetPath := filepath.Join(dataDir, cleanPrefix, cleanRelPath)
			targetDir := filepath.Dir(targetPath)

			if err := os.MkdirAll(targetDir, 0755); err != nil {
				http.Error(w, fmt.Sprintf("Failed to create target directory: %v", err), http.StatusInternalServerError)
				return
			}

			dst, err := os.Create(targetPath)
			if err != nil {
				http.Error(w, fmt.Sprintf("Failed to create file: %v", err), http.StatusInternalServerError)
				return
			}

			if _, err := io.Copy(dst, part); err != nil {
				dst.Close()
				http.Error(w, fmt.Sprintf("Failed to save file: %v", err), http.StatusInternalServerError)
				return
			}
			dst.Close()
			part.Close()
			fileCount++
		}
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"success":    true,
		"file_count": fileCount,
		"prefix":     currentPrefix,
	})
}

// Lists dates or runs within a date
func handleList(w http.ResponseWriter, r *http.Request) {
	dateParam := r.URL.Query().Get("date")

	var results []string

	if dateParam == "" {
		// List top-level date directories
		entries, err := os.ReadDir(dataDir)
		if err != nil {
			http.Error(w, fmt.Sprintf("Failed to read data directory: %v", err), http.StatusInternalServerError)
			return
		}
		for _, entry := range entries {
			if entry.IsDir() {
				// Basic check for YYYY-MM-DD pattern
				name := entry.Name()
				if len(name) == 10 && name[4] == '-' && name[7] == '-' {
					results = append(results, name)
				}
			}
		}
	} else {
		// List upload runs inside the specified date
		cleanDate := filepath.Clean(dateParam)
		if strings.Contains(cleanDate, "..") || strings.HasPrefix(cleanDate, "/") {
			http.Error(w, "Invalid date format", http.StatusBadRequest)
			return
		}

		targetPath := filepath.Join(dataDir, cleanDate)
		entries, err := os.ReadDir(targetPath)
		if err != nil {
			if os.IsNotExist(err) {
				w.Header().Set("Content-Type", "application/json")
				json.NewEncoder(w).Encode([]string{})
				return
			}
			http.Error(w, fmt.Sprintf("Failed to read date directory: %v", err), http.StatusInternalServerError)
			return
		}

		for _, entry := range entries {
			if entry.IsDir() {
				results = append(results, entry.Name())
			}
		}
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(results)
}

// Download file stream or metadata list
func handleDownload(w http.ResponseWriter, r *http.Request) {
	filePath := r.URL.Query().Get("file")
	snapshotPath := r.URL.Query().Get("snapshot")

	if filePath != "" {
		// Stream a specific file
		cleanPath := filepath.Clean(filePath)
		if strings.Contains(cleanPath, "..") || strings.HasPrefix(cleanPath, "/") {
			http.Error(w, "Invalid file path", http.StatusBadRequest)
			return
		}

		fullPath := filepath.Join(dataDir, cleanPath)
		if _, err := os.Stat(fullPath); os.IsNotExist(err) {
			http.Error(w, "File not found", http.StatusNotFound)
			return
		}

		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filepath.Base(fullPath)))
		w.Header().Set("Content-Type", "application/octet-stream")
		http.ServeFile(w, r, fullPath)
		return
	}

	if snapshotPath != "" {
		// List all files within a snapshot to download
		cleanSnapshot := filepath.Clean(snapshotPath)
		if strings.Contains(cleanSnapshot, "..") || strings.HasPrefix(cleanSnapshot, "/") {
			http.Error(w, "Invalid snapshot path", http.StatusBadRequest)
			return
		}

		fullDir := filepath.Join(dataDir, cleanSnapshot)
		var files []string

		err := filepath.Walk(fullDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() {
				rel, err := filepath.Rel(fullDir, path)
				if err == nil {
					// Convert windows backslashes to slash for web transport compatibility
					files = append(files, filepath.ToSlash(rel))
				}
			}
			return nil
		})

		if err != nil {
			if os.IsNotExist(err) {
				http.Error(w, "Snapshot not found", http.StatusNotFound)
				return
			}
			http.Error(w, fmt.Sprintf("Failed to scan snapshot: %v", err), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(files)
		return
	}

	http.Error(w, "Missing query parameter 'file' or 'snapshot'", http.StatusBadRequest)
}