package main

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

	"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()

		// Bidirectional copy
		done := make(chan struct{})
		go func() {
			for {
				_, msg, err := conn.ReadMessage()
				if err != nil {
					break
				}
				uiConn.WriteMessage(websocket.TextMessage, msg)
			}
			close(done)
		}()
		for {
			_, msg, err := uiConn.ReadMessage()
			if err != nil {
				break
			}
			conn.WriteMessage(websocket.TextMessage, msg)
		}
		<-done
	})

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

	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 all other routes
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Serve index.html for SPA routes
		if r.URL.Path == "/" || strings.HasPrefix(r.URL.Path, "/sessions") || strings.HasPrefix(r.URL.Path, "/runs") || strings.HasPrefix(r.URL.Path, "/infrastructure") || strings.HasPrefix(r.URL.Path, "/agents") {
			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)
			return
		}

		http.NotFound(w, r)
	})

	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
}