package cases

import (
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/sandbox0-ai/sandbox0/pkg/apispec"
	"github.com/sandbox0-ai/sandbox0/pkg/framework"
	"github.com/sandbox0-ai/sandbox0/pkg/naming"
	e2eutils "github.com/sandbox0-ai/sandbox0/tests/e2e/utils"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

const (
	sshFixtureNamespace       = "sandbox0-e2e-ssh-fixture"
	sshFixtureDeploymentName  = "ssh-server"
	sshFixtureServiceName     = "ssh-server"
	sshFixtureUserName        = "e2e"
	sshFixtureImageEnvVar     = "E2E_SSH_FIXTURE_IMAGE"
	sshFixtureSourceImageRef  = "lscr.io/linuxserver/openssh-server@sha256:68b605929e83b2efe000da09269688f6d82a44579e8a18e2d9e8c8d272917cf7"
	defaultSSHFixtureImageRef = "sandbox0ai/e2e-openssh-server:68b605929e83"
	sshFixturePullAttempts    = 3
)

const sshFixturePrivateKey = `-----BEGIN OPENSSH PRIVATE KEY-----
b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW
QyNTUxOQAAACBOgie1mnIFBEInHm7z4BVuUiAaFzia5Irj1x5mkBT6ywAAAJA0M/muNDP5
rgAAAAtzc2gtZWQyNTUxOQAAACBOgie1mnIFBEInHm7z4BVuUiAaFzia5Irj1x5mkBT6yw
AAAED3v+5SnUyId7/jDlrMYu5c/Jx18a8SQBSCcxTFtnNHEk6CJ7WacgUEQicebvPgFW5S
IBoXOJrkiuPXHmaQFPrLAAAADHNhbmRib3gwLWUyZQE=
-----END OPENSSH PRIVATE KEY-----`

const sshFixturePublicKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIE6CJ7WacgUEQicebvPgFW5SIBoXOJrkiuPXHmaQFPrL sandbox0-e2e"

const sshProxyFakePrivateKey = `-----BEGIN OPENSSH PRIVATE KEY-----
b3BlbnNzaC1rZXktdjEAAAAABG5vbmUAAAAEbm9uZQAAAAAAAAABAAAAMwAAAAtzc2gtZW
QyNTUxOQAAACASvDR8ptotozntpt2p9CwRrhplUWStWNc6pBZQmd76JQAAAJji6KuR4uir
kQAAAAtzc2gtZWQyNTUxOQAAACASvDR8ptotozntpt2p9CwRrhplUWStWNc6pBZQmd76JQ
AAAEB5X9BIl2QVR3rMJVVSat2i+fYMRxzAiyUnBFLH0WOGlhK8NHym2i2jOe2m3an0LBGu
GmVRZK1Y1zqkFlCZ3volAAAAEXNhbmRib3gwLWUyZS1mYWtlAQIDBA==
-----END OPENSSH PRIVATE KEY-----`

const sshProxyFakePublicKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIBK8NHym2i2jOe2m3an0LBGuGmVRZK1Y1zqkFlCZ3vol sandbox0-e2e-fake"

type sshFixture struct {
	Namespace string
	ServiceIP string
	UserName  string
	ImageRef  string
}

func setupSSHFixture(env *framework.ScenarioEnv) (*sshFixture, func()) {
	imageRef := strings.TrimSpace(os.Getenv(sshFixtureImageEnvVar))
	if imageRef == "" {
		imageRef = defaultSSHFixtureImageRef
	}

	if env.TestCtx != nil && env.TestCtx.Cluster != nil {
		Expect(preloadSSHFixtureImage(env, imageRef)).To(Succeed())
	}

	Expect(framework.EnsureNamespace(env.TestCtx.Context, env.Config.Kubeconfig, sshFixtureNamespace)).To(Succeed())
	Expect(framework.ApplyManifestContent(env.TestCtx.Context, env.Config.Kubeconfig, "sandbox0-e2e-ssh-fixture-", buildSSHFixtureManifest(imageRef))).To(Succeed())
	Expect(framework.WaitForDeployment(env.TestCtx.Context, env.Config.Kubeconfig, sshFixtureNamespace, sshFixtureDeploymentName, "3m")).To(Succeed())

	var serviceIP string
	Eventually(func() string {
		ip, err := framework.KubectlGetJSONPath(
			env.TestCtx.Context,
			env.Config.Kubeconfig,
			sshFixtureNamespace,
			"service",
			sshFixtureServiceName,
			"{.spec.clusterIP}",
		)
		if err != nil {
			return ""
		}
		return strings.TrimSpace(ip)
	}).WithTimeout(30 * time.Second).WithPolling(2 * time.Second).ShouldNot(BeEmpty())

	serviceIP, err := framework.KubectlGetJSONPath(
		env.TestCtx.Context,
		env.Config.Kubeconfig,
		sshFixtureNamespace,
		"service",
		sshFixtureServiceName,
		"{.spec.clusterIP}",
	)
	Expect(err).NotTo(HaveOccurred())

	fixture := &sshFixture{
		Namespace: sshFixtureNamespace,
		ServiceIP: strings.TrimSpace(serviceIP),
		UserName:  sshFixtureUserName,
		ImageRef:  imageRef,
	}

	cleanup := func() {
		_ = framework.Kubectl(
			env.TestCtx.Context,
			env.Config.Kubeconfig,
			"delete",
			"namespace",
			sshFixtureNamespace,
			"--ignore-not-found=true",
			"--wait=false",
		)
	}

	return fixture, cleanup
}

func buildSSHFixtureManifest(imageRef string) string {
	return fmt.Sprintf(`apiVersion: apps/v1
kind: Deployment
metadata:
  name: %s
  namespace: %s
spec:
  replicas: 1
  selector:
    matchLabels:
      app.kubernetes.io/name: %s
  template:
    metadata:
      labels:
        app.kubernetes.io/name: %s
    spec:
      containers:
        - name: sshd
          image: %s
          imagePullPolicy: IfNotPresent
          env:
            - name: PUID
              value: "1000"
            - name: PGID
              value: "1000"
            - name: TZ
              value: "Etc/UTC"
            - name: USER_NAME
              value: %q
            - name: PUBLIC_KEY
              value: %q
            - name: PASSWORD_ACCESS
              value: "false"
            - name: LOG_STDOUT
              value: "true"
          ports:
            - name: ssh
              containerPort: 2222
          readinessProbe:
            tcpSocket:
              port: 2222
            initialDelaySeconds: 2
            periodSeconds: 2
          volumeMounts:
            - name: config
              mountPath: /config
      volumes:
        - name: config
          emptyDir: {}
---
apiVersion: v1
kind: Service
metadata:
  name: %s
  namespace: %s
spec:
  selector:
    app.kubernetes.io/name: %s
  ports:
    - name: ssh
      port: 22
      targetPort: 2222
      protocol: TCP
`, sshFixtureDeploymentName, sshFixtureNamespace, sshFixtureDeploymentName, sshFixtureDeploymentName, imageRef, sshFixtureUserName, sshFixturePublicKey, sshFixtureServiceName, sshFixtureNamespace, sshFixtureDeploymentName)
}

func preloadSSHFixtureImage(env *framework.ScenarioEnv, imageRef string) error {
	loaded, err := sshFixtureImageLoadedInKind(env, imageRef)
	if err != nil {
		fmt.Printf("Unable to inspect Kind image cache for %q: %v\n", imageRef, err)
	} else if loaded {
		fmt.Printf("SSH fixture image %q is already loaded in Kind cluster %q.\n", imageRef, env.TestCtx.Cluster.Name)
		return nil
	}

	if _, err := framework.RunCommandOutput(env.TestCtx.Context, "docker", "image", "inspect", imageRef); err != nil {
		sourceRef := imageRef
		if imageRef == defaultSSHFixtureImageRef {
			sourceRef = sshFixtureSourceImageRef
		}

		if _, err := framework.RunCommandOutput(env.TestCtx.Context, "docker", "image", "inspect", sourceRef); err != nil {
			if err := pullSSHFixtureImageWithRetry(env, sourceRef); err != nil {
				return err
			}
		}
		if sourceRef != imageRef {
			if err := framework.RunCommand(env.TestCtx.Context, "docker", "tag", sourceRef, imageRef); err != nil {
				return err
			}
		}
	}
	return env.TestCtx.Cluster.LoadDockerImage(env.TestCtx.Context, imageRef)
}

func sshFixtureImageLoadedInKind(env *framework.ScenarioEnv, imageRef string) (bool, error) {
	if env == nil || env.TestCtx == nil || env.TestCtx.Cluster == nil || env.TestCtx.Cluster.Name == "" {
		return false, nil
	}

	nodesOutput, err := framework.RunCommandOutput(env.TestCtx.Context, "kind", "get", "nodes", "--name", env.TestCtx.Cluster.Name)
	if err != nil {
		return false, err
	}

	nodes := strings.Fields(nodesOutput)
	if len(nodes) == 0 {
		return false, nil
	}

	candidates := sshFixtureImageNameCandidates(imageRef)
	for _, node := range nodes {
		imagesOutput, err := framework.RunCommandOutput(env.TestCtx.Context, "docker", "exec", node, "ctr", "--namespace=k8s.io", "images", "ls", "-q")
		if err != nil {
			return false, err
		}
		if !sshFixtureImageListContains(imagesOutput, candidates) {
			return false, nil
		}
	}

	return true, nil
}

func sshFixtureImageNameCandidates(imageRef string) []string {
	candidates := []string{imageRef}
	firstPathPart, _, ok := strings.Cut(imageRef, "/")
	if ok && firstPathPart != "localhost" && !strings.ContainsAny(firstPathPart, ".:") {
		candidates = append(candidates, "docker.io/"+imageRef)
	}
	return candidates
}

func sshFixtureImageListContains(imagesOutput string, candidates []string) bool {
	for _, line := range strings.Split(imagesOutput, "\n") {
		imageName := strings.TrimSpace(line)
		if imageName == "" {
			continue
		}
		for _, candidate := range candidates {
			if imageName == candidate {
				return true
			}
		}
	}
	return false
}

func pullSSHFixtureImageWithRetry(env *framework.ScenarioEnv, imageRef string) error {
	var lastErr error
	for attempt := 1; attempt <= sshFixturePullAttempts; attempt++ {
		if err := framework.RunCommand(env.TestCtx.Context, "docker", "pull", imageRef); err != nil {
			lastErr = err
		} else {
			return nil
		}

		if attempt == sshFixturePullAttempts {
			break
		}

		delay := time.Duration(attempt*5) * time.Second
		fmt.Printf("Retrying SSH fixture image pull for %q after %s (attempt %d/%d failed).\n", imageRef, delay, attempt, sshFixturePullAttempts)
		select {
		case <-time.After(delay):
		case <-env.TestCtx.Context.Done():
			return env.TestCtx.Context.Err()
		}
	}
	return lastErr
}

func assertSSHAppProtocolTrafficRules(env *framework.ScenarioEnv, session *e2eutils.Session, sandboxID string, fixture *sshFixture) {
	Expect(sandboxID).NotTo(BeEmpty())
	Expect(fixture).NotTo(BeNil())

	sandbox, _, err := session.GetSandbox(env.TestCtx.Context, GinkgoT(), sandboxID)
	Expect(err).NotTo(HaveOccurred())
	templateNamespace, err := naming.TemplateNamespaceForBuiltin(sandbox.TemplateId)
	Expect(err).NotTo(HaveOccurred())

	sandbox = waitForSandboxPodReadyEventually(env, session, sandboxID, templateNamespace)

	Expect(installSSHFixturePrivateKey(env, templateNamespace, sandbox.PodName)).To(Succeed())

	clearPolicy := apispec.SandboxNetworkPolicy{
		Mode: apispec.AllowAll,
	}
	defer func() {
		_, _, _, _ = session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, clearPolicy)
	}()

	serviceCIDR := fixture.ServiceIP + "/32"
	sshCommand := buildSSHFixtureCommand(fixture, "/tmp/sandbox0-e2e-ssh-key")

	_, status, apiErr, err := session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, clearPolicy)
	Expect(err).NotTo(HaveOccurred())
	Expect(apiErr).To(BeNil())
	Expect(status).To(Equal(200))
	assertSSHFixtureCommandEventuallySucceeds(env, templateNamespace, sandbox.PodName, sshCommand)

	Expect(updateSandboxTrafficRulesPolicy(
		env,
		session,
		sandboxID,
		apispec.BlockAll,
		buildTrafficRule("allow-tls-only", apispec.Allow, serviceCIDR, apispec.TrafficRuleAppProtocolTls),
	)).To(Succeed())
	assertSSHFixtureCommandEventuallyFails(env, templateNamespace, sandbox.PodName, sshCommand)

	Expect(updateSandboxTrafficRulesPolicy(
		env,
		session,
		sandboxID,
		apispec.BlockAll,
		buildTrafficRule("allow-ssh", apispec.Allow, serviceCIDR, apispec.TrafficRuleAppProtocolSsh),
	)).To(Succeed())
	assertSSHFixtureCommandEventuallySucceeds(env, templateNamespace, sandbox.PodName, sshCommand)

	Expect(updateSandboxTrafficRulesPolicy(
		env,
		session,
		sandboxID,
		apispec.AllowAll,
		buildTrafficRule("deny-ssh", apispec.Deny, serviceCIDR, apispec.TrafficRuleAppProtocolSsh),
	)).To(Succeed())
	assertSSHFixtureCommandEventuallyFails(env, templateNamespace, sandbox.PodName, sshCommand)
}

