// Copyright 2025 Daytona Platforms Inc.
// SPDX-License-Identifier: AGPL-3.0

package computeruse

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/daytonaio/daemon/pkg/toolbox/computeruse"
	"github.com/godbus/dbus/v5"
	log "github.com/sirupsen/logrus"
)

type Process struct {
	Name        string
	Command     string
	Args        []string
	User        string
	Priority    int
	Env         map[string]string
	LogFile     string
	ErrFile     string
	AutoRestart bool
	cmd         *exec.Cmd
	ctx         context.Context
	cancel      context.CancelFunc
	mu          sync.Mutex
	running     bool
}

type ComputerUse struct {
	processes map[string]*Process
	mu        sync.RWMutex
	configDir string

	// AT-SPI accessibility bus connection. Lazily established on first call to
	// connectA11y(); protected by atspiMu. Implementation lives in accessibility.go.
	atspiMu   sync.Mutex
	atspiConn *dbus.Conn

	a11yHealth func() bool
	waitDBus   func(string, time.Duration) error

	a11yStatusMu        sync.Mutex
	a11yStatusRunning   bool
	a11yStatusCheckedAt time.Time
}

var _ computeruse.IComputerUse = &ComputerUse{}

func (c *ComputerUse) Initialize() (*computeruse.Empty, error) {
	c.processes = make(map[string]*Process)
	// Create config directory for logs
	homeDir, err := os.UserHomeDir()
	if err != nil {
		return new(computeruse.Empty), fmt.Errorf("failed to get home directory: %v", err)
	}
	c.configDir = filepath.Join(homeDir, ".daytona", "computeruse")
	err = os.MkdirAll(c.configDir, 0755)
	if err != nil {
		return new(computeruse.Empty), fmt.Errorf("failed to create config directory: %v", err)
	}

	// Start a D-Bus session and set env vars globally
	cmd := exec.Command("dbus-launch")
	output, err := cmd.Output()
	if err != nil {
		log.Errorf("Failed to start dbus-launch: %v", err)
	} else {
		for _, line := range strings.Split(string(output), "\n") {
			if strings.HasPrefix(line, "DBUS_SESSION_BUS_ADDRESS=") || strings.HasPrefix(line, "DBUS_SESSION_BUS_PID=") {
				parts := strings.SplitN(line, ";", 2)
				for _, part := range parts {
					kv := strings.SplitN(part, "=", 2)
					if len(kv) == 2 {
						os.Setenv(strings.TrimSpace(kv[0]), strings.TrimSpace(kv[1]))
					}
				}
			}
		}
	}

	c.initializeProcesses(homeDir)

	return new(computeruse.Empty), nil
}

func (c *ComputerUse) Start() (*computeruse.Empty, error) {
	// Set DISPLAY environment variable in the main process
	display := os.Getenv("DISPLAY")
	if display == "" {
		display = ":0"
	}
	os.Setenv("DISPLAY", display)
	log.Infof("Set DISPLAY environment variable to: %s", display)

	// Start all processes in order of priority
	c.startAllProcesses()
	c.invalidateA11yStatus()

	// Check process status after starting
	status, err := c.GetProcessStatus()
	if err != nil {
		return nil, fmt.Errorf("failed to get process status after start: %v", err)
	}

	// Check if all required processes are running. atspi is deliberately
	// excluded — a11y failures surface as 503 A11Y_UNAVAILABLE on the /a11y
	// endpoints, and should not block the rest of computer-use from starting.
	required := []string{"xvfb", "xfce4", "x11vnc", "novnc"}
	var failed []string
	for _, name := range required {
		if s, ok := status[name]; !ok || !s.Running {
			failed = append(failed, name)
		}
	}

	if len(failed) > 0 {
		return nil, fmt.Errorf("failed to start: %v", failed)
	}

	return new(computeruse.Empty), nil
}

