pogdark-api/main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"gopkg.in/yaml.v2"
)

var (
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	redisClient *redis.Client
	ctx         = context.Background()
)

type Message struct {
	Id        string `json:"Id"`
	Name      string `json:"Name"`
	Image     string `json:"Image"`
	Status    string `json:"Status"`
	Timestamp string `json:"Timestamp"`
}

type Config struct {
	Redis struct {
		Host string `yaml:"host"`
		Port string `yaml:"port"`
	}
}

type Client struct {
	conn *websocket.Conn
}

type Server struct {
	clients   map[*Client]bool
	broadcast chan Message
	mu        sync.Mutex
}

func getConfig(configPath string) (*Config, error) {
	// Create config structure
	config := &Config{}

	// Open config file
	file, err := os.Open(configPath)
	if err != nil {
		return nil, err
	}
	defer func(file *os.File) {
		err := file.Close()
		if err != nil {

		}
	}(file)

	// Init new YAML decode
	d := yaml.NewDecoder(file)

	// Start YAML decoding from file
	if err := d.Decode(&config); err != nil {
		return nil, err
	}

	return config, nil
}

func loadConfig() *Config {
	cfg, err := getConfig("config.yaml")
	if err != nil {
		log.Fatal(err)
	}
	return cfg
}

func initRedis(cfg Config) {
	redisClient = redis.NewClient(&redis.Options{
		Addr: cfg.Redis.Host + ":" + cfg.Redis.Port,
	})
	_, err := redisClient.Do(context.Background(), "CONFIG", "SET", "notify-keyspace-events", "KEA").Result()
	if err != nil {
		fmt.Printf("unable to set keyspace events %v", err.Error())
		os.Exit(1)
	}
	redisClient.ConfigSet(ctx, "notify-keyspace-events", "Ex")
}

func newServer() *Server {
	return &Server{
		clients:   make(map[*Client]bool),
		broadcast: make(chan Message),
	}
}

func (s *Server) addClient(client *Client) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.clients[client] = true
	broadcastAllRecordsToClient(client)
}

func (s *Server) removeClient(client *Client) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.clients[client]; ok {
		delete(s.clients, client)
		err := client.conn.Close()
		if err != nil {
			return
		}
	}
}

func (s *Server) handleConnections(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		fmt.Printf("Error upgrading connection: %v\n", err)
		return
	}

	client := &Client{conn: conn}
	s.addClient(client)

	defer s.removeClient(client)

	for {
		_, msgContent, err := conn.ReadMessage()
		if err != nil {
			fmt.Printf("Error reading message: %v\n", err)
			break
		}

		var message Message
		err = json.Unmarshal(msgContent, &message)
		if err != nil {
			log.Fatalf("Unable to marshal JSON due to %s", err)
		}
		s.broadcast <- message
	}
}

func broadcastAllRecords(s *Server) {
	allRecords, err := fetchAllRecords()
	if err != nil {
		log.Println("Error fetching all records:", err)
	}
	if len(allRecords) > 0 {
		var message Message
		for _, msgContent := range allRecords {
			err = json.Unmarshal([]byte(msgContent), &message)
			if err != nil {
				log.Println("Unable to marshal JSON due to: ", err)
			}
			fmt.Printf("Broadcasting %s,%s,%s,%s to %d clients\n", message.Id, message.Status, message.Name, message.Timestamp, len(s.clients))

			for client := range s.clients {
				if err := client.conn.WriteMessage(websocket.TextMessage, []byte(msgContent)); err != nil {
					log.Println("Failed to broadcast update:", err)
				}
			}
		}
	}
}

func broadcastRemovedRecords(s *Server, expired string) {
	var message Message = Message{Id: expired, Name: "", Image: "", Status: "expired", Timestamp: ""}
	var msgJSON, err = json.Marshal(message)
	if err != nil {
		log.Println("Error marshalling json:", err)
		return
	}
	fmt.Printf("Broadcasting expiration: %s\n", string(msgJSON))
	for client := range s.clients {
		if err := client.conn.WriteMessage(websocket.TextMessage, msgJSON); err != nil {
			log.Println("Failed to broadcast update:", err)
			return
		}
	}
}

func broadcastAllRecordsToClient(c *Client) {
	allRecords, err := fetchAllRecords()
	if err != nil {
		log.Println("Error fetching all records:", err)
	}
	fmt.Printf("Broadcasting %d records to client\n", len(allRecords))
	var message Message
	for _, msgContent := range allRecords {
		err = json.Unmarshal([]byte(msgContent), &message)
		if err != nil {
			log.Println("Unable to marshal JSON due to: ", err)
		}
		fmt.Printf("Broadcasting %s,%s,%s,%s\n", message.Id, message.Status, message.Name, message.Timestamp)
		if err := c.conn.WriteMessage(websocket.TextMessage, []byte(msgContent)); err != nil {
			log.Println("Failed to broadcast update:", err)
		}
	}
}