func installSSHFixturePrivateKey(env *framework.ScenarioEnv, namespace, podName string) error {
	return installSSHPrivateKey(env, namespace, podName, "/tmp/sandbox0-e2e-ssh-key", sshFixturePrivateKey)
}

func installSSHProxyFakePrivateKey(env *framework.ScenarioEnv, namespace, podName string) error {
	return installSSHPrivateKey(env, namespace, podName, "/tmp/sandbox0-e2e-ssh-fake-key", sshProxyFakePrivateKey)
}

func installSSHPrivateKey(env *framework.ScenarioEnv, namespace, podName, path, privateKey string) error {
	_, err := execInSandboxPod(env, namespace, podName, fmt.Sprintf(`set -eu
command -v ssh >/dev/null
cat <<'EOF' >%s
%s
EOF
chmod 600 %s
`, path, strings.TrimSpace(privateKey), path))
	return err
}

func buildSSHFixtureCommand(fixture *sshFixture, keyPath string) string {
	return fmt.Sprintf(
		"ssh -F /dev/null -i %s -o IdentitiesOnly=yes -o BatchMode=yes -o ConnectTimeout=5 -o LogLevel=ERROR -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -p 22 %s@%s 'printf ok'",
		keyPath,
		fixture.UserName,
		fixture.ServiceIP,
	)
}

