main.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/hungpdn/llmgo/pkg/agent"
	"github.com/hungpdn/llmgo/pkg/llm"
	"github.com/hungpdn/llmgo/pkg/llm/ollama"
	"github.com/hungpdn/llmgo/pkg/log"
	"github.com/hungpdn/llmgo/pkg/memory"
	"github.com/hungpdn/llmgo/pkg/orchestrator"
	"github.com/hungpdn/llmgo/pkg/tools"
	"github.com/hungpdn/llmgo/pkg/tools/std"
)

// AppContext holds shared dependencies (Dependency Injection).
// This replaces global variables and makes the server thread-safe and testable.
type AppContext struct {
	Provider llm.Provider
	Registry *tools.Registry
	Router   orchestrator.Router
}

// Standardized API Input/Output structures
type ChatRequest struct {
	Message string `json:"message"`
}

type ChatResponse struct {
	Reply string `json:"reply,omitempty"`
	Error string `json:"error,omitempty"`
}

func main() {
	// 1. Setup Dependencies
	// Ensure you have run: `ollama pull qwen2.5:7b-instruct`
	provider := ollama.NewClient("http://localhost:11434", "qwen2.5:7b-instruct")

	registry := tools.NewRegistry()
	_ = registry.Register(std.NewTimeTool())
	_ = registry.Register(std.NewCalculatorTool())

	router := orchestrator.NewLLMRouter(provider)

	app := &AppContext{
		Provider: provider,
		Registry: registry,
		Router:   router,
	}

	// 2. Setup HTTP Multiplexer
	mux := http.NewServeMux()
	mux.HandleFunc("/chat", app.handleChat)

	server := &http.Server{
		Addr:    ":8080",
		Handler: mux,
		// Explicitly set ReadHeaderTimeout to mitigate Slowloris attacks (Fixes gosec G112)
		ReadHeaderTimeout: 5 * time.Second,
	}

	// 3. Run Server in Background
	go func() {
		fmt.Println("🚀 Server is running on http://localhost:8080/chat")
		if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Errorf("HTTP server error: %v", err)
		}
	}()

	// 4. Trap OS Signals for Graceful Shutdown (Ctrl+C, Docker stop)
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
	<-quit
	fmt.Println("\n🛑 Shutting down HTTP server...")

	// Allow ongoing requests 10 seconds to finish before forcing shutdown
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := server.Shutdown(ctx); err != nil {
		log.Errorf("Server forced to shutdown: %v", err)
	}
	fmt.Println("✅ Server exited safely.")
}

// handleChat processes each HTTP request independently (Stateless).
func (a *AppContext) handleChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		respondJSON(w, http.StatusMethodNotAllowed, ChatResponse{Error: "Method POST required"})
		return
	}

	var req ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		respondJSON(w, http.StatusBadRequest, ChatResponse{Error: "Invalid JSON payload"})
		return
	}

	if req.Message == "" {
		respondJSON(w, http.StatusBadRequest, ChatResponse{Error: "Message cannot be empty"})
		return
	}

	// Note: For a stateful Chat App (remembering previous turns), replace BufferMemory
	// with RedisMemory and use a "Session-ID" extracted from the HTTP Headers.
	orchMem := memory.NewTokenWindowMemory(8192, nil)

	coder := agent.NewReActAgent(
		"Coder",
		"Software Engineer",
		`You are a Senior Go Developer.
		- Write the Golang code requested by the user.
		- Output your final code clearly in markdown blocks.`,
		a.Provider,
		memory.NewTokenWindowMemory(2048, nil),
		a.Registry,
		agent.WithMaxTurns(5),
	)

	reviewer := agent.NewReActAgent(
		"Reviewer",
		"Quality Assurance",
		fmt.Sprintf(`You are a strict code reviewer.
		- If the code has errors, point them out.
		- If the code is correct, you MUST reply with exactly: "%s"`, orchestrator.DefaultStopPhrase),
		a.Provider,
		memory.NewTokenWindowMemory(2048, nil),
		nil,
	)

	engine := orchestrator.New(a.Router, orchMem)
	engine.Register(coder)
	engine.Register(reviewer)

	// If the user closes the browser tab or cancels the Postman request, r.Context() will be canceled.
	// This immediately aborts the Coder/Reviewer goroutines and network calls, guaranteeing Zero Memory Leaks!
	go coder.Run(r.Context())
	go reviewer.Run(r.Context())

	// Execute the multi-agent task
	err := engine.Broadcast(r.Context(), req.Message)
	if err != nil {
		// Smart handling: Client disconnected early
		if errors.Is(err, context.Canceled) {
			log.Warnf("Client disconnected early. Task aborted cleanly.")
			return
		}
		log.Errorf("Orchestrator error: %v", err)
		respondJSON(w, http.StatusInternalServerError, ChatResponse{Error: err.Error()})
		return
	}

	// Retrieve the final result from the Orchestrator's memory
	history, _ := orchMem.History(r.Context())
	lastMsg := history[len(history)-1]

	respondJSON(w, http.StatusOK, ChatResponse{Reply: lastMsg.Content})
}

// respondJSON is a helper function to send standardized JSON responses.
func respondJSON(w http.ResponseWriter, status int, payload interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(payload)
}