func (c *ComputerUse) initializeProcesses(homeDir string) {
	// Get environment variables from Dockerfile or use defaults
	vncResolution := os.Getenv("VNC_RESOLUTION")
	if vncResolution == "" {
		vncResolution = "1024x768"
	}

	vncPort := os.Getenv("VNC_PORT")
	if vncPort == "" {
		vncPort = "5901"
	}

	noVncPort := os.Getenv("NO_VNC_PORT")
	if noVncPort == "" {
		noVncPort = "6080"
	}

	display := os.Getenv("DISPLAY")
	if display == "" {
		display = ":0"
	}

	// Get user from environment, fallback to DAYTONA_SANDBOX_USER or default to "root" (just in case, but should not happen)
	user := os.Getenv("VNC_USER")
	if user == "" {
		user = os.Getenv("DAYTONA_SANDBOX_USER")
		if user == "" {
			user = "root"
		}
	}

	// Get D-Bus session address from environment
	dbusAddress := os.Getenv("DBUS_SESSION_BUS_ADDRESS")

	// Process 1: Xvfb (X Virtual Framebuffer)
	c.processes["xvfb"] = &Process{
		Name:        "xvfb",
		Command:     "/usr/bin/Xvfb",
		Args:        []string{display, "-screen", "0", vncResolution + "x24"},
		User:        user,
		Priority:    100,
		AutoRestart: true,
		Env: map[string]string{
			"DISPLAY": display,
		},
		LogFile: filepath.Join(c.configDir, "xvfb.log"),
		ErrFile: filepath.Join(c.configDir, "xvfb.err"),
	}

	// Process 2: xfce4 (Desktop Environment)
	c.processes["xfce4"] = &Process{
		Name:        "xfce4",
		Command:     "/usr/bin/startxfce4",
		Args:        []string{},
		User:        user,
		Priority:    200,
		AutoRestart: true,
		Env: map[string]string{
			"DISPLAY":                  display,
			"HOME":                     homeDir,
			"USER":                     user,
			"DBUS_SESSION_BUS_ADDRESS": dbusAddress,
		},
		LogFile: filepath.Join(c.configDir, "xfce4.log"),
		ErrFile: filepath.Join(c.configDir, "xfce4.err"),
	}

	// Process 2.5: at-spi-bus-launcher (AT-SPI daemon for the accessibility API)
	// Launches the org.a11y.Bus service so GTK/Qt/Electron apps can publish
	// their widget trees. It uses the session D-Bus created during Initialize.
	// Path lives in /usr/libexec on Debian/Ubuntu; some distros ship it under
	// /usr/lib/at-spi2-core/; as a last resort fall back to $PATH so images
	// that put the binary somewhere unusual still work. If we can't find the
	// binary anywhere, we skip registering atspi entirely — the a11y endpoints
	// will return 503 cleanly instead of thrashing on a nonexistent command.
	atspiCommand := ""
	if _, err := os.Stat("/usr/libexec/at-spi-bus-launcher"); err == nil {
		atspiCommand = "/usr/libexec/at-spi-bus-launcher"
	} else if _, err := os.Stat("/usr/lib/at-spi2-core/at-spi-bus-launcher"); err == nil {
		atspiCommand = "/usr/lib/at-spi2-core/at-spi-bus-launcher"
	} else if p, err := exec.LookPath("at-spi-bus-launcher"); err == nil {
		atspiCommand = p
	}
	if atspiCommand == "" {
		log.Warnf("at-spi-bus-launcher not found in any known location; accessibility API will return 503 until the binary is installed")
	} else if dbusAddress == "" {
		log.Warnf("DBUS_SESSION_BUS_ADDRESS is empty; accessibility API will return 503 until session D-Bus is available")
	} else {
		waitDBus := waitForSessionBus
		if c.waitDBus != nil {
			waitDBus = c.waitDBus
		}
		if err := waitDBus(dbusAddress, 5*time.Second); err != nil {
			log.Warnf("session D-Bus check failed for at-spi-bus-launcher; launching anyway so the a11y bus can start if D-Bus becomes usable: %v", err)
		}
		c.processes["atspi"] = &Process{
			Name:        "atspi",
			Command:     atspiCommand,
			Args:        []string{"--launch-immediately"},
			User:        user,
			Priority:    250,
			AutoRestart: false,
			Env: map[string]string{
				"DISPLAY":                  display,
				"HOME":                     homeDir,
				"USER":                     user,
				"DBUS_SESSION_BUS_ADDRESS": dbusAddress,
			},
			LogFile: filepath.Join(c.configDir, "atspi.log"),
			ErrFile: filepath.Join(c.configDir, "atspi.err"),
		}
	}

	// Process 3: x11vnc (VNC Server)
	c.processes["x11vnc"] = &Process{
		Name:        "x11vnc",
		Command:     "/usr/bin/x11vnc",
		Args:        []string{"-display", display, "-forever", "-shared", "-rfbport", vncPort},
		User:        user,
		Priority:    300,
		AutoRestart: true,
		Env: map[string]string{
			"DISPLAY": display,
		},
		LogFile: filepath.Join(c.configDir, "x11vnc.log"),
		ErrFile: filepath.Join(c.configDir, "x11vnc.err"),
	}
	// Determine the best available NoVNC command with fallback options
	var novncCommand string
	var novncArgs []string

	// Priority 1: Try launch.sh (modern NoVNC with enhanced features)
	if _, err := os.Stat("/usr/share/novnc/utils/launch.sh"); err == nil {
		novncCommand = "/usr/share/novnc/utils/launch.sh"
		novncArgs = []string{"--vnc", "localhost:" + vncPort, "--listen", noVncPort}
		log.Infof("Using NoVNC launch.sh (recommended)")
	} else if _, err := os.Stat("/usr/share/novnc/utils/novnc_proxy"); err == nil {
		// Priority 2: Try novnc_proxy (legacy NoVNC script)
		novncCommand = "/usr/share/novnc/utils/novnc_proxy"
		novncArgs = []string{"--vnc", "localhost:" + vncPort, "--listen", noVncPort}
		log.Infof("Using NoVNC novnc_proxy (legacy)")
	} else {
		// Priority 3: Fallback to direct websockify (always available)
		novncCommand = "websockify"
		novncArgs = []string{"--web=/usr/share/novnc/", noVncPort, "localhost:" + vncPort}
		log.Infof("Using direct websockify (fallback)")
	}

	c.processes["novnc"] = &Process{
		Name:        "novnc",
		Command:     novncCommand,
		Args:        novncArgs,
		User:        user,
		Priority:    400,
		AutoRestart: true,
		Env: map[string]string{
			"DISPLAY": display,
		},
		LogFile: filepath.Join(c.configDir, "novnc.log"),
		ErrFile: filepath.Join(c.configDir, "novnc.err"),
	}
}

