package main

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/jaredfolkins/terraform-loading-bar/progress"
	compute "google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/iamcredentials/v1"
	"google.golang.org/api/option"
)

// GCPCredentials represents the structure of a Google Cloud Platform service account key
type GCPCredentials struct {
	Type                    string `json:"type"`
	ProjectID               string `json:"project_id"`
	PrivateKeyID            string `json:"private_key_id"`
	PrivateKey              string `json:"private_key"`
	ClientEmail             string `json:"client_email"`
	ClientID                string `json:"client_id"`
	AuthURI                 string `json:"auth_uri"`
	TokenURI                string `json:"token_uri"`
	AuthProviderX509CertURL string `json:"auth_provider_x509_cert_url"`
	ClientX509CertURL       string `json:"client_x509_cert_url"`
}

var (
	sourceDir = "terraform-config"
)

// TerraformFunction represents the type of Terraform operation to perform
type TerraformFunction string

const (
	FunctionApply     TerraformFunction = "apply"
	FunctionDestroy   TerraformFunction = "destroy"
	FunctionImage     TerraformFunction = "image"
	publicDir                           = "/lemc/public"
	destDir                             = "/lemc/private"
	tfVarsFilename                      = "terraform.tfvars"
	installScriptName                   = "install.sh"
)

// Config holds all the configuration values needed for the application
type Config struct {
	// Directory Configuration
	SourceDir         string
	PublicDir         string
	DestDir           string
	WorkingDir        string
	AbsoluteSourceDir string
	AbsoluteDestDir   string

	// Terraform Configuration
	TerraformFunction string

	// GCP Configuration
	GCPProjectID   string
	GCPRegion      string
	GCPZone        string
	GCPCredentials []byte
	GCPKeyFilePath string

	// LEMC Configuration
	LEMCUUID        string
	LEMCScope       string
	LEMCUsername    string
	LEMCUserID      string
	LEMCMachineType string
	LEMCHTTPBaseURL string

	// Domain Configuration
	RootDomain     string
	RootZone       string
	DomainName     string
	ResourcePrefix string

	// VM Configuration
	ImageName             string
	TargetImageName       string
	Port                  string
	BootDiskSizeGB        string
	BackendTimeoutSeconds string

	// Existing SSL Certificate
	ExistingSSLCertificateName string

	// Docker Config
	GCRHostname       string
	DockerImageToPull string

	// Optional workload configuration.
	WorkloadToken string
}

// serviceAccountKey is used to unmarshal the client_email from the SA key JSON
type serviceAccountKey struct {
	ClientEmail string `json:"client_email"`
}

