package dynamo

import (
	"fmt"
	"sort"
	"strconv"
	"strings"

	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

type TRTLLMBackend struct {
	MpiRunSecretName string
}

func (b *TRTLLMBackend) UpdateContainer(container *corev1.Container, numberOfNodes int32, role Role, component *v1alpha1.DynamoComponentDeploymentSharedSpec, serviceName string, multinodeDeployer MultinodeDeployer) {
	// Check for volumeMounts with useAsCompilationCache=true
	for _, volumeMount := range component.VolumeMounts {
		if volumeMount.UseAsCompilationCache {
			logger := log.Log.WithName("trtllm-backend")
			logger.Info("Compilation cache configured for TensorRT-LLM but not yet fully supported",
				"backend", "trtllm",
				"status", "partial-support",
				"use-as-compilation-cache", true,
				"env-vars-set", false,
				"next-steps", "upstream TensorRT-LLM changes needed")
		}
	}

	// For single node, nothing to do
	if numberOfNodes <= 1 {
		return
	}

	// Configure probes for multinode deployments
	if role == RoleWorker {
		// For workers: remove liveness and startup probes, set readiness to check SSH port
		container.LivenessProbe = nil
		container.StartupProbe = nil
		container.ReadinessProbe = &corev1.Probe{
			ProbeHandler: corev1.ProbeHandler{
				TCPSocket: &corev1.TCPSocketAction{
					Port: intstr.FromInt(commonconsts.MpiRunSshPort),
				},
			},
			InitialDelaySeconds: 20,
			PeriodSeconds:       20,
			TimeoutSeconds:      5,
			FailureThreshold:    10,
		}
	}
	// For leaders: leave all probes untouched

	// Add SSH keypair volume mount for multinode deployments
	b.addSSHVolumeMount(container)

	// Add OpenMPI environment variable to keep FQDN hostnames
	envVar := corev1.EnvVar{
		Name:  "OMPI_MCA_orte_keep_fqdn_hostnames",
		Value: "1",
	}
	container.Env = append(container.Env, envVar)

	// Update container command based on role
	switch role {
	case RoleLeader:
		b.setupLeaderContainer(container, numberOfNodes, serviceName, component, multinodeDeployer)
	case RoleWorker:
		b.setupWorkerContainer(container)
	}
}

func (b *TRTLLMBackend) UpdatePodSpec(podSpec *corev1.PodSpec, numberOfNodes int32, role Role, component *v1alpha1.DynamoComponentDeploymentSharedSpec, serviceName string, multinodeDeployer MultinodeDeployer) {
	// Add SSH keypair volume for TRTLLM multinode deployments
	if numberOfNodes > 1 {
		sshVolume := corev1.Volume{
			Name: b.MpiRunSecretName,
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName:  b.MpiRunSecretName,
					DefaultMode: func() *int32 { mode := int32(0644); return &mode }(),
				},
			},
		}
		podSpec.Volumes = append(podSpec.Volumes, sshVolume)
	}
}

// addSSHVolumeMount adds the SSH keypair secret volume mount to the container
func (b *TRTLLMBackend) addSSHVolumeMount(container *corev1.Container) {
	sshVolumeMount := corev1.VolumeMount{
		Name:      b.MpiRunSecretName,
		MountPath: "/ssh-pk",
		ReadOnly:  true,
	}
	container.VolumeMounts = append(container.VolumeMounts, sshVolumeMount)
}