func waitForSessionBus(address string, timeout time.Duration) error {
	if address == "" {
		return fmt.Errorf("DBUS_SESSION_BUS_ADDRESS is empty")
	}

	deadline := time.Now().Add(timeout)
	var lastErr error
	for {
		conn, err := dbus.Connect(address)
		if err == nil {
			_ = conn.Close()
			return nil
		}
		lastErr = err

		if time.Now().After(deadline) {
			return fmt.Errorf("session D-Bus did not become ready: %w", lastErr)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func (c *ComputerUse) startAllProcesses() {
	// Sort processes by priority and start them
	processes := c.getProcessesByPriority()

	for _, process := range processes {
		go c.startProcess(process)
		// Wait a bit between starting processes to ensure proper initialization
		time.Sleep(2 * time.Second)
	}
}

func (c *ComputerUse) getProcessesByPriority() []*Process {
	c.mu.RLock()
	defer c.mu.RUnlock()

	var processes []*Process
	for _, p := range c.processes {
		processes = append(processes, p)
	}

	// Sort by priority (lower number = higher priority)
	for i := 0; i < len(processes)-1; i++ {
		for j := i + 1; j < len(processes); j++ {
			if processes[i].Priority > processes[j].Priority {
				processes[i], processes[j] = processes[j], processes[i]
			}
		}
	}

	return processes
}

func (c *ComputerUse) startProcess(process *Process) {
	process.mu.Lock()
	if process.running {
		process.mu.Unlock()
		return
	}
	process.running = true
	process.mu.Unlock()

	for {
		log.Infof("Starting process: %s", process.Name)

		// Create context for the process
		process.ctx, process.cancel = context.WithCancel(context.Background())

		// Create command
		process.cmd = exec.CommandContext(process.ctx, process.Command, process.Args...)

		// Set environment variables
		if len(process.Env) > 0 {
			process.cmd.Env = os.Environ()
			for key, value := range process.Env {
				process.cmd.Env = append(process.cmd.Env, fmt.Sprintf("%s=%s", key, value))
			}
		}

		// Set up logging
		if process.LogFile != "" {
			logFile, err := os.OpenFile(process.LogFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
			if err != nil {
				log.Errorf("Failed to open log file for %s: %v", process.Name, err)
			} else {
				process.cmd.Stdout = logFile
				defer logFile.Close()
			}
		}

		if process.ErrFile != "" {
			errFile, err := os.OpenFile(process.ErrFile, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
			if err != nil {
				log.Errorf("Failed to open error file for %s: %v", process.Name, err)
			} else {
				process.cmd.Stderr = errFile
				defer errFile.Close()
			}
		}

		// Start the process
		err := process.cmd.Start()
		if err != nil {
			log.Errorf("Failed to start process %s: %v", process.Name, err)
			if !process.AutoRestart {
				break
			}
			time.Sleep(5 * time.Second)
			continue
		}

		log.Infof("Process %s started with PID: %d", process.Name, process.cmd.Process.Pid)

		// Wait for the process to finish
		err = process.cmd.Wait()
		if err != nil {
			log.Errorf("Process %s exited with error: %v", process.Name, err)
		} else {
			log.Infof("Process %s exited normally", process.Name)
		}

		// Check if we should restart
		if !process.AutoRestart {
			break
		}

		log.Infof("Restarting process %s in 5 seconds...", process.Name)
		time.Sleep(5 * time.Second)
	}

	process.mu.Lock()
	process.running = false
	process.mu.Unlock()
}

func (c *ComputerUse) Stop() (*computeruse.Empty, error) {
	log.Info("Stopping all computer use processes...")

	c.mu.RLock()
	processes := make([]*Process, 0, len(c.processes))
	for _, p := range c.processes {
		processes = append(processes, p)
	}
	c.mu.RUnlock()

	// Stop processes in reverse priority order
	for i := len(processes) - 1; i >= 0; i-- {
		process := processes[i]
		c.stopProcess(process)
	}

	// Release the cached AT-SPI bus connection so a later Start() dials a
	// fresh one — the bus address can change across launcher restarts.
	c.atspiMu.Lock()
	if c.atspiConn != nil {
		_ = c.atspiConn.Close()
		c.atspiConn = nil
	}
	c.atspiMu.Unlock()
	c.setA11yStatus(false)

	return new(computeruse.Empty), nil
}

func (c *ComputerUse) stopProcess(process *Process) {
	process.mu.Lock()
	defer process.mu.Unlock()

	if !process.running {
		return
	}

	log.Infof("Stopping process: %s", process.Name)

	// Cancel the context to stop the process
	if process.cancel != nil {
		process.cancel()
	}

	// Kill the process if it's still running
	if process.cmd != nil && process.cmd.Process != nil {
		err := process.cmd.Process.Kill()
		if err != nil {
			log.Errorf("Failed to kill process %s: %v", process.Name, err)
		}
	}

	process.running = false
}

func (c *ComputerUse) GetProcessStatus() (map[string]computeruse.ProcessStatus, error) {
	c.mu.RLock()
	processes := make(map[string]*Process, len(c.processes))
	for name, process := range c.processes {
		processes[name] = process
	}
	c.mu.RUnlock()

	status := make(map[string]computeruse.ProcessStatus, len(processes))
	for name, process := range processes {
		processStatus := getProcessStatus(process)
		if name == "atspi" {
			processStatus.Running = c.cachedA11yAvailable()
			processStatus.Pid = nil
		}
		status[name] = processStatus
	}

	return status, nil
}

func getProcessStatus(process *Process) computeruse.ProcessStatus {
	process.mu.Lock()
	defer process.mu.Unlock()

	status := computeruse.ProcessStatus{
		Running:     false,
		Priority:    process.Priority,
		AutoRestart: process.AutoRestart,
	}

	if process.cmd != nil && process.cmd.Process != nil {
		if err := process.cmd.Process.Signal(syscall.Signal(0)); err == nil {
			pid := process.cmd.Process.Pid
			status.Running = true
			status.Pid = &pid
		}
	}

	return status
}

// IsProcessRunning checks if a specific process is running
func (c *ComputerUse) IsProcessRunning(req *computeruse.ProcessRequest) (bool, error) {
	c.mu.RLock()
	process, exists := c.processes[req.ProcessName]
	c.mu.RUnlock()
	if !exists {
		return false, fmt.Errorf("process %s not found", req.ProcessName)
	}
	if req.ProcessName == "atspi" {
		return c.cachedA11yAvailable(), nil
	}

	process.mu.Lock()
	defer process.mu.Unlock()
	return process.running, nil
}

// RestartProcess restarts a specific process
func (c *ComputerUse) RestartProcess(req *computeruse.ProcessRequest) (*computeruse.Empty, error) {
	c.mu.RLock()
	process, exists := c.processes[req.ProcessName]
	c.mu.RUnlock()

	if !exists {
		return nil, fmt.Errorf("process %s not found", req.ProcessName)
	}

	// Stop the process first
	c.stopProcess(process)

	// Wait a moment for cleanup
	time.Sleep(1 * time.Second)

	// Start the process again
	go c.startProcess(process)

	return new(computeruse.Empty), nil
}

// GetProcessLogs returns the logs for a specific process
func (c *ComputerUse) GetProcessLogs(req *computeruse.ProcessRequest) (string, error) {
	c.mu.RLock()
	process, exists := c.processes[req.ProcessName]
	c.mu.RUnlock()

	if !exists {
		return "", fmt.Errorf("process %s not found", req.ProcessName)
	}

	if process.LogFile == "" {
		return "", fmt.Errorf("no log file configured for process %s", req.ProcessName)
	}

	content, err := os.ReadFile(process.LogFile)
	if err != nil {
		return "", fmt.Errorf("failed to read log file for %s: %v", req.ProcessName, err)
	}

	return string(content), nil
}

// GetProcessErrors returns the error logs for a specific process
func (c *ComputerUse) GetProcessErrors(req *computeruse.ProcessRequest) (string, error) {
	c.mu.RLock()
	process, exists := c.processes[req.ProcessName]
	c.mu.RUnlock()

	if !exists {
		return "", fmt.Errorf("process %s not found", req.ProcessName)
	}

	if process.ErrFile == "" {
		return "", fmt.Errorf("no error file configured for process %s", req.ProcessName)
	}

	content, err := os.ReadFile(process.ErrFile)
	if err != nil {
		return "", fmt.Errorf("failed to read error file for %s: %v", req.ProcessName, err)
	}

	return string(content), nil
}

func (c *ComputerUse) GetStatus() (*computeruse.ComputerUseStatusResponse, error) {
	// Get the current process status
	processStatus, err := c.GetProcessStatus()
	if err != nil {
		return &computeruse.ComputerUseStatusResponse{
			Status: "error",
		}, err
	}

	// Check if all required processes are running
	requiredProcesses := []string{"xvfb", "xfce4", "x11vnc", "novnc"}
	allRunning := true

	for _, processName := range requiredProcesses {
		if status, exists := processStatus[processName]; !exists || !status.Running {
			allRunning = false
			break
		}
	}

	if allRunning {
		return &computeruse.ComputerUseStatusResponse{
			Status: "active",
		}, nil
	}

	// Check if any processes are running
	anyRunning := false
	for _, status := range processStatus {
		if status.Running {
			anyRunning = true
			break
		}
	}

	if anyRunning {
		return &computeruse.ComputerUseStatusResponse{
			Status: "partial",
		}, nil
	}

	return &computeruse.ComputerUseStatusResponse{
		Status: "inactive",
	}, nil
}