// GetConfig validates and returns the application configuration from environment variables
func GetConfig() (*Config, error) {
	config := &Config{}

	// Set up directory configuration
	config.SourceDir = "terraform-config"
	config.PublicDir = "/lemc/public"
	config.DestDir = "/lemc/private"

	// Get current working directory

	cwd, err := os.Getwd()
	if err != nil {
		return nil, fmt.Errorf("error getting current working directory: %v", err)
	}
	config.WorkingDir = cwd

	// Resolve absolute paths
	if !filepath.IsAbs(config.SourceDir) {
		config.AbsoluteSourceDir = filepath.Join(cwd, config.SourceDir)
	} else {
		config.AbsoluteSourceDir = config.SourceDir
	}

	if !filepath.IsAbs(config.DestDir) {
		config.AbsoluteDestDir = filepath.Join(cwd, config.DestDir)
	} else {
		config.AbsoluteDestDir = config.DestDir
	}

	// Create required directories
	if err := os.MkdirAll(config.AbsoluteDestDir, os.ModePerm); err != nil {
		return nil, fmt.Errorf("error creating destination directory %s: %v", config.AbsoluteDestDir, err)
	}

	if err := os.MkdirAll(config.PublicDir, os.ModePerm); err != nil {
		return nil, fmt.Errorf("error creating public directory %s: %v", config.PublicDir, err)
	}

	// GCP Configuration
	config.GCPKeyFilePath = os.Getenv("GCP_KEY_FILE")
	googleCredsEnv := os.Getenv("GOOGLE_CREDENTIALS")
	if googleCredsEnv != "" {
		config.GCPCredentials = []byte(googleCredsEnv)
		// Parse the credentials JSON to extract project_id
		var creds GCPCredentials
		if err := json.Unmarshal([]byte(googleCredsEnv), &creds); err != nil {
			return nil, fmt.Errorf("error parsing GOOGLE_CREDENTIALS JSON: %v", err)
		}
		if creds.ProjectID != "" {
			config.GCPProjectID = creds.ProjectID
			lg(fmt.Sprintf("Using project_id from GOOGLE_CREDENTIALS: %s", config.GCPProjectID))
		}
	}
	// Only use environment variable if not set from credentials
	if config.GCPProjectID == "" {
		config.GCPProjectID = os.Getenv("GCP_PROJECT_ID")
	}
	config.GCPRegion = os.Getenv("GCP_REGION")
	config.GCPZone = os.Getenv("GCP_ZONE")

	// LEMC Configuration
	config.LEMCUUID = os.Getenv("LEMC_UUID")
	if config.LEMCUUID == "" {
		config.LEMCUUID = "default-uuid"
	}
	config.LEMCScope = os.Getenv("LEMC_SCOPE")
	config.LEMCUsername = os.Getenv("LEMC_USERNAME")
	config.LEMCUserID = os.Getenv("LEMC_USER_ID")
	config.LEMCMachineType = os.Getenv("MACHINE_TYPE")
	config.LEMCHTTPBaseURL = os.Getenv("LEMC_HTTP_DOWNLOAD_BASE_URL")

	// Domain Configuration
	config.RootDomain = os.Getenv("ROOT_DOMAIN")
	config.RootZone = os.Getenv("ROOT_ZONE")

	// VM Configuration
	config.ImageName = os.Getenv("CUSTOM_IMAGE")
	if config.ImageName == "" {
		config.ImageName = os.Getenv("AMI")
	}
	if config.ImageName == "" {
		config.ImageName = os.Getenv("IMAGE")
	}
	config.TargetImageName = os.Getenv("IMAGE_NAME")
	config.Port = os.Getenv("PORT_FROM")
	config.BootDiskSizeGB = os.Getenv("BOOT_DISK_SIZE_GB")
	config.BackendTimeoutSeconds = os.Getenv("BACKEND_TIMEOUT_SECONDS")
	if config.BackendTimeoutSeconds == "" {
		config.BackendTimeoutSeconds = "7200" // Default to 7200 seconds
	}

	// Existing SSL Certificate
	config.ExistingSSLCertificateName = os.Getenv("EXISTING_SSL_CERTIFICATE_NAME")

	// Docker Config
	config.GCRHostname = os.Getenv("GCR_HOSTNAME")
	if config.GCRHostname == "" {
		config.GCRHostname = "us-central1-docker.pkg.dev" // Default
	}
	config.DockerImageToPull = os.Getenv("DOCKER_IMAGE_TO_PULL")

	// Optional workload configuration.
	config.WorkloadToken = os.Getenv("WORKLOAD_TOKEN")

	// Validate required fields
	if config.LEMCHTTPBaseURL == "" {
		return nil, fmt.Errorf("LEMC_HTTP_DOWNLOAD_BASE_URL environment variable is required")
	}

	// Validate LEMCMachineType only if TERRAFORM_FUNCTION is "apply"
	terraformFunction := os.Getenv("TERRAFORM_FUNCTION")
	if (terraformFunction == string(FunctionApply) || terraformFunction == "") && config.LEMCMachineType == "" { // Default to apply if not specified
		return nil, fmt.Errorf("MACHINE_TYPE environment variable is required when TERRAFORM_FUNCTION is 'apply'")
	}

	// Construct resource prefix
	uuidPrefix := config.LEMCUUID
	if len(uuidPrefix) >= 8 {
		uuidPrefix = uuidPrefix[:8]
	}

	scopePrefix := strings.ToLower(config.LEMCScope)
	if scopePrefix == "individual" {
		scopePrefix = "ind"
	} else if scopePrefix == "shared" {
		scopePrefix = "shd"
	}

	config.ResourcePrefix = fmt.Sprintf("lemc-%s-%s-%s-%s",
		uuidPrefix,
		config.LEMCUsername,
		config.LEMCUserID,
		scopePrefix)

	// Construct domain name if possible
	if config.RootDomain != "" && config.ResourcePrefix != "" {
		config.DomainName = fmt.Sprintf("%s.%s", config.ResourcePrefix, config.RootDomain)
	} else {
		config.DomainName = "default.example.com" // Fallback
	}

	// Ensure LEMC_HTTP_DOWNLOAD_BASE_URL ends with a slash
	if !strings.HasSuffix(config.LEMCHTTPBaseURL, "/") {
		config.LEMCHTTPBaseURL += "/"
	}

	return config, nil
}

