counters.go

package main

import (
	"fmt"
	"math"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"
)

// Write method for countingWriter to track the number of bytes written
func (w *countingWriter) Write(data []byte) (int, error) {
	n, err := w.ResponseWriter.Write(data)
	w.bytesWritten += int64(n)
	return n, err
}

// normalizeStatsPath makes map keys consistent with API display paths (slash-separated).
func normalizeStatsPath(p string) string {
	return filepath.ToSlash(strings.TrimSpace(p))
}

func lookupDownloadCount(displayPath string) int64 {
	statsMutex.Lock()
	defer statsMutex.Unlock()
	key := normalizeStatsPath(displayPath)
	if s, ok := downloadStats[key]; ok {
		return s.Count
	}
	return 0
}

// Called after serving a file to update global stats and print download progress
func (w *countingWriter) finish() {
	var elapsed time.Duration
	if !w.startedAt.IsZero() {
		elapsed = time.Since(w.startedAt)
	}
	recordServedDownload(w.path, w.clientIP, w.bytesWritten, w.fileSize, w.isRangeRequest, elapsed)
}

func formatServedDuration(d time.Duration) string {
	if d < 0 {
		d = 0
	}
	if d < time.Second {
		return fmt.Sprintf("%.1fs", d.Seconds())
	}
	d = d.Round(time.Second)
	sec := int64(d / time.Second)
	if sec < 60 {
		return fmt.Sprintf("%ds", sec)
	}
	m := sec / 60
	s := sec % 60
	if s == 0 {
		return fmt.Sprintf("%dm", m)
	}
	return fmt.Sprintf("%dm %ds", m, s)
}

// formatServedDataSize uses decimal (SI) units: prefer GB only when >= 1 GB so e.g. 900 MB stays MB; whole GB when rounded.
func formatServedDataSize(bytes int64) string {
	if bytes < 0 {
		bytes = 0
	}
	const kb, mb, gb = 1000.0, 1e6, 1e9
	f := float64(bytes)
	if f >= gb {
		g := f / gb
		ig := int64(math.Round(g))
		if math.Abs(g-float64(ig)) < 0.0005 {
			return fmt.Sprintf("%dGB", ig)
		}
		return fmt.Sprintf("%.2fGB", g)
	}
	if f >= mb {
		return fmt.Sprintf("%.0fMB", f/mb)
	}
	if f >= kb {
		return fmt.Sprintf("%.0fKB", f/kb)
	}
	return fmt.Sprintf("%dB", bytes)
}

func formatServedMbps(bytesWritten int64, elapsed time.Duration) string {
	if elapsed < time.Millisecond {
		return "?mbit/s"
	}
	sec := elapsed.Seconds()
	if sec <= 0 {
		return "?mbit/s"
	}
	mbps := (float64(bytesWritten) * 8.0) / sec / 1e6
	if mbps >= 100 {
		return fmt.Sprintf("%.0fmbit/s", mbps)
	}
	if mbps >= 10 {
		return fmt.Sprintf("%.1fmbit/s", mbps)
	}
	return fmt.Sprintf("%.2fmbit/s", mbps)
}

func recordServedDownload(relPath, clientIP string, bytesWritten, fileSize int64, isRangeRequest bool, elapsed time.Duration) {
	statsMutex.Lock()
	defer statsMutex.Unlock()

	key := normalizeStatsPath(relPath)
	stats := downloadStats[key]
	stats.Bytes += bytesWritten
	stats.Count++
	downloadStats[key] = stats

	partial := isRangeRequest
	if fileSize > 0 {
		partial = partial || bytesWritten < fileSize
	}
	perClientMu.Lock()
	if _, ok := perClientFileStats[clientIP]; !ok {
		perClientFileStats[clientIP] = make(map[string]clientFileStat)
	}
	cf := perClientFileStats[clientIP][key]
	if partial {
		cf.Partial++
	} else {
		cf.Full++
	}
	perClientFileStats[clientIP][key] = cf
	perClientMu.Unlock()

	kind := "full"
	if partial {
		kind = "partial"
	}
	sizeStr := formatServedDataSize(bytesWritten)
	timeStr := formatServedDuration(elapsed)
	rateStr := formatServedMbps(bytesWritten, elapsed)
	outPrintf("Served file %s to %s (%s), sent %d bytes in %s (%s@%s)\n", key, clientIP, kind, bytesWritten, timeStr, sizeStr, rateStr)
	if serverCfg.ByteLimit > 0 {
		atomic.AddInt64(&globalBytesTransferred, bytesWritten)
	}
	appendServerEvent("download", clientIP, fmt.Sprintf("path=%s kind=%s bytes=%d", key, kind, bytesWritten))
}

// Write method for rateLimitedWriter that throttles writes based on bandwidth limit
func (w *rateLimitedWriter) Write(data []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()

	dataLen := int64(len(data))
	if dataLen == 0 {
		return w.ResponseWriter.Write(data)
	}

	now := time.Now()
	elapsed := now.Sub(w.lastWrite).Seconds()

	// Reset counter if more than 1 second has passed
	if elapsed >= 1.0 {
		w.bytesWritten = 0
		w.lastWrite = now
		elapsed = 0
	}

	// Calculate how long to wait if we've exceeded the limit
	if w.bytesWritten >= w.bytesPerSecond {
		waitTime := time.Duration((1.0 - elapsed) * float64(time.Second))
		if waitTime > 0 {
			time.Sleep(waitTime)
			w.bytesWritten = 0
			w.lastWrite = time.Now()
		}
	}

	// Write data in chunks to respect the bandwidth limit
	totalWritten := 0
	chunkSize := int64(8192) // 8KB chunks for smoother throttling
	if chunkSize > w.bytesPerSecond {
		chunkSize = w.bytesPerSecond
	}

	for int64(totalWritten) < dataLen {
		remaining := dataLen - int64(totalWritten)
		chunk := chunkSize
		if chunk > remaining {
			chunk = remaining
		}

		// Check if we can write this chunk without exceeding the limit
		if w.bytesWritten+chunk > w.bytesPerSecond {
			// Wait until we can write more
			elapsed := time.Since(w.lastWrite).Seconds()
			if elapsed < 1.0 {
				waitTime := time.Duration((1.0 - elapsed) * float64(time.Second))
				time.Sleep(waitTime)
			}
			w.bytesWritten = 0
			w.lastWrite = time.Now()
		}

		n, err := w.ResponseWriter.Write(data[totalWritten : totalWritten+int(chunk)])
		totalWritten += n
		w.bytesWritten += int64(n)
		if err != nil {
			return totalWritten, err
		}
	}

	return totalWritten, nil
}