// setupLeaderContainer configures the leader node with SSH setup and mpirun command
func (b *TRTLLMBackend) setupLeaderContainer(container *corev1.Container, numberOfNodes int32, serviceName string, component *v1alpha1.DynamoComponentDeploymentSharedSpec, multinodeDeployer MultinodeDeployer) {
	// Generate the list of all hostnames
	hostNamesList := b.hostNamesList(numberOfNodes, serviceName, multinodeDeployer)
	allHostnames := strings.Join(hostNamesList, ",")

	// Store original command/args for later use
	var originalCommand string

	if len(container.Command) > 0 && isPythonCommand(container.Command[0]) {
		// Direct Python command: combine command + args
		// Shell-quote each part to handle args with spaces (e.g., JSON in --override-engine-args)
		var quotedParts []string
		for _, part := range container.Command {
			quotedParts = append(quotedParts, shellQuoteForBashC(part))
		}
		for _, part := range container.Args {
			quotedParts = append(quotedParts, shellQuoteForBashC(part))
		}
		originalCommand = strings.Join(quotedParts, " ")
	} else if len(container.Args) > 0 {
		// Shell command (sh -c): args contains the full command
		originalCommand = strings.Join(container.Args, " ")
	} else if len(container.Command) > 0 {
		// Fallback: just command
		originalCommand = strings.Join(container.Command, " ")
	}

	// Setup SSH and run mpirun command
	// Use $HOME instead of ~ because the container may set HOME=/home/dynamo via Dockerfile
	// while the shell user is root (from securityContext.runAsUser: 0). ~ follows /etc/passwd
	// but $HOME follows the environment, and SSH/mpirun need to find keys where we put them.
	sshSetupCommands := []string{
		"mkdir -p $HOME/.ssh",
		"ls -la /ssh-pk/", // Debug: list files in ssh-pk directory
		"cp /ssh-pk/private.key $HOME/.ssh/id_rsa",
		"cp /ssh-pk/private.key.pub $HOME/.ssh/id_rsa.pub",
		"cp /ssh-pk/private.key.pub $HOME/.ssh/authorized_keys",
		"chmod 600 $HOME/.ssh/id_rsa $HOME/.ssh/authorized_keys",
		"chmod 644 $HOME/.ssh/id_rsa.pub",
		fmt.Sprintf("printf 'Host *\\nIdentityFile '$HOME'/.ssh/id_rsa\\nStrictHostKeyChecking no\\nPort %d\\n' > $HOME/.ssh/config", commonconsts.MpiRunSshPort),
	}

	// Calculate total number of GPUs across all nodes
	gpusPerNode := getGPUsPerNode(component.Resources)
	totalGPUs := numberOfNodes * gpusPerNode

	// Build mpirun command with explicit SSH configuration and environment variables
	// Wrap the entire command (trtllm-llmapi-launch + original command) in bash -c for proper shell interpretation
	wrappedCommand := fmt.Sprintf("bash -c 'trtllm-llmapi-launch %s'", originalCommand)

	// Generate environment variable flags for mpirun
	envVarsStr := generateEnvVarFlags(container.Env)

	// Use --allow-run-as-root only when the container is running as root (UID 0).
	// When running as a non-root user, mpirun works without this flag and omitting
	// it avoids masking accidental root execution.
	mpirunCmd := fmt.Sprintf("mpirun $([ \"$(id -u)\" = \"0\" ] && echo --allow-run-as-root) --oversubscribe -n %d -H %s --mca pml ob1 --mca plm_rsh_args \"-p %d -o StrictHostKeyChecking=no -i $HOME/.ssh/id_rsa\" %s %s",
		totalGPUs,
		allHostnames,
		commonconsts.MpiRunSshPort,
		envVarsStr,
		wrappedCommand)

	// Combine SSH setup and mpirun command, optionally adding DNS wait for deployers that need it
	var allCommands []string
	if multinodeDeployer.NeedsDNSWait() {
		// Wait for DNS resolution of all worker nodes (needed for LWS)
		workerHosts := strings.Join(hostNamesList[1:], " ")
		dnsWaitCmd := fmt.Sprintf(`TIMEOUT=300; START_TIME=$(date +%%s); for worker in %s; do echo "Waiting for DNS: $worker"; until getent hosts $worker >/dev/null 2>&1; do CURRENT_TIME=$(date +%%s); if [ $((CURRENT_TIME - START_TIME)) -gt $TIMEOUT ]; then echo "ERROR: Timeout waiting for DNS: $worker"; exit 1; fi; echo "DNS not ready for $worker, retrying..."; sleep 2; done; echo "✓ DNS resolved: $worker"; done; echo "All workers DNS ready"`, workerHosts)

		allCommands = append(sshSetupCommands, dnsWaitCmd, mpirunCmd)
	} else {
		allCommands = append(sshSetupCommands, mpirunCmd)
	}
	fullCommand := strings.Join(allCommands, " && ")

	// Update container to use bash with the full command
	container.Command = []string{"/bin/sh", "-c"}
	container.Args = []string{fullCommand}
}