// AdjustBootDiskSize checks the source image size and increases BootDiskSizeGB if necessary
func (c *Config) AdjustBootDiskSize(ctx context.Context) error {
	if c.ImageName == "" {
		return nil
	}

	// Parse Image String
	// Supported formats:
	// - projects/{project}/global/images/{image}
	// - projects/{project}/global/images/family/{family}
	// - {project}/{image} (or family)
	// - {image} (or family) - assumes current project

	var project, resource, kind string

	if strings.HasPrefix(c.ImageName, "projects/") {
		parts := strings.Split(c.ImageName, "/")
		if len(parts) >= 5 {
			project = parts[1]
			// parts[2] = global, parts[3] = images
			if len(parts) == 6 && parts[4] == "family" {
				kind = "family"
				resource = parts[5]
			} else {
				kind = "image" // or mixed, but usually image
				resource = parts[len(parts)-1]
			}
		}
	} else if strings.Contains(c.ImageName, "/") {
		parts := strings.Split(c.ImageName, "/")
		if len(parts) == 2 {
			project = parts[0]
			resource = parts[1]
			// We don't know if it's family or image yet, try both
			kind = "unknown"
		}
	} else {
		project = c.GCPProjectID
		resource = c.ImageName
		kind = "unknown"
	}

	if project == "" || resource == "" {
		// Fallback or skip if we can't parse, might be a complex self-link we don't handle
		lg(fmt.Sprintf("Warning: Could not parse project/resource from image name '%s'. Skipping disk size adjustment.", c.ImageName))
		return nil
	}

	// Setup Compute Service
	opts := []option.ClientOption{}
	if len(c.GCPCredentials) > 0 {
		opts = append(opts, option.WithCredentialsJSON(c.GCPCredentials))
	}
	computeService, err := compute.NewService(ctx, opts...)
	if err != nil {
		return fmt.Errorf("failed to create compute service: %w", err)
	}

	var imageSizeGb int64

	getImageSize := func(p, r string) (int64, error) {
		img, err := computeService.Images.Get(p, r).Do()
		if err == nil {
			return img.DiskSizeGb, nil
		}
		return 0, err
	}

	getFamilySize := func(p, r string) (int64, error) {
		img, err := computeService.Images.GetFromFamily(p, r).Do()
		if err == nil {
			return img.DiskSizeGb, nil
		}
		return 0, err
	}

	lg(fmt.Sprintf("Checking size for image: Project=%s, Resource=%s...", project, resource))

	// Attempt resolution
	if kind == "family" {
		imageSizeGb, err = getFamilySize(project, resource)
	} else if kind == "image" {
		imageSizeGb, err = getImageSize(project, resource)
	} else {
		// Try image first
		imageSizeGb, err = getImageSize(project, resource)
		if err != nil {
			// Check if 404
			if isNotFound(err) {
				// Try family
				imageSizeGb, err = getFamilySize(project, resource)
			}
		}
	}

	if err != nil {
		// Log warning but don't fail, maybe the user has access via Terraform but not this account, or other issues.
		lg(fmt.Sprintf("Warning: Failed to resolve image size for %s/%s: %v. Using configured default.", project, resource, err))
		return nil
	}

	// Parse current requested size
	var currentSize int64 = 0
	if c.BootDiskSizeGB != "" {
		currentSize, _ = strconv.ParseInt(c.BootDiskSizeGB, 10, 64)
	}

	if imageSizeGb > currentSize {
		lg(fmt.Sprintf("Adjusting BOOT_DISK_SIZE_GB from %d to %d (required by image)", currentSize, imageSizeGb))
		c.BootDiskSizeGB = fmt.Sprintf("%d", imageSizeGb)
	} else {
		lg(fmt.Sprintf("Image size (%dGB) is within requested size (%dGB).", imageSizeGb, currentSize))
	}

	return nil
}

func isNotFound(err error) bool {
	if gErr, ok := err.(*googleapi.Error); ok {
		return gErr.Code == 404
	}
	return false
}

// Copies a single file from src to dst.
func copyFile(src, dst string) error {
	sourceFileStat, err := os.Stat(src)
	if err != nil {
		return err
	}

	if !sourceFileStat.Mode().IsRegular() {
		return fmt.Errorf("%s is not a regular file", src)
	}

	source, err := os.Open(src)
	if err != nil {
		return err
	}
	defer source.Close()

	destination, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer destination.Close()

	_, err = io.Copy(destination, source)
	return err
}

func lt(s string) {
	fmt.Printf("lemc.html.trunc; <pre>%s</pre><br>\n", s)
}

func lg(s string) {
	fmt.Printf("lemc.html.append; <pre>%s</pre><br>\n", s)
}

// TerraformInit initializes a Terraform working directory
func TerraformInit(workingDir string) error {
	lg("Running Terraform init...")
	initCmd := exec.Command("terraform", "init")
	initCmd.Stdout = os.Stdout
	initCmd.Stderr = os.Stderr
	initCmd.Dir = workingDir
	return initCmd.Run()
}

// TerraformPlan creates a Terraform execution plan
func TerraformPlan(workingDir string) error {
	lg("Running Terraform plan...")
	planCmd := exec.Command("terraform", "plan", "-out", "terraform.plan")
	planCmd.Stdout = os.Stdout
	planCmd.Stderr = os.Stderr
	planCmd.Dir = workingDir
	return planCmd.Run()
}

