skyview/internal/server/server.go

package server

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"path"
	"strconv"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	
	"skyview/internal/merger"
)

// OriginConfig represents the reference point configuration
type OriginConfig struct {
	Latitude  float64 `json:"latitude"`
	Longitude float64 `json:"longitude"`
	Name      string  `json:"name,omitempty"`
}

// Server handles HTTP requests and WebSocket connections
type Server struct {
	port        int
	merger      *merger.Merger
	staticFiles embed.FS
	server      *http.Server
	origin      OriginConfig
	
	// WebSocket management
	wsClients   map[*websocket.Conn]bool
	wsClientsMu sync.RWMutex
	upgrader    websocket.Upgrader
	
	// Broadcast channels
	broadcastChan chan []byte
	stopChan      chan struct{}
}

// WebSocketMessage represents messages sent over WebSocket
type WebSocketMessage struct {
	Type      string      `json:"type"`
	Timestamp int64       `json:"timestamp"`
	Data      interface{} `json:"data"`
}

// AircraftUpdate represents aircraft data for WebSocket
type AircraftUpdate struct {
	Aircraft map[string]*merger.AircraftState `json:"aircraft"`
	Sources  []*merger.Source                 `json:"sources"`
	Stats    map[string]interface{}           `json:"stats"`
}

// NewServer creates a new HTTP server
func NewServer(port int, merger *merger.Merger, staticFiles embed.FS, origin OriginConfig) *Server {
	return &Server{
		port:          port,
		merger:        merger,
		staticFiles:   staticFiles,
		origin:        origin,
		wsClients:     make(map[*websocket.Conn]bool),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true // Allow all origins in development
			},
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
		},
		broadcastChan: make(chan []byte, 100),
		stopChan:      make(chan struct{}),
	}
}

// Start starts the HTTP server
func (s *Server) Start() error {
	// Start broadcast routine
	go s.broadcastRoutine()
	
	// Start periodic updates
	go s.periodicUpdateRoutine()
	
	// Setup routes
	router := s.setupRoutes()
	
	s.server = &http.Server{
		Addr:    fmt.Sprintf(":%d", s.port),
		Handler: router,
	}
	
	return s.server.ListenAndServe()
}

// Stop gracefully stops the server
func (s *Server) Stop() {
	close(s.stopChan)
	
	if s.server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		s.server.Shutdown(ctx)
	}
}

func (s *Server) setupRoutes() http.Handler {
	router := mux.NewRouter()
	
	// API routes
	api := router.PathPrefix("/api").Subrouter()
	api.HandleFunc("/aircraft", s.handleGetAircraft).Methods("GET")
	api.HandleFunc("/aircraft/{icao}", s.handleGetAircraftDetails).Methods("GET")
	api.HandleFunc("/sources", s.handleGetSources).Methods("GET")
	api.HandleFunc("/stats", s.handleGetStats).Methods("GET")
	api.HandleFunc("/origin", s.handleGetOrigin).Methods("GET")
	api.HandleFunc("/coverage/{sourceId}", s.handleGetCoverage).Methods("GET")
	api.HandleFunc("/heatmap/{sourceId}", s.handleGetHeatmap).Methods("GET")
	
	// WebSocket
	router.HandleFunc("/ws", s.handleWebSocket)
	
	// Static files
	router.PathPrefix("/static/").Handler(s.staticFileHandler())
	router.HandleFunc("/favicon.ico", s.handleFavicon)
	
	// Main page
	router.HandleFunc("/", s.handleIndex)
	
	// Enable CORS
	return s.enableCORS(router)
}

func (s *Server) handleGetAircraft(w http.ResponseWriter, r *http.Request) {
	aircraft := s.merger.GetAircraft()
	
	// Convert ICAO keys to hex strings for JSON
	aircraftMap := make(map[string]*merger.AircraftState)
	for icao, state := range aircraft {
		aircraftMap[fmt.Sprintf("%06X", icao)] = state
	}
	
	response := map[string]interface{}{
		"timestamp": time.Now().Unix(),
		"aircraft":  aircraftMap,
		"count":     len(aircraft),
	}
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
}

