vetrag/controllers.go

package main

import (
	"net/http"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

// Controller struct to hold dependencies
// Add more dependencies as needed
// e.g. templates, services, etc.
type Controller struct {
	repo                 ChatRepositoryAPI
	llm                  LLMClientAPI
	visitDB              *VisitDB
	template             TemplateExecutor
	uiDBEditTemplate     TemplateExecutor
	uiAdminChatsTemplate TemplateExecutor
	uiAdminLoginTemplate TemplateExecutor
}

// NewController creates a new Controller instance
func NewController(repo ChatRepositoryAPI, llm LLMClientAPI, visitDB *VisitDB, template, uiDBEditTemplate, uiAdminChatsTemplate, uiAdminLoginTemplate TemplateExecutor) *Controller {
	return &Controller{
		repo:                 repo,
		llm:                  llm,
		visitDB:              visitDB,
		template:             template,
		uiDBEditTemplate:     uiDBEditTemplate,
		uiAdminChatsTemplate: uiAdminChatsTemplate,
		uiAdminLoginTemplate: uiAdminLoginTemplate,
	}
}

// Handler for root UI
func (ctrl *Controller) RootUI(c *gin.Context) {
	c.Status(http.StatusOK)
	if err := ctrl.template.Execute(c.Writer, nil); err != nil {
		logrus.Errorf("Failed to execute ui.html template: %v", err)
	}
}

// Handler for health check
func (ctrl *Controller) Health(c *gin.Context) {
	c.Status(http.StatusOK)
	c.JSON(http.StatusOK, gin.H{"status": "ok"})
}

// Handler for chat
func (ctrl *Controller) HandleChat(c *gin.Context) {
	// Use your existing chatService logic here
	// You may want to move NewChatService to a service file for full MVC
	chatService := NewChatService(ctrl.llm, ctrl.visitDB, ctrl.repo)
	chatService.HandleChat(c)
}

// Handler for admin UI
func (ctrl *Controller) AdminUI(c *gin.Context) {
	c.Status(http.StatusOK)
	if err := ctrl.uiDBEditTemplate.Execute(c.Writer, nil); err != nil {
		logrus.Errorf("Failed to execute ui_dbedit.html template: %v", err)
	}
}

// Handler for download db.yaml
func (ctrl *Controller) DownloadDB(c *gin.Context) {
	c.File("db.yaml")
}

// Handler for saving knowledgeModel (snapshot)
func (ctrl *Controller) SaveKnowledgeModel(c *gin.Context) {
	if ctrl.repo == nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "repository not configured"})
		return
	}
	var req struct {
		Text string `json:"text"`
	}
	if err := c.BindJSON(&req); err != nil || req.Text == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
		return
	}
	if err := ctrl.repo.SaveKnowledgeModel(c.Request.Context(), req.Text); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to save snapshot"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "ok"})
}

// Handler for listing chat interactions
func (ctrl *Controller) ListChatInteractions(c *gin.Context) {
	if ctrl.repo == nil {
		c.JSON(http.StatusOK, gin.H{"items": []ChatInteraction{}, "pagination": gin.H{"limit": 0, "offset": 0, "count": 0}, "warning": "repository not configured"})
		return
	}
	limit := 50
	if ls := c.Query("limit"); ls != "" {
		if v, err := strconv.Atoi(ls); err == nil {
			limit = v
		}
	}
	offset := 0
	if osf := c.Query("offset"); osf != "" {
		if v, err := strconv.Atoi(osf); err == nil {
			offset = v
		}
	}
	items, err := ctrl.repo.ListChatInteractions(c.Request.Context(), limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list interactions"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": items, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(items)}})
}

// Handler for listing LLM events
func (ctrl *Controller) ListLLMEvents(c *gin.Context) {
	corr := c.Query("correlation_id")
	if corr == "" {
		if strings.Contains(c.GetHeader("Accept"), "application/json") {
			c.JSON(http.StatusBadRequest, gin.H{"error": "missing correlation_id"})
			return
		}
		c.Status(http.StatusOK)
		if err := ctrl.uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
			logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
		}
		return
	}
	if ctrl.repo == nil {
		c.JSON(http.StatusOK, gin.H{"items": []RawLLMEvent{}, "pagination": gin.H{"limit": 0, "offset": 0, "count": 0}, "warning": "repository not configured"})
		return
	}
	limit := 100
	if ls := c.Query("limit"); ls != "" {
		if v, err := strconv.Atoi(ls); err == nil {
			limit = v
		}
	}
	offset := 0
	if osf := c.Query("offset"); osf != "" {
		if v, err := strconv.Atoi(osf); err == nil {
			offset = v
		}
	}
	events, err := ctrl.repo.ListLLMRawEvents(c.Request.Context(), corr, limit, offset)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list events"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": events, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(events)}})
}