// runTerraformCommandWithProgress executes a Terraform command and handles its progress output
func runTerraformCommandWithProgress(cmd *exec.Cmd, operation string) error {
	output, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("error creating stdout pipe: %v", err)
	}
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("error starting %s command: %v", operation, err)
	}

	progressHandler := progress.NewProgressHandler(output)

	for {
		line, err := progressHandler.ReadLine()
		if err != nil {
			if err == io.EOF {
				break
			}
			lg(fmt.Sprintf("Error: Terraform %s failed. %v", operation, err))
			for _, line := range progressHandler.GetOriginalOutput() {
				lg(line)
			}
			return fmt.Errorf("error reading %s output: %v", operation, err)
		}

		if line != "" {
			formattedLine := strings.TrimSpace(line)
			lt(formattedLine)
		}
	}

	if err := cmd.Wait(); err != nil {
		lg(fmt.Sprintf("Error: Terraform %s failed. %v", operation, err))
		for _, line := range progressHandler.GetOriginalOutput() {
			lg(line)
		}
		return fmt.Errorf("error waiting for %s command: %v", operation, err)
	}

	lg(fmt.Sprintf("Terraform %s completed.", operation))
	lg(fmt.Sprintf("Complete Terraform %s JSON output:", operation))

	for _, line := range progressHandler.GetOriginalOutput() {
		log.Println(line)
	}

	return nil
}

// TerraformApply applies the Terraform plan
func TerraformApply(workingDir string) error {
	lg("Running Terraform apply...")
	applyCmd := exec.Command("terraform", "apply", "-auto-approve", "-json", "terraform.plan")
	applyCmd.Dir = workingDir
	return runTerraformCommandWithProgress(applyCmd, "apply")
}

// TerraformDestroy destroys the Terraform-managed infrastructure
func TerraformDestroy(workingDir string) error {
	lg("Running Terraform destroy...")
	destroyCmd := exec.Command("terraform", "destroy", "-auto-approve", "-json")
	destroyCmd.Dir = workingDir
	return runTerraformCommandWithProgress(destroyCmd, "destroy")
}

// GetTerraformOutputs retrieves the Terraform outputs as a map
func GetTerraformOutputs(workingDir string) (map[string]struct {
	Value     interface{} `json:"value"`
	Type      string      `json:"type"`
	Sensitive bool        `json:"sensitive"`
}, error) {
	lg("Retrieving Terraform outputs...")
	outputCmd := exec.Command("terraform", "output", "-json")
	outputCmd.Dir = workingDir
	outputJSON, err := outputCmd.Output()
	if err != nil {
		return nil, fmt.Errorf("error running terraform output -json: %v", err)
	}

	var outputs map[string]struct {
		Value     interface{} `json:"value"`
		Type      string      `json:"type"`
		Sensitive bool        `json:"sensitive"`
	}
	if err := json.Unmarshal(outputJSON, &outputs); err != nil {
		return nil, fmt.Errorf("error parsing terraform output JSON: %v", err)
	}

	return outputs, nil
}

// GenerateTerraformVars generates the content for terraform.tfvars
func (c *Config) GenerateTerraformVars() string {
	var content strings.Builder

	// Map of environment variables to Terraform variables
	varMap := map[string]string{
		"GCP_PROJECT_ID": "gcp_project_id",
		"GCP_REGION":     "gcp_region",
		"GCP_ZONE":       "gcp_zone",
		"LEMC_SCOPE":     "lemc_scope",
		"LEMC_USERNAME":  "lemc_username",
		"LEMC_USER_ID":   "lemc_user_id",
		"PORT_FROM":      "port",
	}

	// Add mapped variables
	for envKey, tfKey := range varMap {
		var value string
		// Prioritize project_id from key file
		if envKey == "GCP_PROJECT_ID" && c.GCPProjectID != "" {
			value = c.GCPProjectID
			lg(fmt.Sprintf("Using GCP_PROJECT_ID from key file for tfvar %s: %s", tfKey, value))
		} else {
			value = os.Getenv(envKey)
		}

		if value != "" {
			content.WriteString(fmt.Sprintf("%s = \"%s\"\n", tfKey, value))
		} else {
			// Log warning only if not GCP_PROJECT_ID (which might be intentionally overridden by key)
			// or if it's GCP_PROJECT_ID but was not found in key either.
			if envKey != "GCP_PROJECT_ID" || (envKey == "GCP_PROJECT_ID" && c.GCPProjectID == "") {
				lg(fmt.Sprintf("Warning: Environment variable %s not found, and not overridden by key file.", envKey))
			}
		}
	}

	// Add lemc_uuid
	content.WriteString(fmt.Sprintf("lemc_uuid = \"%s\"\n", c.LEMCUUID))

	// Add resource prefix
	content.WriteString(fmt.Sprintf("resource_prefix = \"%s\"\n", c.ResourcePrefix))

	// Add machine type if specified
	if c.LEMCMachineType != "" {
		content.WriteString(fmt.Sprintf("machine_type = \"%s\"\n", c.LEMCMachineType))
	} else {
		lg("Warning: MACHINE_TYPE not found in environment. Terraform will use its default.")
	}

	// Add boot disk size if specified
	if c.BootDiskSizeGB != "" {
		content.WriteString(fmt.Sprintf("boot_disk_size_gb = %s\n", c.BootDiskSizeGB))
	} else {
		lg("Warning: BOOT_DISK_SIZE_GB not found in environment. Terraform will use its default.")
	}

	// Add image if specified
	if c.ImageName != "" {
		content.WriteString(fmt.Sprintf("image = \"%s\"\n", c.ImageName))
	} else {
		lg("Warning: AMI not found in environment. Terraform will use its default.")
	}

	// Add domain name
	content.WriteString(fmt.Sprintf("domain_name = \"%s\"\n", c.DomainName))

	// Add DNS zone name
	if c.RootZone != "" {
		content.WriteString(fmt.Sprintf("dns_zone_name = \"%s\"\n", c.RootZone))
	} else {
		lg("Warning: ROOT_ZONE not found in environment, required for dns_zone_name.")
		content.WriteString(fmt.Sprintf("dns_zone_name = \"%s\"\n", "default-zone-name")) // Placeholder
	}

	// Add backend timeout
	content.WriteString(fmt.Sprintf("backend_timeout_seconds = %s\n", c.BackendTimeoutSeconds))

	// Add existing SSL certificate name if provided
	if c.ExistingSSLCertificateName != "" {
		content.WriteString(fmt.Sprintf("existing_ssl_certificate_name = \"%s\"\n", c.ExistingSSLCertificateName))
	} else {
		// If not provided, Terraform will expect it to be an empty string or not set, which is fine as it's optional there.
		// We could also choose to not write the variable at all if empty.
		lg("Warning: EXISTING_SSL_CERTIFICATE_NAME not provided. If you intend to use an existing certificate, ensure this variable is set in Terraform or the environment.")
	}

	return content.String()
}