func (s *Server) handleGetAircraftDetails(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	icaoStr := vars["icao"]
	
	// Parse ICAO hex string
	icao, err := strconv.ParseUint(icaoStr, 16, 32)
	if err != nil {
		http.Error(w, "Invalid ICAO address", http.StatusBadRequest)
		return
	}
	
	aircraft := s.merger.GetAircraft()
	if state, exists := aircraft[uint32(icao)]; exists {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(state)
	} else {
		http.Error(w, "Aircraft not found", http.StatusNotFound)
	}
}

func (s *Server) handleGetSources(w http.ResponseWriter, r *http.Request) {
	sources := s.merger.GetSources()
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"sources": sources,
		"count":   len(sources),
	})
}

func (s *Server) handleGetStats(w http.ResponseWriter, r *http.Request) {
	stats := s.merger.GetStatistics()
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(stats)
}

func (s *Server) handleGetOrigin(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(s.origin)
}

func (s *Server) handleGetCoverage(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	sourceID := vars["sourceId"]
	
	// Generate coverage data based on signal strength
	aircraft := s.merger.GetAircraft()
	coveragePoints := make([]map[string]interface{}, 0)
	
	for _, state := range aircraft {
		if srcData, exists := state.Sources[sourceID]; exists {
			coveragePoints = append(coveragePoints, map[string]interface{}{
				"lat":      state.Latitude,
				"lon":      state.Longitude,
				"signal":   srcData.SignalLevel,
				"distance": srcData.Distance,
				"altitude": state.Altitude,
			})
		}
	}
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"source": sourceID,
		"points": coveragePoints,
	})
}

func (s *Server) handleGetHeatmap(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	sourceID := vars["sourceId"]
	
	// Generate heatmap data grid
	aircraft := s.merger.GetAircraft()
	heatmapData := make(map[string]interface{})
	
	// Simple grid-based heatmap
	grid := make([][]float64, 100)
	for i := range grid {
		grid[i] = make([]float64, 100)
	}
	
	// Find bounds
	minLat, maxLat := 90.0, -90.0
	minLon, maxLon := 180.0, -180.0
	
	for _, state := range aircraft {
		if _, exists := state.Sources[sourceID]; exists {
			if state.Latitude < minLat {
				minLat = state.Latitude
			}
			if state.Latitude > maxLat {
				maxLat = state.Latitude
			}
			if state.Longitude < minLon {
				minLon = state.Longitude
			}
			if state.Longitude > maxLon {
				maxLon = state.Longitude
			}
		}
	}
	
	// Fill grid
	for _, state := range aircraft {
		if srcData, exists := state.Sources[sourceID]; exists {
			latIdx := int((state.Latitude - minLat) / (maxLat - minLat) * 99)
			lonIdx := int((state.Longitude - minLon) / (maxLon - minLon) * 99)
			
			if latIdx >= 0 && latIdx < 100 && lonIdx >= 0 && lonIdx < 100 {
				grid[latIdx][lonIdx] += srcData.SignalLevel
			}
		}
	}
	
	heatmapData["grid"] = grid
	heatmapData["bounds"] = map[string]float64{
		"minLat": minLat,
		"maxLat": maxLat,
		"minLon": minLon,
		"maxLon": maxLon,
	}
	
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(heatmapData)
}

func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}
	defer conn.Close()
	
	// Register client
	s.wsClientsMu.Lock()
	s.wsClients[conn] = true
	s.wsClientsMu.Unlock()
	
	// Send initial data
	s.sendInitialData(conn)
	
	// Handle client messages (ping/pong)
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
	
	// Unregister client
	s.wsClientsMu.Lock()
	delete(s.wsClients, conn)
	s.wsClientsMu.Unlock()
}

