package main

import (
	"embed"
	"encoding/json"
	"flag"
	"io/fs"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

//go:embed static
var staticFiles embed.FS

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// ---------------------------------------------------------------------------
// Wire message types
// ---------------------------------------------------------------------------

type InMsg struct {
	Type string          `json:"type"`
	Data json.RawMessage `json:"data,omitempty"`
}

type ConnectData struct {
	Host    string `json:"host"`
	Port    int    `json:"port"`
	UDPPort int    `json:"udp_port"`
	LogPort int    `json:"log_port"`
}

type CmdData struct {
	Cmd string `json:"cmd"`
}

// ---------------------------------------------------------------------------
// Hub – fan-out WS messages to all connected browsers
// ---------------------------------------------------------------------------

type Hub struct {
	mu      sync.RWMutex
	clients map[*WSClient]struct{}
	marte   *MarteClient
}

func newHub() *Hub {
	return &Hub{clients: make(map[*WSClient]struct{})}
}

func (h *Hub) add(c *WSClient) {
	h.mu.Lock()
	h.clients[c] = struct{}{}
	h.mu.Unlock()
}

func (h *Hub) remove(c *WSClient) {
	h.mu.Lock()
	delete(h.clients, c)
	h.mu.Unlock()
}

func (h *Hub) broadcast(msg []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		select {
		case c.send <- msg:
		default:
		}
	}
}

// ---------------------------------------------------------------------------
// WSClient – one browser tab
// ---------------------------------------------------------------------------

type WSClient struct {
	hub  *Hub
	conn *websocket.Conn
	send chan []byte
}

func (c *WSClient) writePump() {
	ticker := time.NewTicker(30 * time.Second)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case msg, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (c *WSClient) readPump() {
	defer func() {
		c.hub.remove(c)
		c.conn.Close()
	}()
	c.conn.SetReadLimit(256 * 1024)
	c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})
	for {
		_, raw, err := c.conn.ReadMessage()
		if err != nil {
			break
		}
		c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))

		var in InMsg
		if err := json.Unmarshal(raw, &in); err != nil {
			continue
		}
		c.handleIncoming(in)
	}
}

func (c *WSClient) handleIncoming(in InMsg) {
	m := c.hub.marte
	switch in.Type {
	case "connect":
		var d ConnectData
		if err := json.Unmarshal(in.Data, &d); err != nil {
			return
		}
		if d.Host == "" {
			d.Host = "127.0.0.1"
		}
		if d.Port == 0 {
			d.Port = 8080
		}
		if d.UDPPort == 0 {
			d.UDPPort = 8081
		}
		if d.LogPort == 0 {
			d.LogPort = 8082
		}
		m.Connect(d.Host, d.Port, d.UDPPort, d.LogPort)

	case "disconnect":
		m.Disconnect()

	case "cmd":
		var d CmdData
		if err := json.Unmarshal(in.Data, &d); err != nil {
			return
		}
		m.SendCommand(d.Cmd)
	}
}

// ---------------------------------------------------------------------------
// HTTP handler
// ---------------------------------------------------------------------------

func wsHandler(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("upgrade:", err)
		return
	}
	c := &WSClient{hub: hub, conn: conn, send: make(chan []byte, 256)}
	hub.add(c)

	// Send current connection status to newly joined client.
	if hub.marte.IsConnected() {
		b, _ := json.Marshal(map[string]any{"type": "connected"})
		c.send <- b
		// Also push current signal list if already discovered.
		hub.marte.SendCachedState(c)
	}

	go c.writePump()
	c.readPump() // blocks until client disconnects
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------

func main() {
	addr := flag.String("addr", ":7777", "HTTP listen address")
	flag.Parse()

	hub := newHub()
	hub.marte = newMarteClient(hub)

	sub, err := fs.Sub(staticFiles, "static")
	if err != nil {
		log.Fatal(err)
	}
	fileServer := http.FileServer(http.FS(sub))
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		wsHandler(hub, w, r)
	})
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Serve index.html for the root path explicitly.
		if r.URL.Path == "/" {
			data, err := staticFiles.ReadFile("static/index.html")
			if err != nil {
				http.Error(w, "not found", http.StatusNotFound)
				return
			}
			w.Header().Set("Content-Type", "text/html; charset=utf-8")
			w.Write(data)
			return
		}
		fileServer.ServeHTTP(w, r)
	})

	log.Printf("MARTe2 Web Debug Client listening on %s", *addr)
	log.Fatal(http.ListenAndServe(*addr, nil))
}