ai-servers/llm-gateway/internal/dashboard/handler.go

package dashboard

import (
	"embed"
	"fmt"
	"html/template"
	"net/http"
	"time"

	"llm-gateway/internal/auth"
)

//go:embed templates/*.html templates/partials/*.html
var templateFiles embed.FS

var templateFuncs = template.FuncMap{
	"formatTime": func(ts int64) string {
		if ts == 0 {
			return "never"
		}
		return time.Unix(ts, 0).Format("2006-01-02")
	},
	"addInt": func(a, b int) int {
		return a + b
	},
	"formatCost": func(v float64) string {
		if v == 0 {
			return "$0.00"
		}
		if v < 0.01 {
			return fmt.Sprintf("$%.6f", v)
		}
		return fmt.Sprintf("$%.4f", v)
	},
}

// PageData is the common data passed to all templates.
type PageData struct {
	ActivePage string
	User       *auth.User
	// Page-specific data
	Summary    *SummaryResult
	Models     []ModelStats
	Providers  []ProviderStats
	TokenStats []TokenUsageStats
	Tokens     []auth.APIToken
	Users      []auth.User
}

// Dashboard serves the HTMX-based dashboard pages.
type Dashboard struct {
	templates *template.Template
	authStore *auth.Store
	statsAPI  *StatsAPI
}

// NewDashboard creates a new Dashboard handler.
func NewDashboard(authStore *auth.Store, statsAPI *StatsAPI) *Dashboard {
	tmpl := template.Must(
		template.New("").Funcs(templateFuncs).ParseFS(templateFiles,
			"templates/*.html",
			"templates/partials/*.html",
		),
	)

	return &Dashboard{
		templates: tmpl,
		authStore: authStore,
		statsAPI:  statsAPI,
	}
}

// LoginPage serves the login page.
func (d *Dashboard) LoginPage(w http.ResponseWriter, r *http.Request) {
	if !d.authStore.HasAnyUser() {
		http.Redirect(w, r, "/setup", http.StatusFound)
		return
	}
	if user := d.getSessionUser(r); user != nil {
		http.Redirect(w, r, "/dashboard", http.StatusFound)
		return
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	d.templates.ExecuteTemplate(w, "login", nil)
}

// SetupPage serves the initial setup page.
func (d *Dashboard) SetupPage(w http.ResponseWriter, r *http.Request) {
	if d.authStore.HasAnyUser() {
		http.Redirect(w, r, "/login", http.StatusFound)
		return
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	d.templates.ExecuteTemplate(w, "setup", nil)
}

// DashboardPage serves the main dashboard view.
func (d *Dashboard) DashboardPage(w http.ResponseWriter, r *http.Request) {
	user := auth.UserFromContext(r.Context())
	tokenNames := d.statsAPI.TokenNamesForUser(user)

	data := PageData{
		ActivePage: "dashboard",
		User:       user,
		Summary:    d.statsAPI.GetSummary(tokenNames),
		Models:     d.statsAPI.GetModels(tokenNames),
		Providers:  d.statsAPI.GetProviders(tokenNames),
		TokenStats: d.statsAPI.GetTokenUsage(tokenNames),
	}

	d.renderDashboardPage(w, r, "partials/dashboard.html", data)
}

// TokensPage serves the tokens management view.
func (d *Dashboard) TokensPage(w http.ResponseWriter, r *http.Request) {
	user := auth.UserFromContext(r.Context())

	var userID int64
	if !user.IsAdmin {
		userID = user.ID
	}

	tokens, _ := d.authStore.ListAPITokens(userID)
	if tokens == nil {
		tokens = []auth.APIToken{}
	}

	d.renderDashboardPage(w, r, "partials/tokens.html", PageData{
		ActivePage: "tokens",
		User:       user,
		Tokens:     tokens,
	})
}

// UsersPage serves the user management view (admin only).
func (d *Dashboard) UsersPage(w http.ResponseWriter, r *http.Request) {
	user := auth.UserFromContext(r.Context())
	users, _ := d.authStore.ListUsers()

	d.renderDashboardPage(w, r, "partials/users.html", PageData{
		ActivePage: "users",
		User:       user,
		Users:      users,
	})
}

// SettingsPage serves the settings view.
func (d *Dashboard) SettingsPage(w http.ResponseWriter, r *http.Request) {
	user := auth.UserFromContext(r.Context())
	user, _ = d.authStore.GetUserByID(user.ID)

	d.renderDashboardPage(w, r, "partials/settings.html", PageData{
		ActivePage: "settings",
		User:       user,
	})
}

// renderDashboardPage renders either the full layout or just the content partial.
func (d *Dashboard) renderDashboardPage(w http.ResponseWriter, r *http.Request, partialFile string, data PageData) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")

	if r.Header.Get("HX-Request") == "true" {
		tmpl := template.Must(
			template.New("").Funcs(templateFuncs).ParseFS(templateFiles, "templates/"+partialFile),
		)
		tmpl.ExecuteTemplate(w, "content", data)
	} else {
		tmpl := template.Must(
			template.New("").Funcs(templateFuncs).ParseFS(templateFiles,
				"templates/layout.html",
				"templates/"+partialFile,
			),
		)
		tmpl.ExecuteTemplate(w, "layout", data)
	}
}

func (d *Dashboard) getSessionUser(r *http.Request) *auth.User {
	cookie, err := r.Cookie("llmgw_session")
	if err != nil || cookie.Value == "" {
		return nil
	}
	sess, err := d.authStore.GetSession(cookie.Value)
	if err != nil {
		return nil
	}
	user, err := d.authStore.GetUserByID(sess.UserID)
	if err != nil {
		return nil
	}
	return user
}