package main

import (
	"encoding/json"
	"errors"
	"log/slog"
	"os"
	"sync"
	"time"
)

// RequestCounter tracks request counts per path in memory with periodic snapshots
type RequestCounter struct {
	counts       map[string]int64
	mu           sync.RWMutex
	snapshotPath string
	stopChan     chan struct{}
	wg           sync.WaitGroup
}

// NewRequestCounter creates a new request counter with periodic snapshots
func NewRequestCounter(snapshotInterval time.Duration) (*RequestCounter, error) {
	snapshotPath := os.Getenv("REQUEST_COUNTS_PATH")
	if snapshotPath == "" {
		err := errors.New("REQUEST_COUNTS_PATH environment variable must be set and non-empty")
		slog.Error("failed to initialize request counter", "error", err)
		os.Exit(1)
	}
	c := &RequestCounter{
		counts:       make(map[string]int64),
		snapshotPath: snapshotPath,
		stopChan:     make(chan struct{}),
	}

	// Load existing counts from snapshot file
	if err := c.load(); err != nil {
		slog.Warn("failed to load counter snapshot, starting fresh", "error", err, "path", snapshotPath)
	}

	// Start periodic snapshot goroutine
	c.wg.Add(1)
	go c.periodicSnapshot(snapshotInterval)

	return c, nil
}

// Increment increments the counter for a given path and returns the new count
func (c *RequestCounter) Increment(path string) int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.counts[path]++
	return c.counts[path]
}

// Get returns the current count for a given path
func (c *RequestCounter) Get(path string) int64 {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.counts[path]
}

func (c *RequestCounter) GetTotal() int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	var total int64
	for _, count := range c.counts {
		total += count
	}
	return total
}

// load reads the snapshot file and loads counts into memory
func (c *RequestCounter) load() error {
	data, err := os.ReadFile(c.snapshotPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil // File doesn't exist yet, that's okay
		}
		return err
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	if err := json.Unmarshal(data, &c.counts); err != nil {
		return err
	}

	slog.Info("loaded request counter snapshot", "paths", len(c.counts), "file", c.snapshotPath)
	return nil
}

// snapshot writes current counts to disk
func (c *RequestCounter) snapshot() error {
	c.mu.RLock()
	data, err := json.MarshalIndent(c.counts, "", "  ")
	c.mu.RUnlock()

	if err != nil {
		return err
	}

	// Write to temp file first, then rename for atomicity
	tempPath := c.snapshotPath + ".tmp"
	if err := os.WriteFile(tempPath, data, 0644); err != nil {
		return err
	}

	if err := os.Rename(tempPath, c.snapshotPath); err != nil {
		return err
	}

	slog.Debug("saved request counter snapshot", "paths", len(c.counts), "file", c.snapshotPath)
	return nil
}

// periodicSnapshot runs in a goroutine and saves snapshots at regular intervals
func (c *RequestCounter) periodicSnapshot(interval time.Duration) {
	defer c.wg.Done()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if err := c.snapshot(); err != nil {
				slog.Error("failed to save counter snapshot", "error", err)
			}
		case <-c.stopChan:
			// Final snapshot before shutdown
			if err := c.snapshot(); err != nil {
				slog.Error("failed to save final counter snapshot", "error", err)
			}
			return
		}
	}
}

// Close stops the periodic snapshot goroutine and saves a final snapshot
func (c *RequestCounter) Close() error {
	close(c.stopChan)
	c.wg.Wait()
	return nil
}