// generateGCPAccessToken generates a short-lived GCP access token.
func generateGCPAccessToken(ctx context.Context, serviceAccountKeyJSON string) (string, error) {
	var keyData serviceAccountKey
	if err := json.Unmarshal([]byte(serviceAccountKeyJSON), &keyData); err != nil {
		return "", fmt.Errorf("failed to parse service account key JSON: %w", err)
	}
	if keyData.ClientEmail == "" {
		return "", fmt.Errorf("client_email not found in service account key JSON")
	}

	iamService, err := iamcredentials.NewService(ctx, option.WithCredentialsJSON([]byte(serviceAccountKeyJSON)))
	if err != nil {
		return "", fmt.Errorf("failed to create IAM credentials service: %w", err)
	}

	resourceName := "projects/-/serviceAccounts/" + keyData.ClientEmail
	request := &iamcredentials.GenerateAccessTokenRequest{
		Lifetime: "3600s", // 1 hour (Default max allowed without Org Policy change)
		Scope: []string{
			"https://www.googleapis.com/auth/cloud-platform",
		},
	}

	resp, err := iamService.Projects.ServiceAccounts.GenerateAccessToken(resourceName, request).Do()
	if err != nil {
		return "", fmt.Errorf("failed to generate access token: %w", err)
	}

	return resp.AccessToken, nil
}

// executeCommandsInSSHSession connects to a remote VM via SSH and executes a series of commands.
func executeCommandsInSSHSession(sshUsername, vmIP, sshKeyPath string, commands []string) error {
	if len(commands) == 0 {
		lg("No commands provided to executeCommandsInSSHSession.")
		return nil
	}

	fullCommand := strings.Join(commands, " && ")
	lg(fmt.Sprintf("Preparing to execute remote commands on %s@%s: %s", sshUsername, vmIP, fullCommand))

	cmdArgs := []string{
		"-i", sshKeyPath,
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		fmt.Sprintf("%s@%s", sshUsername, vmIP),
		fullCommand,
	}

	cmd := exec.Command("ssh", cmdArgs...)

	// Create a pipe for stdin and close it immediately to prevent hanging
	stdinPipe, err := cmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("failed to create stdin pipe: %w", err)
	}
	stdinPipe.Close()

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("failed to create stdout pipe: %w", err)
	}
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("failed to create stderr pipe: %w", err)
	}

	lg(fmt.Sprintf("Executing SSH command: ssh %s", strings.Join(cmdArgs, " ")))

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start SSH command: %w", err)
	}

	var wg sync.WaitGroup
	wg.Add(2)

	// Stream stdout
	go func() {
		defer wg.Done()
		scanner := bufio.NewScanner(stdoutPipe)
		for scanner.Scan() {
			lg(fmt.Sprintf("[remote] %s", scanner.Text()))
		}
	}()

	// Stream stderr
	go func() {
		defer wg.Done()
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			lg(fmt.Sprintf("[remote-err] %s", scanner.Text()))
		}
	}()

	if err := cmd.Wait(); err != nil {
		wg.Wait() // Ensure all output is processed
		return fmt.Errorf("failed to execute remote commands on %s@%s: %w", sshUsername, vmIP, err)
	}
	wg.Wait()

	lg(fmt.Sprintf("Successfully executed remote commands on %s@%s.", sshUsername, vmIP))
	return nil
}

