KAOS/code/internal/server/console.go

package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

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

// sessionState represents the state of a console session.
type sessionState struct {
	mu        sync.Mutex
	status    string // "idle" or "running"
	cmd       *exec.Cmd
	execID    string
	taskID    string // which task is being worked on (if any)
	history   []historyEntry
	output    []string
	listeners map[chan string]bool
}

// historyEntry represents a command in the session history.
type historyEntry struct {
	Command   string `json:"command"`
	ExecID    string `json:"exec_id"`
	Timestamp string `json:"timestamp"`
	Status    string `json:"status"` // "running", "done", "error", "killed"
}

// execRequest is the JSON body for starting a command.
type execRequest struct {
	Cmd     string `json:"cmd"`
	Session int    `json:"session"`
}

// execResponse is the JSON response after starting a command.
type execResponse struct {
	ExecID  string `json:"exec_id"`
	Session int    `json:"session"`
}

// sessionStatus represents the status of a session for the API.
type sessionStatus struct {
	Session int    `json:"session"`
	Status  string `json:"status"`
	TaskID  string `json:"task_id,omitempty"`
	ExecID  string `json:"exec_id,omitempty"`
}

// consoleManager manages the two console sessions.
type consoleManager struct {
	sessions [2]*sessionState
	mu       sync.Mutex
	counter  int
}

// newConsoleManager creates a new console manager with two idle sessions.
func newConsoleManager() *consoleManager {
	return &consoleManager{
		sessions: [2]*sessionState{
			{status: "idle", listeners: make(map[chan string]bool)},
			{status: "idle", listeners: make(map[chan string]bool)},
		},
	}
}

// nextExecID generates a unique execution ID.
func (cm *consoleManager) nextExecID() string {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.counter++
	return fmt.Sprintf("exec-%d-%d", time.Now().Unix(), cm.counter)
}

// getSession returns a session by index (0 or 1).
func (cm *consoleManager) getSession(idx int) *sessionState {
	if idx < 0 || idx > 1 {
		return nil
	}
	return cm.sessions[idx]
}

// handleConsoleExec starts a command in a session.
func (s *Server) handleConsoleExec(c *gin.Context) {
	var req execRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "nevalidan JSON: " + err.Error()})
		return
	}

	if req.Session < 1 || req.Session > 2 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "sesija mora biti 1 ili 2"})
		return
	}

	sessionIdx := req.Session - 1
	session := s.console.getSession(sessionIdx)

	session.mu.Lock()
	if session.status == "running" {
		session.mu.Unlock()
		c.JSON(http.StatusConflict, gin.H{"error": "sesija je zauzeta"})
		return
	}

	execID := s.console.nextExecID()
	session.status = "running"
	session.execID = execID
	session.output = nil
	session.mu.Unlock()

	// Add to history
	entry := historyEntry{
		Command:   req.Cmd,
		ExecID:    execID,
		Timestamp: timeNow(),
		Status:    "running",
	}

	session.mu.Lock()
	session.history = append(session.history, entry)
	if len(session.history) > 50 {
		session.history = session.history[len(session.history)-50:]
	}
	session.mu.Unlock()

	// Start the command in background
	go s.runCommand(session, req.Cmd, execID)

	c.JSON(http.StatusOK, execResponse{
		ExecID:  execID,
		Session: req.Session,
	})
}

// cleanEnv returns the current environment with CLAUDECODE removed,
// so child claude processes don't inherit the parent's session.
func cleanEnv() []string {
	var env []string
	for _, e := range os.Environ() {
		if !strings.HasPrefix(e, "CLAUDECODE=") {
			env = append(env, e)
		}
	}
	return env
}

// runCommand executes a command and streams output to listeners.
func (s *Server) runCommand(session *sessionState, command, execID string) {
	// Build the claude command
	cmd := exec.Command("claude", "--permission-mode", "dontAsk", "-p", command)
	cmd.Dir = s.projectRoot()
	cmd.Env = cleanEnv()

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		s.sendToSession(session, "[greška: "+err.Error()+"]")
		s.finishSession(session, execID, "error")
		return
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		s.sendToSession(session, "[greška: "+err.Error()+"]")
		s.finishSession(session, execID, "error")
		return
	}

	session.mu.Lock()
	session.cmd = cmd
	session.mu.Unlock()

	if err := cmd.Start(); err != nil {
		s.sendToSession(session, "[greška pri pokretanju: "+err.Error()+"]")
		s.finishSession(session, execID, "error")
		return
	}

	// Read stdout and stderr concurrently
	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		s.streamReader(session, stdout)
	}()

	go func() {
		defer wg.Done()
		s.streamReader(session, stderr)
	}()

	wg.Wait()

	err = cmd.Wait()
	status := "done"
	if err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			status = "error"
		}
	}

	s.finishSession(session, execID, status)
}

// streamReader reads from a reader line by line and sends to session.
func (s *Server) streamReader(session *sessionState, reader io.Reader) {
	scanner := bufio.NewScanner(reader)
	scanner.Buffer(make([]byte, 64*1024), 256*1024)
	for scanner.Scan() {
		line := scanner.Text()
		s.sendToSession(session, line)
	}
}