func (s *Server) listenForExpirationEvents() {
	ps := redisClient.PSubscribe(ctx, "__keyevent@0__:expired")
	defer func(ps *redis.PubSub) {
		err := ps.Close()
		if err != nil {

		}
	}(ps)

	for msg := range ps.Channel() {
		// Extract expired payload ID
		expiredID := msg.Payload
		fmt.Println(expiredID)
		// Broadcast expiration event
		broadcastRemovedRecords(s, expiredID)
		fmt.Printf("Done Broadcasting Expiration\n")
	}
}

func fetchAllRecords() (map[string]string, error) {
	allKeys, err := redisClient.Keys(ctx, "*").Result()
	if err != nil {
		return nil, fmt.Errorf("could not fetch keys: %v", err)
	}

	records := make(map[string]string)
	for _, key := range allKeys {
		value, err := redisClient.Get(ctx, key).Result()
		if errors.Is(err, redis.Nil) {
			continue // skip if key does not exist
		} else if err != nil {
			return nil, fmt.Errorf("could not fetch value for key %s: %v", key, err)
		}
		records[key] = value
	}

	return records, nil
}

func statusCheck(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed)
		return
	}

	var request struct {
		Id string `json:"Id"`
	}
	err := json.NewDecoder(r.Body).Decode(&request)
	if err != nil {
		http.Error(w, "Invalid JSON format", http.StatusBadRequest)
		return
	}

	if request.Id == "" {
		http.Error(w, "Missing or empty Id field", http.StatusBadRequest)
		return
	}

	value, err := redisClient.Get(ctx, request.Id).Result()
	if errors.Is(err, redis.Nil) {
		message := Message{Id: request.Id, Status: "none", Timestamp: time.Now().Format(time.RFC3339)}
		w.Header().Set("Content-Type", "application/json")
		err := json.NewEncoder(w).Encode(message)
		if err != nil {
			return
		}
		return
	} else if err != nil {
		log.Printf("Redis error: %v", err)
		http.Error(w, "Error connecting to Redis", http.StatusInternalServerError)
		return
	}

	var message Message
	err = json.Unmarshal([]byte(value), &message)
	if err != nil {
		log.Printf("Failed to decode Redis data for Id: %s - %v", request.Id, err)
		http.Error(w, "Failed to parse data", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	err = json.NewEncoder(w).Encode(message)
	if err != nil {
		http.Error(w, "Failed to encode response", http.StatusInternalServerError)
	}
}

func setState(w http.ResponseWriter, r *http.Request, s *Server) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST method is allowed", http.StatusMethodNotAllowed)
		return
	}

	var message Message
	err := json.NewDecoder(r.Body).Decode(&message)
	if err != nil || message.Id == "" || message.Status == "" {
		http.Error(w, "Invalid request format", http.StatusBadRequest)
		return
	}

	if message.Status == "none" {
		if err := redisClient.Del(ctx, message.Id).Err(); err != nil {
			log.Printf("Error deleting key from Redis: %v", err)
			http.Error(w, "Failed to delete key from Redis", http.StatusInternalServerError)
			return
		}
		broadcastRemovedRecords(s, message.Id)
		w.WriteHeader(http.StatusOK)
		return
	}

	msgJSON, err := json.Marshal(message)
	if err != nil {
		http.Error(w, "Failed to encode message", http.StatusInternalServerError)
		return
	}

	timeout := 20 * time.Second
	err = redisClient.Set(ctx, message.Id, msgJSON, timeout).Err()
	if err != nil {
		log.Printf("Failed to set key in Redis: %v", err)
		http.Error(w, "Failed to store data", http.StatusInternalServerError)
		return
	}

	broadcastAllRecords(s)
	w.WriteHeader(http.StatusOK)
}

func enableCORS(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*") // Allow all origins
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

		// Allow preflight requests for the OPTIONS method
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	})
}

func main() {
	initRedis(*loadConfig())

	server := newServer()
	router := mux.NewRouter()

	// Register routes on the mux router
	router.HandleFunc("/ws", server.handleConnections).Methods("GET")
	router.HandleFunc("/status", statusCheck).Methods("POST")
	router.HandleFunc("/set", func(w http.ResponseWriter, r *http.Request) {
		setState(w, r, server)
	}).Methods("POST")
	corsRouter := enableCORS(router)

	// Start server and other necessary goroutines
	go server.listenForExpirationEvents()

	// Pass the mux router to ListenAndServe
	fmt.Println("Server started on :8080")
	err := http.ListenAndServe(":8080", corsRouter)
	if err != nil {
		fmt.Printf("Server error: %v\n", err)
	}
}