// runImage creates a GCP disk image from the running VM using the Compute API
func runImage(config *Config) error {
	lg("Starting image creation process...")

	if config.TargetImageName == "" {
		return fmt.Errorf("IMAGE_NAME environment variable is required for image creation")
	}

	ctx := context.Background()

	// Setup Compute Service
	opts := []option.ClientOption{}
	if len(config.GCPCredentials) > 0 {
		opts = append(opts, option.WithCredentialsJSON(config.GCPCredentials))
	}
	computeService, err := compute.NewService(ctx, opts...)
	if err != nil {
		return fmt.Errorf("failed to create compute service: %w", err)
	}

	// Construct Source Disk URL
	// Format: projects/{project}/zones/{zone}/disks/{disk}
	sourceDiskURL := fmt.Sprintf("projects/%s/zones/%s/disks/%s", config.GCPProjectID, config.GCPZone, config.ResourcePrefix)

	lg(fmt.Sprintf("Preparing to create image '%s' from source disk '%s'...", config.TargetImageName, config.ResourcePrefix))

	image := &compute.Image{
		Name:       config.TargetImageName,
		SourceDisk: sourceDiskURL,
	}

	// Insert the image. ForceCreate=true is required if the instance is running.
	op, err := computeService.Images.Insert(config.GCPProjectID, image).ForceCreate(true).Do()
	if err != nil {
		return fmt.Errorf("failed to start image creation: %w", err)
	}

	lg(fmt.Sprintf("Image creation started. Operation ID: %s", op.Name))

	// Poll for completion
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	// Global Operations service for Images (Images are global resources, but sometimes operations are zonal?
	// Images.Insert returns a GlobalOperation)
	globalOperationsService := computeService.GlobalOperations

	for {
		select {
		case <-ticker.C:
			lg("Polling image creation status...")
			pollOp, err := globalOperationsService.Get(config.GCPProjectID, op.Name).Do()
			if err != nil {
				return fmt.Errorf("failed to poll operation status: %w", err)
			}

			if pollOp.Status == "DONE" {
				if pollOp.Error != nil {
					var errorMsgs []string
					for _, e := range pollOp.Error.Errors {
						errorMsgs = append(errorMsgs, fmt.Sprintf("%s: %s", e.Code, e.Message))
					}
					return fmt.Errorf("image creation failed with errors: %s", strings.Join(errorMsgs, "; "))
				}

				lg(fmt.Sprintf("Successfully created image: %s", config.TargetImageName))
				return nil
			}
			lg(fmt.Sprintf("Current status: %s. Waiting...", pollOp.Status))

		case <-time.After(30 * time.Minute): // Safety timeout
			return errors.New("image creation timed out after 30 minutes")
		}
	}
}

