package agent

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/fleetforge/forge-agent/internal/client"
	"github.com/fleetforge/forge-agent/internal/config"
	"github.com/fleetforge/forge-agent/internal/inventory"
	"github.com/fleetforge/forge-agent/internal/metrics"
	"github.com/fleetforge/forge-agent/internal/pkgmanager"
	"github.com/fleetforge/forge-agent/internal/svcmgr"
	"github.com/fleetforge/forge-agent/internal/tunnel"
	"github.com/fleetforge/forge-agent/internal/vnc"
)

const AgentVersion = "0.1.6"

var (
	tunnelMu sync.Mutex
	tunnels  = map[string]*tunnel.Tunnel{}
)

func storeTunnel(sessionID string, t *tunnel.Tunnel) {
	tunnelMu.Lock()
	defer tunnelMu.Unlock()
	tunnels[sessionID] = t
}

func removeTunnel(sessionID string) {
	tunnelMu.Lock()
	defer tunnelMu.Unlock()
	delete(tunnels, sessionID)
}

func stopTunnel(sessionID string) {
	tunnelMu.Lock()
	t := tunnels[sessionID]
	tunnelMu.Unlock()
	if t != nil {
		t.Stop()
	}
}

type Agent struct {
	cfg    *config.Config
	client *client.Client
}

func New(cfg *config.Config) *Agent {
	return &Agent{
		cfg:    cfg,
		client: client.New(cfg.ServerURL, cfg.Credential),
	}
}

func (a *Agent) Run() error {
	log.Println("Forge Agent starting...")
	log.Printf("Server: %s", a.cfg.ServerURL)
	log.Printf("Device ID: %s", a.cfg.DeviceID)

	ticker := time.NewTicker(time.Duration(a.cfg.HeartbeatSec) * time.Second)
	defer ticker.Stop()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	if a.cfg.DeviceID == "" {
		log.Println("No DeviceID configured - waiting for enrollment")
	} else {
		if err := pkgmanager.BootstrapManagers(); err != nil {
			log.Printf("WARNING: Manager bootstrap failed: %v", err)
		}
		a.sendHeartbeat()
	}

	for {
		select {
		case <-ticker.C:
			if a.cfg.DeviceID == "" || a.cfg.Credential == "" {
				a.tryEnroll()
			} else {
				a.sendHeartbeat()
				a.checkJobs()
			}
		case sig := <-sigCh:
			log.Printf("Received signal %v, shutting down...", sig)
			return nil
		}
	}
}

func (a *Agent) tryEnroll() {
	if a.cfg.EnrollmentToken == "" {
		return
	}
	log.Println("Attempting auto-enrollment...")

	info, err := inventory.CollectSystemInfo()
	if err != nil {
		log.Printf("Enrollment failed (inventory): %v", err)
		return
	}

	c := client.New(a.cfg.ServerURL, "")
	resp, err := c.Enroll(client.EnrollRequest{
		Token:            a.cfg.EnrollmentToken,
		Hostname:         info.Hostname,
		Platform:         info.Platform,
		OSName:           info.OSName,
		OSVersion:        info.OSVersion,
		Architecture:     info.Architecture,
		KernelVersion:    info.KernelVersion,
		CPUModel:         info.CPUModel,
		CPUCores:         info.CPUCores,
		TotalMemoryBytes: info.TotalMemory,
		TotalDiskBytes:   info.TotalDisk,
		IPAddress:        info.IPAddress,
		AgentVersion:     AgentVersion,
	})
	if err != nil {
		log.Printf("Enrollment failed: %v", err)
		return
	}

	a.cfg.DeviceID = resp.DeviceID
	a.cfg.Credential = resp.Credential
	a.cfg.EnrollmentToken = ""
	if err := config.Save(a.cfg); err != nil {
		log.Printf("Failed to save config after enrollment: %v", err)
		return
	}
	a.client = client.New(a.cfg.ServerURL, a.cfg.Credential)
	log.Printf("Enrollment successful! Device ID: %s", resp.DeviceID)
}

