ZovOS_AI/internal/ui/server.go

package ui

import (
	"embed"
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/agent"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/config"
)

//go:embed static/*
var staticFiles embed.FS

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

// Server is the HTTP/WebSocket server for the UI.
type Server struct {
	cfg       config.ServerConfig
	aiManager *ai.Manager
	agent     *agent.Agent
	memory    *agent.Memory
	log       *slog.Logger
}

// NewServer creates a new UI server.
func NewServer(cfg config.ServerConfig, aiMgr *ai.Manager, ag *agent.Agent, mem *agent.Memory, log *slog.Logger) *Server {
	return &Server{
		cfg:       cfg,
		aiManager: aiMgr,
		agent:     ag,
		memory:    mem,
		log:       log,
	}
}

// Start begins serving HTTP requests.
func (s *Server) Start() error {
	mux := http.NewServeMux()

	// Static files
	mux.Handle("/", http.FileServer(http.FS(staticFiles)))

	// API endpoints
	mux.HandleFunc("/api/providers", s.handleProviders)
	mux.HandleFunc("/api/provider", s.handleSwitchProvider)
	mux.HandleFunc("/api/sessions", s.handleSessions)
	mux.HandleFunc("/api/session/clear", s.handleClearSession)

	// WebSocket for chat
	mux.HandleFunc("/ws", s.handleWebSocket)

	addr := fmt.Sprintf("%s:%d", s.cfg.Host, s.cfg.Port)
	s.log.Info("starting UI server", "address", addr)
	return http.ListenAndServe(addr, mux)
}

// --- REST Handlers ---

func (s *Server) handleProviders(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"providers": s.aiManager.List(),
		"active":    s.aiManager.ActiveName(),
	})
}

func (s *Server) handleSwitchProvider(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPut && r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var req struct {
		Provider string `json:"provider"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if err := s.aiManager.SetActive(req.Provider); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	// Update agent's provider
	p, _ := s.aiManager.Active()
	s.agent.SetProvider(p)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "ok", "active": req.Provider})
}

func (s *Server) handleSessions(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"sessions": s.memory.ListSessions(),
	})
}

func (s *Server) handleClearSession(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var req struct {
		SessionID string `json:"session_id"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	s.memory.Clear(req.SessionID)
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

// --- WebSocket Handler ---

type wsMessage struct {
	Type      string `json:"type"` // "message", "tool_call", "tool_result", "stream", "done", "error"
	Content   string `json:"content,omitempty"`
	SessionID string `json:"session_id,omitempty"`
	ToolName  string `json:"tool_name,omitempty"`
	ToolArgs  string `json:"tool_args,omitempty"`
	ToolID    string `json:"tool_id,omitempty"`
}

func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		s.log.Error("websocket upgrade failed", "error", err)
		return
	}
	defer conn.Close()

	var mu sync.Mutex
	writeJSON := func(msg wsMessage) {
		mu.Lock()
		defer mu.Unlock()
		conn.WriteJSON(msg)
	}

	for {
		var incoming wsMessage
		if err := conn.ReadJSON(&incoming); err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				s.log.Error("websocket read error", "error", err)
			}
			return
		}

		if incoming.Type != "message" {
			continue
		}

		sessionID := incoming.SessionID
		if sessionID == "" {
			sessionID = "default"
		}

		history := s.memory.Get(sessionID)

		go func() {
			callback := func(chunk string, toolCall *ai.ToolCall, toolResult *string, done bool) {
				if done {
					writeJSON(wsMessage{Type: "done", SessionID: sessionID})
					return
				}
				if chunk != "" {
					writeJSON(wsMessage{Type: "stream", Content: chunk, SessionID: sessionID})
				}
				if toolCall != nil && toolResult == nil {
					writeJSON(wsMessage{
						Type:      "tool_call",
						ToolName:  toolCall.Function.Name,
						ToolArgs:  toolCall.Function.Arguments,
						ToolID:    toolCall.ID,
						SessionID: sessionID,
					})
				}
				if toolCall != nil && toolResult != nil {
					writeJSON(wsMessage{
						Type:      "tool_result",
						ToolName:  toolCall.Function.Name,
						Content:   *toolResult,
						ToolID:    toolCall.ID,
						SessionID: sessionID,
					})
				}
			}

			messages, err := s.agent.Run(r.Context(), incoming.Content, history, callback)
			if err != nil {
				writeJSON(wsMessage{Type: "error", Content: err.Error(), SessionID: sessionID})
				return
			}

			// Save updated history (strip system prompt)
			var filtered []ai.Message
			for _, m := range messages {
				if m.Role != "system" {
					filtered = append(filtered, m)
				}
			}
			s.memory.Set(sessionID, filtered)
		}()
	}
}