func buildTrafficRule(name string, action apispec.TrafficRuleAction, cidr string, appProtocol apispec.TrafficRuleAppProtocol) apispec.TrafficRule {
	cidrs := []string{cidr}
	appProtocols := []apispec.TrafficRuleAppProtocol{appProtocol}
	ports := []apispec.PortSpec{{
		Port:     22,
		Protocol: ptrTo("tcp"),
	}}
	return apispec.TrafficRule{
		Name:         ptrTo(name),
		Action:       action,
		Cidrs:        &cidrs,
		Ports:        &ports,
		AppProtocols: &appProtocols,
	}
}

func updateSandboxTrafficRulesPolicy(
	env *framework.ScenarioEnv,
	session *e2eutils.Session,
	sandboxID string,
	mode apispec.SandboxNetworkPolicyMode,
	rule apispec.TrafficRule,
) error {
	rules := []apispec.TrafficRule{rule}
	policy, status, apiErr, err := session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, apispec.SandboxNetworkPolicy{
		Mode: mode,
		Egress: &apispec.NetworkEgressPolicy{
			TrafficRules: &rules,
		},
	})
	if err != nil {
		return err
	}
	if apiErr != nil {
		return fmt.Errorf("update network policy returned API error: %s", apiErr.Error.Message)
	}
	if status != 200 {
		return fmt.Errorf("update network policy returned status %d", status)
	}
	if policy == nil || policy.Egress == nil || policy.Egress.TrafficRules == nil || len(*policy.Egress.TrafficRules) != 1 {
		return fmt.Errorf("updated network policy did not persist the expected traffic rule")
	}
	return nil
}

