package main

import (
	"embed"
	"io"
	"io/fs"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
)

//go:embed static
var staticFiles embed.FS

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

func main() {
	addr := envDefault("AGENTMON_UI_ADDR", ":8082")
	queryAPIBase := envDefault("AGENTMON_QUERY_BASE", "http://query-api")

	mux := http.NewServeMux()

	// Health check
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	// API proxy to query-api
	queryURL, _ := url.Parse(queryAPIBase)
	proxy := httputil.NewSingleHostReverseProxy(queryURL)
	mux.HandleFunc("/api/", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = strings.TrimPrefix(r.URL.Path, "/api")
		r.Host = queryURL.Host
		proxy.ServeHTTP(w, r)
	})

	// WebSocket proxy to query-api
	mux.HandleFunc("/api/v1/ws", func(w http.ResponseWriter, r *http.Request) {
		queryWSURL := "ws://" + queryURL.Host + "/v1/ws"
		conn, _, err := websocket.DefaultDialer.Dial(queryWSURL, nil)
		if err != nil {
			http.Error(w, "failed to connect to upstream WebSocket", http.StatusBadGateway)
			return
		}
		defer conn.Close()

		uiConn, err := wsUpgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer uiConn.Close()

		var uiMu, upstreamMu sync.Mutex

		// upstream → UI
		done := make(chan struct{})
		go func() {
			defer close(done)
			for {
				_, msg, err := conn.ReadMessage()
				if err != nil {
					break
				}
				uiMu.Lock()
				err = uiConn.WriteMessage(websocket.TextMessage, msg)
				uiMu.Unlock()
				if err != nil {
					break
				}
			}
		}()

		// UI → upstream
		for {
			_, msg, err := uiConn.ReadMessage()
			if err != nil {
				break
			}
			upstreamMu.Lock()
			err = conn.WriteMessage(websocket.TextMessage, msg)
			upstreamMu.Unlock()
			if err != nil {
				break
			}
		}
		<-done
	})

	// Static files
	staticFS, _ := fs.Sub(staticFiles, "static")
	fileServer := http.FileServer(http.FS(staticFS))

	mux.HandleFunc("/favicon.svg", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = "/favicon.svg"
		fileServer.ServeHTTP(w, r)
	})

	mux.HandleFunc("/favicon.ico", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/favicon.svg", http.StatusMovedPermanently)
	})

	mux.HandleFunc("/static/", func(w http.ResponseWriter, r *http.Request) {
		r.URL.Path = strings.TrimPrefix(r.URL.Path, "/static")
		fileServer.ServeHTTP(w, r)
	})

	// SPA catch-all: serve index.html for routes without file extensions
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// If the path contains a dot, it's likely a missing static asset
		if strings.Contains(r.URL.Path, ".") {
			http.NotFound(w, r)
			return
		}

		f, err := staticFiles.Open("static/index.html")
		if err != nil {
			http.Error(w, "index.html not found", http.StatusInternalServerError)
			return
		}
		defer f.Close()

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		_, _ = io.Copy(w, f)
	})

	log.Printf("web-ui listening on %s", addr)
	log.Fatal(http.ListenAndServe(addr, mux))
}

func envDefault(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}