// setupWorkerContainer configures worker nodes with SSH setup and daemon
func (b *TRTLLMBackend) setupWorkerContainer(container *corev1.Container) {
	// Setup SSH for worker nodes
	// Use $HOME instead of ~ for the same reasons as setupLeaderContainer (see comment above).
	sshSetupCommands := []string{
		"mkdir -p $HOME/.ssh $HOME/.ssh/host_keys $HOME/.ssh/run /run/sshd",
		"ls -la /ssh-pk/", // Debug: list files in ssh-pk directory
		"cp /ssh-pk/private.key $HOME/.ssh/id_rsa",
		"cp /ssh-pk/private.key.pub $HOME/.ssh/id_rsa.pub",
		"cp /ssh-pk/private.key.pub $HOME/.ssh/authorized_keys",
		"chmod 600 $HOME/.ssh/id_rsa $HOME/.ssh/authorized_keys",
		"chmod 644 $HOME/.ssh/id_rsa.pub",
		fmt.Sprintf("printf 'Host *\\nIdentityFile '$HOME'/.ssh/id_rsa\\nStrictHostKeyChecking no\\nPort %d\\n' > $HOME/.ssh/config", commonconsts.MpiRunSshPort),
		// Generate host keys in user writable directory
		"ssh-keygen -t rsa -f $HOME/.ssh/host_keys/ssh_host_rsa_key -N ''",
		"ssh-keygen -t ecdsa -f $HOME/.ssh/host_keys/ssh_host_ecdsa_key -N ''",
		"ssh-keygen -t ed25519 -f $HOME/.ssh/host_keys/ssh_host_ed25519_key -N ''",
		// Create SSH daemon config using $HOME for absolute paths.
		// sshd expands ~ via /etc/passwd (-> /root/) not the HOME env var,
		// so we break out of single quotes to let the shell expand $HOME.
		// AuthorizedKeysFile also needs absolute $HOME path because sshd resolves
		// relative paths from the connecting user's /etc/passwd home (-> /root/).
		// StrictModes disabled because /home/dynamo may be owned by a non-root UID
		// while sshd runs as root, causing permission check failures.
		fmt.Sprintf("printf 'Port %d\\nHostKey '$HOME'/.ssh/host_keys/ssh_host_rsa_key\\nHostKey '$HOME'/.ssh/host_keys/ssh_host_ecdsa_key\\nHostKey '$HOME'/.ssh/host_keys/ssh_host_ed25519_key\\nPidFile '$HOME'/.ssh/run/sshd.pid\\nStrictModes no\\nPermitRootLogin yes\\nPasswordAuthentication no\\nPubkeyAuthentication yes\\nAuthorizedKeysFile '$HOME'/.ssh/authorized_keys\\n' > $HOME/.ssh/sshd_config", commonconsts.MpiRunSshPort),
		"/usr/sbin/sshd -D -f $HOME/.ssh/sshd_config",
	}

	fullCommand := strings.Join(sshSetupCommands, " && ")

	// Update container to use bash with the SSH setup and daemon
	container.Command = []string{"/bin/sh", "-c"}
	container.Args = []string{fullCommand}
}

// hostNamesList generates the list of hostnames for all nodes in the multinode deployment
func (b *TRTLLMBackend) hostNamesList(numberOfNodes int32, serviceName string, multinodeDeployer MultinodeDeployer) []string {
	return multinodeDeployer.GetHostNames(serviceName, numberOfNodes)
}

// getGPUsPerNode extracts the number of GPUs per node from resources
func getGPUsPerNode(resources *v1alpha1.Resources) int32 {
	if resources != nil && resources.Requests != nil && resources.Requests.GPU != "" {
		if gpus, err := strconv.ParseInt(resources.Requests.GPU, 10, 32); err == nil {
			return int32(gpus)
		}
	}
	if resources != nil && resources.Limits != nil && resources.Limits.GPU != "" {
		if gpus, err := strconv.ParseInt(resources.Limits.GPU, 10, 32); err == nil {
			return int32(gpus)
		}
	}
	return 0 // Default to 0 GPUs if not specified
}

// getCommonTRTLLMEnvVars returns a map of common environment variables for TRTLLM deployments
func getCommonTRTLLMEnvVars() map[string]bool {
	return map[string]bool{
		"CUDA_VISIBLE_DEVICES": true, "MODEL_PATH": true, "HF_TOKEN": true, "HUGGING_FACE_HUB_TOKEN": true, "HF_ENDPOINT": true,
		"TOKENIZERS_PARALLELISM": true, "NCCL_DEBUG": true, "NCCL_IB_DISABLE": true, "NCCL_P2P_DISABLE": true,
		"TENSORRT_LLM_CACHE_DIR": true, "HF_HOME": true, "TRANSFORMERS_CACHE": true, "HF_DATASETS_CACHE": true,
		"PATH": true, "LD_LIBRARY_PATH": true, "PYTHONPATH": true, "HOME": true, "USER": true, "TRTLLM_USE_UCX_KVCACHE": true,
	}
}

// collectAllEnvVars combines explicit container env vars with common TRTLLM env vars, removing duplicates
func collectAllEnvVars(containerEnvVars []corev1.EnvVar) []string {
	// Initialize set with common environment variables
	envVarSet := getCommonTRTLLMEnvVars()

	// Add explicit environment variables from container
	for _, env := range containerEnvVars {
		envVarSet[env.Name] = true
	}

	// Convert set to sorted slice for consistent output
	envVarNames := make([]string, 0, len(envVarSet))
	for envVar := range envVarSet {
		envVarNames = append(envVarNames, envVar)
	}
	sort.Strings(envVarNames)

	return envVarNames
}

// formatEnvVarFlags converts environment variable names to mpirun -x flags
func formatEnvVarFlags(envVarNames []string) string {
	envVars := make([]string, 0, len(envVarNames))
	for _, envVar := range envVarNames {
		envVars = append(envVars, fmt.Sprintf("-x %s", envVar))
	}
	return strings.Join(envVars, " ")
}

// generateEnvVarFlags generates the complete environment variable flags string for mpirun
func generateEnvVarFlags(containerEnvVars []corev1.EnvVar) string {
	envVarNames := collectAllEnvVars(containerEnvVars)
	return formatEnvVarFlags(envVarNames)
}