// Handler for admin chats UI
func (ctrl *Controller) AdminChatsUI(c *gin.Context) {
	c.Status(http.StatusOK)
	if err := ctrl.uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
		logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
	}
}

// Handler for listing knowledgeModel entries
func (ctrl *Controller) ListKnowledgeModels(c *gin.Context) {
	if ctrl.repo == nil {
		c.JSON(http.StatusOK, gin.H{"items": []interface{}{}})
		return
	}
	models, err := ctrl.repo.ListKnowledgeModels(c.Request.Context(), 100, 0)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list snapshots"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"items": models})
}

// Handler for getting a specific knowledgeModel entry
func (ctrl *Controller) GetKnowledgeModel(c *gin.Context) {
	if ctrl.repo == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "repository not configured"})
		return
	}
	idStr := c.Param("id")
	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
		return
	}
	text, err := ctrl.repo.GetKnowledgeModelText(c.Request.Context(), id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "snapshot not found"})
		return
	}
	c.Data(http.StatusOK, "text/yaml; charset=utf-8", []byte(text))
}

// AdminLoginPage serves the admin login page
func (ctrl *Controller) AdminLoginPage(c *gin.Context) {
	c.Status(http.StatusOK)
	if err := ctrl.uiAdminLoginTemplate.Execute(c.Writer, gin.H{"Error": c.Query("error")}); err != nil {
		logrus.Errorf("Failed to execute ui_admin_login.html template: %v", err)
	}
}

// AdminLogin handles admin login POST
func (ctrl *Controller) AdminLogin(c *gin.Context) {
	username := c.PostForm("username")
	password := c.PostForm("password")
	user, err := ctrl.repo.GetUserByUsername(c.Request.Context(), username)
	if err != nil || !CheckPasswordHash(password, user.PasswordHash) {
		c.Redirect(http.StatusFound, "/admin/login?error=Invalid+credentials")
		return
	}
	token, err := GenerateJWT(user.Username)
	if err != nil {
		c.Redirect(http.StatusFound, "/admin/login?error=Internal+error")
		return
	}
	c.SetCookie("admin_jwt", token, 86400, "/", "", false, true)
	c.Redirect(http.StatusFound, "/admin")
}

// AdminLogout logs out the admin
func (ctrl *Controller) AdminLogout(c *gin.Context) {
	c.SetCookie("admin_jwt", "", -1, "/", "", false, true)
	c.Redirect(http.StatusFound, "/admin/login")
}

// AdminAuthMiddleware checks for a valid admin JWT cookie
func AdminAuthMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Request.URL.Path == "/admin/login" || c.Request.URL.Path == "/admin/logout" {
			c.Next()
			return
		}
		jwtToken, err := c.Cookie("admin_jwt")
		if err != nil {
			c.Redirect(http.StatusFound, "/admin/login")
			c.Abort()
			return
		}
		username, err := ValidateJWT(jwtToken)
		if err != nil || username == "" {
			c.Redirect(http.StatusFound, "/admin/login")
			c.Abort()
			return
		}
		c.Set("admin_username", username)
		c.Next()
	}
}

// CreateInitialAdmin allows creation of the first admin user if none exist
func (ctrl *Controller) CreateInitialAdmin(c *gin.Context) {
	// Only allow if no users exist
	var count int
	err := ctrl.repo.CountUsers(&count)
	if err != nil || count > 0 {
		c.JSON(http.StatusForbidden, gin.H{"error": "Admin user already exists or error"})
		return
	}
	username := c.PostForm("username")
	password := c.PostForm("password")
	if username == "" || password == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "Username and password required"})
		return
	}
	hash, err := HashPassword(password)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to hash password"})
		return
	}
	err = ctrl.repo.CreateUser(username, hash)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create user"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "Admin user created"})
}

// TemplateExecutor is an interface for template execution
// This allows for easier testing and decoupling
// You can implement this interface for your templates

type TemplateExecutor interface {
	Execute(wr http.ResponseWriter, data interface{}) error
}