func (s *Server) sendInitialData(conn *websocket.Conn) {
	aircraft := s.merger.GetAircraft()
	sources := s.merger.GetSources()
	stats := s.merger.GetStatistics()
	
	// Convert ICAO keys to hex strings
	aircraftMap := make(map[string]*merger.AircraftState)
	for icao, state := range aircraft {
		aircraftMap[fmt.Sprintf("%06X", icao)] = state
	}
	
	update := AircraftUpdate{
		Aircraft: aircraftMap,
		Sources:  sources,
		Stats:    stats,
	}
	
	msg := WebSocketMessage{
		Type:      "initial_data",
		Timestamp: time.Now().Unix(),
		Data:      update,
	}
	
	conn.WriteJSON(msg)
}

func (s *Server) broadcastRoutine() {
	for {
		select {
		case <-s.stopChan:
			return
		case data := <-s.broadcastChan:
			s.wsClientsMu.RLock()
			for conn := range s.wsClients {
				if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
					conn.Close()
					delete(s.wsClients, conn)
				}
			}
			s.wsClientsMu.RUnlock()
		}
	}
}

func (s *Server) periodicUpdateRoutine() {
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()
	
	for {
		select {
		case <-s.stopChan:
			return
		case <-ticker.C:
			s.broadcastUpdate()
		}
	}
}

func (s *Server) broadcastUpdate() {
	aircraft := s.merger.GetAircraft()
	sources := s.merger.GetSources()
	stats := s.merger.GetStatistics()
	
	// Convert ICAO keys to hex strings
	aircraftMap := make(map[string]*merger.AircraftState)
	for icao, state := range aircraft {
		aircraftMap[fmt.Sprintf("%06X", icao)] = state
	}
	
	update := AircraftUpdate{
		Aircraft: aircraftMap,
		Sources:  sources,
		Stats:    stats,
	}
	
	msg := WebSocketMessage{
		Type:      "aircraft_update",
		Timestamp: time.Now().Unix(),
		Data:      update,
	}
	
	if data, err := json.Marshal(msg); err == nil {
		select {
		case s.broadcastChan <- data:
		default:
			// Channel full, skip this update
		}
	}
}

func (s *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
	data, err := s.staticFiles.ReadFile("static/index.html")
	if err != nil {
		http.Error(w, "Page not found", http.StatusNotFound)
		return
	}
	
	w.Header().Set("Content-Type", "text/html")
	w.Write(data)
}

func (s *Server) handleFavicon(w http.ResponseWriter, r *http.Request) {
	data, err := s.staticFiles.ReadFile("static/favicon.ico")
	if err != nil {
		http.Error(w, "Favicon not found", http.StatusNotFound)
		return
	}
	
	w.Header().Set("Content-Type", "image/x-icon")
	w.Write(data)
}

func (s *Server) staticFileHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Remove /static/ prefix from URL path to get the actual file path
		filePath := "static" + r.URL.Path[len("/static"):]
		
		data, err := s.staticFiles.ReadFile(filePath)
		if err != nil {
			http.NotFound(w, r)
			return
		}
		
		// Set content type
		ext := path.Ext(filePath)
		contentType := getContentType(ext)
		w.Header().Set("Content-Type", contentType)
		
		// Cache control
		w.Header().Set("Cache-Control", "public, max-age=3600")
		
		w.Write(data)
	})
}

func getContentType(ext string) string {
	switch ext {
	case ".html":
		return "text/html"
	case ".css":
		return "text/css"
	case ".js":
		return "application/javascript"
	case ".json":
		return "application/json"
	case ".svg":
		return "image/svg+xml"
	case ".png":
		return "image/png"
	case ".jpg", ".jpeg":
		return "image/jpeg"
	case ".ico":
		return "image/x-icon"
	default:
		return "application/octet-stream"
	}
}

func (s *Server) enableCORS(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
		
		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}
		
		handler.ServeHTTP(w, r)
	})
}