// runApply executes the Terraform apply operation and handles all related tasks
func runApply(config *Config) error {
	// Initialize Terraform
	if err := TerraformInit(config.AbsoluteDestDir); err != nil {
		return fmt.Errorf("error initializing Terraform: %v", err)
	}

	// Create and apply Terraform plan
	if err := TerraformPlan(config.AbsoluteDestDir); err != nil {
		return fmt.Errorf("error creating Terraform plan: %v", err)
	}

	if err := TerraformApply(config.AbsoluteDestDir); err != nil {
		return fmt.Errorf("error applying Terraform plan: %v", err)
	}

	// Get Terraform outputs
	outputs, err := GetTerraformOutputs(config.AbsoluteDestDir)
	if err != nil {
		return fmt.Errorf("error getting Terraform outputs: %v", err)
	}

	privateKeyOutput, ok := outputs["private_ssh_key"]
	if !ok || privateKeyOutput.Value == nil {
		return fmt.Errorf("private_ssh_key not found in Terraform outputs")
	}

	privateKey, ok := privateKeyOutput.Value.(string)
	if !ok {
		return fmt.Errorf("private_ssh_key in Terraform outputs is not a string")
	}
	sshKeyFilename := "private-ssh-key"
	sshKeyPath := filepath.Join(config.PublicDir, sshKeyFilename)

	lg(fmt.Sprintf("Saving private SSH key to %s...", sshKeyPath))

	if err := os.WriteFile(sshKeyPath, []byte(privateKey), 0600); err != nil {
		return fmt.Errorf("error writing private SSH key to file: %v", err)
	}

	lg(fmt.Sprintf("Private SSH key saved to %s", sshKeyPath))

	downloadURL := config.LEMCHTTPBaseURL + sshKeyFilename
	lg(fmt.Sprintf(`SSH Private Key available for download: <a href="%s" target="_blank" style="color: blue; text-decoration: none;">%s</a>`, downloadURL, sshKeyFilename))
	lg(fmt.Sprintf("IMPORT_FROMANT: Secure this key. It provides root access to the created VM. Download it and then ensure it is removed from the public folder if this is a shared environment or if the link is accessible by others."))

	// Print SSH command
	publicIPOutput, ipOk := outputs["public_ip"]
	var publicIP string
	if ipOk && publicIPOutput.Value != nil && config.LEMCUsername != "" {
		var castOk bool
		publicIP, castOk = publicIPOutput.Value.(string)
		if castOk && publicIP != "" {
			sshCommand := fmt.Sprintf("ssh -i %s %s@%s", sshKeyFilename, config.LEMCUsername, publicIP)
			lg(fmt.Sprintf("To SSH into the VM (once key is downloaded and in current directory): <code>%s</code>", sshCommand))
		} else {
			lg("Warning: Could not retrieve or cast public_ip from Terraform outputs to string.")
		}
	} else {
		if !ipOk || publicIPOutput.Value == nil {
			lg("Warning: public_ip not found in Terraform outputs, cannot print SSH command.")
		}
		if config.LEMCUsername == "" {
			lg("Warning: LEMC_USERNAME not set, cannot print SSH command.")
		}
	}

	// Create .env file
	lg("Creating .env file...")

	s := `SSH_USERNAME=%s
VM_IP=%s
DOMAIN_NAME=https://%s
WORKLOAD_TOKEN=%s
`

	envContent := fmt.Sprintf(s, config.LEMCUsername, publicIP, config.DomainName, config.WorkloadToken)
	envFilePath := filepath.Join(config.PublicDir, "dotenv")

	if err := os.WriteFile(envFilePath, []byte(envContent), 0644); err != nil {
		lg(fmt.Sprintf("Error writing dotenv file: %v", err))
	} else {
		lg(fmt.Sprintf(".env file saved to %s", envFilePath))
		envFileDownloadURL := config.LEMCHTTPBaseURL + "dotenv"
		lg(fmt.Sprintf(`.env file available for download: <a href="%s" target="_blank" style="color: blue; text-decoration: none;">dotenv</a>`, envFileDownloadURL))
	}

	// --- Start Docker Install and 12-Hour Token Logic ---
	if publicIP != "" {
		lg("Proceeding with Docker installation and login...")

		// Transfer .env file to remote VM
		remoteEnvTmpPath := "/tmp/.env"
		lg(fmt.Sprintf("Transferring .env to %s@%s:%s...", config.LEMCUsername, publicIP, remoteEnvTmpPath))

		scpEnvCmd := exec.Command("scp",
			"-i", sshKeyPath,
			"-o", "StrictHostKeyChecking=no",
			"-o", "UserKnownHostsFile=/dev/null",
			envFilePath,
			fmt.Sprintf("%s@%s:%s", config.LEMCUsername, publicIP, remoteEnvTmpPath),
		)
		scpEnvCmd.Stdout = os.Stdout
		scpEnvCmd.Stderr = os.Stderr
		if err := scpEnvCmd.Run(); err != nil {
			lg(fmt.Sprintf("Warning: failed to scp .env file: %v", err))
		} else {
			lg(".env file transferred successfully.")

			// Move .env to user home and root home
			setupEnvCmd := fmt.Sprintf(
				"sudo cp %s /home/%s/.env && sudo chown %s:%s /home/%s/.env && sudo chmod 644 /home/%s/.env && "+
					"sudo cp %s /root/.env && sudo chown root:root /root/.env && sudo chmod 644 /root/.env && "+
					"rm %s",
				remoteEnvTmpPath, config.LEMCUsername, config.LEMCUsername, config.LEMCUsername, config.LEMCUsername, config.LEMCUsername,
				remoteEnvTmpPath,
				remoteEnvTmpPath,
			)

			if err := executeCommandsInSSHSession(config.LEMCUsername, publicIP, sshKeyPath, []string{setupEnvCmd}); err != nil {
				lg(fmt.Sprintf("Warning: failed to setup .env file on remote VM: %v", err))
			} else {
				lg(".env file configured in user and root homes.")
			}
		}

		// 1. Generate 12h Token (Actually 1h default, but part of 12h recipe)
		lg("Generating 1-hour GCP Access Token...")
		ctx := context.Background()
		accessToken, err := generateGCPAccessToken(ctx, string(config.GCPCredentials))
		if err != nil {
			lg(fmt.Sprintf("Error generating GCP access token: %v", err))
			// Fail because this is a key part of this recipe
			return fmt.Errorf("failed to generate access token: %v", err)
		}
		lg("GCP access token generated successfully.")

		// 2. Transfer install.sh
		remoteInstallPath := fmt.Sprintf("/tmp/%s", installScriptName)
		localInstallPath := filepath.Join(config.AbsoluteDestDir, installScriptName) // main() copies it here

		lg(fmt.Sprintf("Transferring %s to %s@%s:%s...", localInstallPath, config.LEMCUsername, publicIP, remoteInstallPath))

		// Use scp instead of rsync because rsync might not be installed in the runner image
		scpCmd := exec.Command("scp",
			"-i", sshKeyPath,
			"-o", "StrictHostKeyChecking=no",
			"-o", "UserKnownHostsFile=/dev/null",
			localInstallPath,
			fmt.Sprintf("%s@%s:%s", config.LEMCUsername, publicIP, remoteInstallPath),
		)
		scpCmd.Stdout = os.Stdout
		scpCmd.Stderr = os.Stderr
		if err := scpCmd.Run(); err != nil {
			return fmt.Errorf("failed to scp install script: %v", err)
		}
		lg("Install script transferred successfully.")

		// 3. Execute install.sh
		lg("Executing install script on remote VM...")
		installCmd := fmt.Sprintf("chmod +x %s && sudo -n %s '%s' '%s' '%s'",
			remoteInstallPath,
			remoteInstallPath,
			accessToken,
			config.GCRHostname,
			config.DockerImageToPull,
		)

		if err := executeCommandsInSSHSession(config.LEMCUsername, publicIP, sshKeyPath, []string{installCmd}); err != nil {
			return fmt.Errorf("failed to execute install script on remote VM: %v", err)
		}
		lg("Docker installation and login completed successfully.")
	}

	return nil
}