func (a *Agent) sendHeartbeat() {
	info, err := inventory.CollectSystemInfo()
	if err != nil {
		log.Printf("Failed to collect system info: %v", err)
		return
	}

	met, err := metrics.Collect()
	if err != nil {
		log.Printf("Failed to collect metrics: %v", err)
		return
	}

	req := client.HeartbeatRequest{
		Hostname:      info.Hostname,
		OSVersion:     info.OSVersion,
		KernelVersion: info.KernelVersion,
		AgentVersion:  AgentVersion,
		IPAddress:     info.IPAddress,
		Metrics: &client.MetricsData{
			CPUPercent:    met.CPUPercent,
			MemoryPercent: met.MemoryPercent,
			MemoryUsed:    met.MemoryUsed,
			MemoryTotal:   met.MemoryTotal,
			DiskPercent:   met.DiskPercent,
			DiskUsed:      met.DiskUsed,
			DiskTotal:     met.DiskTotal,
			Load1:         met.Load1,
			Load5:         met.Load5,
			Load15:        met.Load15,
			UptimeSeconds: met.UptimeSeconds,
		},
	}

	if len(info.Packages) > 0 || len(info.Services) > 0 || len(info.NetworkInterfaces) > 0 {
		inv := &client.InventoryData{}
		for _, p := range info.Packages {
			inv.Packages = append(inv.Packages, client.PackageData{
				Name:    p.Name,
				Version: p.Version,
				Manager: p.Manager,
			})
		}
		for _, s := range info.Services {
			inv.Services = append(inv.Services, client.ServiceData{
				Name:    s.Name,
				Status:  s.Status,
				Enabled: s.Enabled,
			})
		}
		for _, ni := range info.NetworkInterfaces {
			inv.NetworkInterfaces = append(inv.NetworkInterfaces, client.NetworkInterfaceData{
				Name:       ni.Name,
				Addresses:  ni.Addresses,
				MacAddress: ni.MAC,
			})
		}
		req.Inventory = inv
	}

	if err := a.client.Heartbeat(req); err != nil {
		log.Printf("Heartbeat failed: %v", err)
	} else {
		log.Printf("Heartbeat sent successfully")
	}
}

func (a *Agent) checkJobs() {
	jobs, err := a.client.FetchJobs()
	if err != nil {
		log.Printf("Failed to fetch jobs: %v", err)
		return
	}

	for _, job := range jobs {
		a.executeJob(job)
	}
}