func assertSSHTransparentEgressAuthProxy(env *framework.ScenarioEnv, session *e2eutils.Session, sandboxID string, fixture *sshFixture) {
	Expect(sandboxID).NotTo(BeEmpty())
	Expect(fixture).NotTo(BeNil())

	sandbox, _, err := session.GetSandbox(env.TestCtx.Context, GinkgoT(), sandboxID)
	Expect(err).NotTo(HaveOccurred())
	templateNamespace, err := naming.TemplateNamespaceForBuiltin(sandbox.TemplateId)
	Expect(err).NotTo(HaveOccurred())

	sandbox = waitForSandboxPodReadyEventually(env, session, sandboxID, templateNamespace)
	Expect(clearNetdRedisTeamBandwidthKeys(env, session.SelectedTeamID())).To(Succeed())

	sourceName := fmt.Sprintf("e2e-ssh-key-%d", time.Now().UnixNano())
	refName := "api-ssh-proxy"
	ruleName := "api-ssh-proxy-rule"
	serviceCIDR := fixture.ServiceIP + "/32"
	ports := []apispec.PortSpec{{
		Port:     22,
		Protocol: ptrTo("tcp"),
	}}

	clearPolicy := apispec.SandboxNetworkPolicy{
		Mode:               apispec.AllowAll,
		CredentialBindings: &[]apispec.CredentialBinding{},
	}
	defer func() {
		_, _, _, _ = session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, clearPolicy)
		status, _, deleteErr := session.DeleteCredentialSource(env.TestCtx.Context, sourceName)
		if deleteErr == nil {
			Expect(status).To(Or(Equal(http.StatusOK), Equal(http.StatusNotFound)))
		}
	}()

	_, status, apiErr, err := session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, clearPolicy)
	Expect(err).NotTo(HaveOccurred())
	Expect(apiErr).To(BeNil())
	Expect(status).To(Equal(http.StatusOK))

	Expect(installSSHProxyFakePrivateKey(env, templateNamespace, sandbox.PodName)).To(Succeed())
	fakeKeyCommand := buildSSHFixtureCommand(fixture, "/tmp/sandbox0-e2e-ssh-fake-key")
	assertSSHFixtureCommandEventuallyFails(env, templateNamespace, sandbox.PodName, fakeKeyCommand)

	knownHosts := scanSSHFixtureKnownHosts(env, templateNamespace, sandbox.PodName, fixture)
	created, err := session.CreateCredentialSource(env.TestCtx.Context, GinkgoT(), apispec.CredentialSourceWriteRequest{
		Name:         sourceName,
		ResolverKind: apispec.StaticSshPrivateKey,
		Spec: apispec.CredentialSourceWriteSpec{
			StaticSSHPrivateKey: &apispec.StaticSSHPrivateKeySourceSpec{
				PrivateKeyPem: ptrTo(sshFixturePrivateKey),
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())
	Expect(created).NotTo(BeNil())
	Expect(created.Name).To(Equal(sourceName))

	policy, status, apiErr, err := session.UpdateNetworkPolicy(env.TestCtx.Context, GinkgoT(), sandboxID, apispec.SandboxNetworkPolicy{
		Mode: apispec.BlockAll,
		Egress: &apispec.NetworkEgressPolicy{
			TrafficRules: &[]apispec.TrafficRule{
				buildTrafficRule("allow-ssh-proxy-fixture", apispec.Allow, serviceCIDR, apispec.TrafficRuleAppProtocolSsh),
			},
			CredentialRules: &[]apispec.EgressCredentialRule{{
				Name:          &ruleName,
				CredentialRef: refName,
				Ports:         &ports,
				Protocol:      ptrTo(apispec.EgressAuthProtocolSsh),
				Rollout:       ptrTo(apispec.EgressAuthRolloutModeEnabled),
			}},
		},
		CredentialBindings: &[]apispec.CredentialBinding{{
			Ref:       refName,
			SourceRef: sourceName,
			Projection: apispec.ProjectionSpec{
				Type: apispec.SshProxy,
				SshProxy: &apispec.SSHProxyProjection{
					SandboxPublicKeys: &[]string{sshProxyFakePublicKey},
					UpstreamUsername:  &fixture.UserName,
					KnownHosts:        &knownHosts,
				},
			},
		}},
	})
	Expect(err).NotTo(HaveOccurred())
	Expect(apiErr).To(BeNil())
	Expect(status).To(Equal(http.StatusOK))
	Expect(policy).NotTo(BeNil())

	assertSSHFixtureCommandEventuallySucceeds(env, templateNamespace, sandbox.PodName, fakeKeyCommand)
}

func scanSSHFixtureKnownHosts(env *framework.ScenarioEnv, namespace, podName string, fixture *sshFixture) []string {
	output, err := execInSandboxPod(env, namespace, podName, fmt.Sprintf(`set -eu
command -v ssh-keyscan >/dev/null
ssh-keyscan -T 5 -p 22 %s 2>/dev/null | sed '/^#/d'
`, fixture.ServiceIP))
	Expect(err).NotTo(HaveOccurred())

	var lines []string
	for _, line := range strings.Split(output, "\n") {
		line = strings.TrimSpace(line)
		if line != "" {
			lines = append(lines, line)
		}
	}
	Expect(lines).NotTo(BeEmpty())
	return lines
}

func assertSSHFixtureCommandEventuallySucceeds(env *framework.ScenarioEnv, namespace, podName, command string) {
	Eventually(func() error {
		output, err := execInSandboxPod(env, namespace, podName, command)
		if err != nil {
			return err
		}
		if sshFixtureSuccessOutput(output) != "ok" {
			return fmt.Errorf("unexpected ssh output: %q", output)
		}
		return nil
	}).WithTimeout(45 * time.Second).WithPolling(3 * time.Second).Should(Succeed())
}

func assertSSHFixtureCommandEventuallyFails(env *framework.ScenarioEnv, namespace, podName, command string) {
	Eventually(func() error {
		output, err := execInSandboxPod(env, namespace, podName, command)
		if err == nil {
			return fmt.Errorf("ssh unexpectedly succeeded with output %q", output)
		}
		return nil
	}).WithTimeout(45 * time.Second).WithPolling(3 * time.Second).Should(Succeed())
}

func sshFixtureSuccessOutput(output string) string {
	lines := strings.Split(strings.ReplaceAll(output, "\r\n", "\n"), "\n")
	for i := len(lines) - 1; i >= 0; i-- {
		line := strings.TrimSpace(lines[i])
		if line != "" {
			return line
		}
	}
	return ""
}