// runDestroy executes the Terraform destroy operation
func runDestroy(config *Config) error {
	// Initialize Terraform
	if err := TerraformInit(config.AbsoluteDestDir); err != nil {
		return fmt.Errorf("error initializing Terraform: %v", err)
	}

	if err := TerraformDestroy(config.AbsoluteDestDir); err != nil {
		return fmt.Errorf("error destroying Terraform resources: %v", err)
	}
	lg("Terraform destroy completed successfully.")
	return nil
}

// RunTerraformFunction executes the specified Terraform function
func RunTerraformFunction(config *Config, function TerraformFunction) error {
	switch function {
	case FunctionApply:
		return runApply(config)
	case FunctionDestroy:
		return runDestroy(config)
	case FunctionImage:
		return runImage(config)
	default:
		return fmt.Errorf("unsupported Terraform function: %s", function)
	}
}

func main() {
	lt("Starting Terraform setup...")

	// Get configuration
	config, err := GetConfig()
	if err != nil {
		log.Fatalf("Error getting configuration: %v", err)
	}

	lg(fmt.Sprintf("Running in directory: %s", config.WorkingDir))
	lg("Starting Terraform setup...")

	// Adjust boot disk size if image requires more space
	if err := config.AdjustBootDiskSize(context.Background()); err != nil {
		lg(fmt.Sprintf("Warning: Failed to adjust boot disk size from image: %v", err))
	}

	// Copy files from source directory
	lg(fmt.Sprintf("Copying files from %s to %s...", config.AbsoluteSourceDir, config.AbsoluteDestDir))
	filesToCopy := []string{"main.tf", "outputs.tf", "variables.tf"}
	for _, filename := range filesToCopy {
		srcPath := filepath.Join(config.AbsoluteSourceDir, filename)
		dstPath := filepath.Join(config.AbsoluteDestDir, filename)
		lg(fmt.Sprintf("  Copying %s to %s", srcPath, dstPath))
		if err := copyFile(srcPath, dstPath); err != nil {
			log.Fatalf("Error copying file %s: %v", filename, err)
		}
	}

	// Copy install.sh from files/ to destDir (so it's available for rsync in runApply)
	installSrcPath := filepath.Join(config.WorkingDir, "files", installScriptName)
	installDstPath := filepath.Join(config.AbsoluteDestDir, installScriptName)
	lg(fmt.Sprintf("  Copying %s to %s", installSrcPath, installDstPath))
	if err := copyFile(installSrcPath, installDstPath); err != nil {
		log.Fatalf("Error copying %s: %v", installScriptName, err)
	}

	lg("Finished copying base Terraform files.")

	// Generate terraform.tfvars content
	lg("Generating terraform.tfvars content...")
	tfVarsContent := config.GenerateTerraformVars()

	// Write terraform.tfvars file
	tfVarsPath := filepath.Join(config.AbsoluteDestDir, tfVarsFilename)
	lg(fmt.Sprintf("Writing %s...", tfVarsPath))
	err = os.WriteFile(tfVarsPath, []byte(tfVarsContent), 0644)
	if err != nil {
		log.Fatalf("Error writing %s: %v", tfVarsPath, err)
	}

	lg(fmt.Sprintf("Successfully created %s", tfVarsPath))

	// Get Terraform function from environment
	tfFunction := TerraformFunction(os.Getenv("TERRAFORM_FUNCTION"))
	if tfFunction == "" {
		tfFunction = FunctionApply // Default to apply if not specified
	}

	// Run the specified Terraform function
	if err := RunTerraformFunction(config, tfFunction); err != nil {
		log.Fatalf("Error running Terraform function: %v", err)
	}
}