func (a *Agent) executeJob(job client.JobInfo) {
	log.Printf("Executing job %s (type: %s)", job.JobID, job.JobType)

	result := client.JobResult{
		Status: "succeeded",
		Result: map[string]interface{}{},
	}
	result.TargetID = job.TargetID

	if err := a.client.ReportStart(job.JobID, job.TargetID); err != nil {
		log.Printf("Failed to report start for job %s: %v", job.JobID, err)
	}

	switch job.JobType {
	case "update_inventory":
		a.sendHeartbeat()
		result.Result["inventory_refreshed"] = true

	case "update_system_info":
		met, err := metrics.Collect()
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		} else {
			result.Result["cpu_percent"] = met.CPUPercent
			result.Result["memory_percent"] = met.MemoryPercent
			result.Result["disk_percent"] = met.DiskPercent
			result.Result["uptime"] = met.UptimeSeconds
		}

	case "install_package":
		pkgName, _ := job.Parameters["package_name"].(string)
		if pkgName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: package_name"
			break
		}
		output, err := pkgmanager.InstallPackage(pkgName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["package"] = pkgName
		result.Result["output"] = truncateOutput(output)

	case "update_package":
		pkgName, _ := job.Parameters["package_name"].(string)
		if pkgName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: package_name"
			break
		}
		output, err := pkgmanager.UpdatePackage(pkgName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["package"] = pkgName
		result.Result["output"] = truncateOutput(output)

	case "remove_package":
		pkgName, _ := job.Parameters["package_name"].(string)
		if pkgName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: package_name"
			break
		}
		source, _ := job.Parameters["source"].(string)
		var manager string
		switch source {
		case "winget":
			manager = "winget"
		case "chocolatey":
			manager = "choco"
		}
		var output string
		var err error
		if manager != "" {
			output, err = pkgmanager.RemovePackageWithManager(pkgName, manager)
		} else {
			output, err = pkgmanager.RemovePackage(pkgName)
		}
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["package"] = pkgName
		result.Result["output"] = truncateOutput(output)

	case "start_service":
		svcName, _ := job.Parameters["service_name"].(string)
		if svcName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: service_name"
			break
		}
		output, err := svcmgr.StartService(svcName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["service"] = svcName
		result.Result["output"] = truncateOutput(output)

	case "stop_service":
		svcName, _ := job.Parameters["service_name"].(string)
		if svcName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: service_name"
			break
		}
		output, err := svcmgr.StopService(svcName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["service"] = svcName
		result.Result["output"] = truncateOutput(output)

	case "restart_service":
		svcName, _ := job.Parameters["service_name"].(string)
		if svcName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: service_name"
			break
		}
		output, err := svcmgr.RestartService(svcName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["service"] = svcName
		result.Result["output"] = truncateOutput(output)

	case "check_service_status":
		svcName, _ := job.Parameters["service_name"].(string)
		if svcName == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: service_name"
			break
		}
		status, enabled, err := svcmgr.GetServiceStatus(svcName)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["service"] = svcName
		result.Result["status"] = status
		result.Result["enabled"] = enabled

	case "reboot_system":
		log.Println("System reboot requested, initiating reboot...")
		result.Result["message"] = "reboot initiated"
		if err := a.client.SubmitResult(job.JobID, result); err != nil {
			log.Printf("Failed to submit result for reboot job: %v", err)
			return
		}
		log.Println("Rebooting system in 5 seconds...")
		time.Sleep(5 * time.Second)
		svcmgr.Reboot()
		return

	case "shutdown_system":
		log.Println("System shutdown requested, initiating shutdown...")
		result.Result["message"] = "shutdown initiated"
		if err := a.client.SubmitResult(job.JobID, result); err != nil {
			log.Printf("Failed to submit result for shutdown job: %v", err)
			return
		}
		log.Println("Shutting down system in 5 seconds...")
		time.Sleep(5 * time.Second)
		svcmgr.Shutdown()
		return

	case "install_registry_package":
		source, _ := job.Parameters["source"].(string)
		pkgID, _ := job.Parameters["package_id"].(string)
		if source == "" || pkgID == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: source/package_id"
			break
		}
		var manager string
		switch source {
		case "winget":
			manager = "winget"
		case "chocolatey":
			manager = "choco"
		default:
			result.Status = "failed"
			result.ErrorMessage = fmt.Sprintf("unsupported source: %s", source)
			break
		}
		if result.Status == "failed" {
			break
		}
		version, _ := job.Parameters["package_version"].(string)
		output, err := pkgmanager.InstallPackageWithManager(pkgID, manager, version)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["package"] = pkgID
		result.Result["source"] = source
		result.Result["output"] = truncateOutput(output)

	case "start_remote_session":
		sessionID, _ := job.Parameters["session_id"].(string)
		if sessionID == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: session_id"
			break
		}
		port := 5900
		if p, ok := job.Parameters["vnc_port"].(float64); ok && p > 0 {
			port = int(p)
		}
		password, _ := job.Parameters["vnc_password"].(string)

		vncServer, err := vnc.Ensure(port, password)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
			break
		}
		log.Printf("VNC server started on port %d for session %s", port, sessionID)

		result.Result["session_id"] = sessionID
		result.Result["vnc_port"] = port
		result.Result["vnc_started"] = true

		// Keep the session alive in a background goroutine; VNC is stopped when
		// the tunnel closes.
		go func(sid string, p int, vs *vnc.Server) {
			t := tunnel.New(a.cfg.ServerURL, a.cfg.Credential, sid, p)
			storeTunnel(sid, t)
			defer func() {
				removeTunnel(sid)
				vs.Stop()
				log.Printf("VNC server stopped after tunnel close for session %s", sid)
			}()
			if err := t.Run(); err != nil {
				log.Printf("Tunnel for session %s closed: %v", sid, err)
			}
		}(sessionID, port, vncServer)

	case "stop_remote_session":
		sessionID, _ := job.Parameters["session_id"].(string)
		if sessionID == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: session_id"
			break
		}
		stopTunnel(sessionID)
		result.Result["session_id"] = sessionID
		result.Result["vnc_stopped"] = true

	case "install_from_file":
		filePath, _ := job.Parameters["file_path"].(string)
		if filePath == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: file_path"
			break
		}
		output, err := pkgmanager.InstallPackageFromFile(filePath)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		result.Result["file_path"] = filePath
		result.Result["output"] = truncateOutput(output)

	case "download_and_install":
		downloadURL, _ := job.Parameters["download_url"].(string)
		packageId, _ := job.Parameters["package_id"].(string)
		if downloadURL == "" && packageId == "" {
			result.Status = "failed"
			result.ErrorMessage = "missing parameter: download_url or package_id"
			break
		}
		if downloadURL == "" && packageId != "" {
			downloadURL = a.cfg.ServerURL + "/api/v1/agent/software/" + packageId + "/download"
		}
		tmpDir := os.TempDir()
		destPath := tmpDir + "/forge-download-" + packageId + ".pkg"
		log.Printf("Downloading package from %s to %s", downloadURL, destPath)
		if err := pkgmanager.DownloadFile(downloadURL, destPath); err != nil {
			result.Status = "failed"
			result.ErrorMessage = fmt.Sprintf("download failed: %v", err)
			break
		}
		output, err := pkgmanager.InstallPackageFromFile(destPath)
		if err != nil {
			result.Status = "failed"
			result.ErrorMessage = err.Error()
		}
		os.Remove(destPath)
		result.Result["download_url"] = downloadURL
		result.Result["output"] = truncateOutput(output)

	default:
		result.Status = "failed"
		result.ErrorMessage = fmt.Sprintf("unsupported job type: %s", job.JobType)
	}

	if a.client.IsJobCancelled(job.JobID) {
		log.Printf("Job %s cancelled on server; discarding result", job.JobID)
		result.Status = "cancelled"
	}

	if err := a.client.SubmitResult(job.JobID, result); err != nil {
		log.Printf("Failed to submit result for job %s: %v", job.JobID, err)
	} else {
		log.Printf("Job %s completed with status: %s", job.JobID, result.Status)
	}
}

func truncateOutput(output string) string {
	const maxLen = 4096
	if len(output) > maxLen {
		return output[:maxLen] + "... (truncated)"
	}
	return output
}