// sendToSession sends a line to all listeners and stores in output buffer.
func (s *Server) sendToSession(session *sessionState, line string) {
	session.mu.Lock()
	defer session.mu.Unlock()

	session.output = append(session.output, line)

	for ch := range session.listeners {
		select {
		case ch <- line:
		default:
			// Skip if channel is full
		}
	}
}

// finishSession marks a session as idle and notifies listeners.
func (s *Server) finishSession(session *sessionState, execID, status string) {
	session.mu.Lock()
	defer session.mu.Unlock()

	session.status = "idle"
	session.cmd = nil

	// Update history entry status
	for i := len(session.history) - 1; i >= 0; i-- {
		if session.history[i].ExecID == execID {
			session.history[i].Status = status
			break
		}
	}

	// Notify listeners that stream is done
	for ch := range session.listeners {
		select {
		case ch <- "[DONE]":
		default:
		}
	}
}

// handleConsoleStream serves an SSE stream for a command execution.
func (s *Server) handleConsoleStream(c *gin.Context) {
	execID := c.Param("id")

	// Find which session has this exec ID
	var session *sessionState
	for i := 0; i < 2; i++ {
		sess := s.console.getSession(i)
		sess.mu.Lock()
		if sess.execID == execID {
			session = sess
			sess.mu.Unlock()
			break
		}
		sess.mu.Unlock()
	}

	if session == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "sesija nije pronađena"})
		return
	}

	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")

	// Create listener channel
	ch := make(chan string, 100)

	session.mu.Lock()
	// Send buffered output first
	for _, line := range session.output {
		fmt.Fprintf(c.Writer, "data: %s\n\n", line)
	}
	c.Writer.Flush()

	// If already done, send done event and return
	if session.status == "idle" && session.execID == execID {
		session.mu.Unlock()
		fmt.Fprintf(c.Writer, "event: done\ndata: finished\n\n")
		c.Writer.Flush()
		return
	}

	session.listeners[ch] = true
	session.mu.Unlock()

	// Clean up on disconnect
	defer func() {
		session.mu.Lock()
		delete(session.listeners, ch)
		session.mu.Unlock()
	}()

	notify := c.Request.Context().Done()

	for {
		select {
		case <-notify:
			return
		case line := <-ch:
			if line == "[DONE]" {
				fmt.Fprintf(c.Writer, "event: done\ndata: finished\n\n")
				c.Writer.Flush()
				return
			}
			fmt.Fprintf(c.Writer, "data: %s\n\n", line)
			c.Writer.Flush()
		}
	}
}

// handleConsoleKill kills the running process in a session.
func (s *Server) handleConsoleKill(c *gin.Context) {
	sessionNum, err := strconv.Atoi(c.Param("session"))
	if err != nil || sessionNum < 1 || sessionNum > 2 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "nevalidna sesija"})
		return
	}

	session := s.console.getSession(sessionNum - 1)

	session.mu.Lock()
	defer session.mu.Unlock()

	if session.status != "running" || session.cmd == nil {
		c.JSON(http.StatusOK, gin.H{"status": "idle", "message": "sesija nije aktivna"})
		return
	}

	if session.cmd.Process != nil {
		session.cmd.Process.Kill()
	}

	// Update history
	for i := len(session.history) - 1; i >= 0; i-- {
		if session.history[i].ExecID == session.execID {
			session.history[i].Status = "killed"
			break
		}
	}

	session.status = "idle"
	session.cmd = nil

	c.JSON(http.StatusOK, gin.H{"status": "killed"})
}

// handleConsoleSessions returns the status of both sessions.
func (s *Server) handleConsoleSessions(c *gin.Context) {
	statuses := make([]sessionStatus, 2)

	for i := 0; i < 2; i++ {
		sess := s.console.getSession(i)
		sess.mu.Lock()
		statuses[i] = sessionStatus{
			Session: i + 1,
			Status:  sess.status,
			TaskID:  sess.taskID,
			ExecID:  sess.execID,
		}
		sess.mu.Unlock()
	}

	c.JSON(http.StatusOK, statuses)
}

// handleConsoleHistory returns command history for a session.
func (s *Server) handleConsoleHistory(c *gin.Context) {
	sessionNum, err := strconv.Atoi(c.Param("session"))
	if err != nil || sessionNum < 1 || sessionNum > 2 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "nevalidna sesija"})
		return
	}

	session := s.console.getSession(sessionNum - 1)

	session.mu.Lock()
	history := make([]historyEntry, len(session.history))
	copy(history, session.history)
	session.mu.Unlock()

	data, _ := json.Marshal(history)
	c.Header("Content-Type", "application/json")
	c.String(http.StatusOK, string(data))
}

// timeNow returns the current time formatted as HH:MM:SS.
func timeNow() string {
	return time.Now().Format("15:04:05")
}

// handleConsolePage serves the console HTML page.
func (s *Server) handleConsolePage(c *gin.Context) {
	c.Header("Content-Type", "text/html; charset=utf-8")
	c.String(http.StatusOK, renderConsolePage())
}