//go:build linux

package microvm

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

// ssh ends here and gets bridged into the failed microvm over the guest's agent conn,
// which keeps keys off the guest. a key only gets in when the spindle says its owner
// can push to the job's repo

const debugAuthTimeout = 5 * time.Second

type DebugAuthorizer interface {
	DebugKeyAllowed(ctx context.Context, repoDid syntax.DID, key gossh.PublicKey) bool
}

// the ssh listener starts with the engine, before the spindle exists, so until
// this is set every key is refused
func (e *Engine) SetDebugAuthorizer(a DebugAuthorizer) {
	e.debugAuth.Store(&a)
}

func (e *Engine) serveDebugSSH(ctx context.Context) {
	dbg := e.cfg.MicroVMPipelines.DebugSSH
	addr := dbg.ListenAddr

	srv := &ssh.Server{
		Addr:    addr,
		Handler: e.debugHandle,
		PublicKeyHandler: func(c ssh.Context, key ssh.PublicKey) bool {
			return e.checkDebugAuth(c, c.User(), key)
		},
	}
	keyPath := e.cfg.MicroVMPipelines.DebugSSH.HostKeyPath
	if keyPath == "" {
		keyPath = filepath.Join(filepath.Dir(e.cfg.Server.DBPath), "debug_ssh_host_key")
	}
	if err := ensureDebugHostKey(keyPath); err != nil {
		e.l.Error("debug ssh: ensure host key", "path", keyPath, "err", err)
		return
	}
	if err := srv.SetOption(ssh.HostKeyFile(keyPath)); err != nil {
		e.l.Error("debug ssh: load host key", "path", keyPath, "err", err)
		return
	}

	go func() {
		<-ctx.Done()
		_ = srv.Close()
	}()

	e.l.Info("starting debug ssh server", "address", addr)
	if err := srv.ListenAndServe(); err != nil && err != ssh.ErrServerClosed {
		e.l.Error("debug ssh server stopped", "err", err)
	}
}

func ensureDebugHostKey(path string) error {
	if _, err := os.Stat(path); err == nil {
		return nil
	} else if !os.IsNotExist(err) {
		return fmt.Errorf("stat host key: %w", err)
	}

	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return fmt.Errorf("generate host key: %w", err)
	}
	block, err := gossh.MarshalPrivateKey(priv, "")
	if err != nil {
		return fmt.Errorf("marshal host key: %w", err)
	}

	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return fmt.Errorf("create host key dir: %w", err)
	}
	if err := os.WriteFile(path, pem.EncodeToMemory(block), 0o600); err != nil {
		return fmt.Errorf("write host key: %w", err)
	}
	return nil
}

func (e *Engine) debugHandle(sess ssh.Session) {
	ptyReq, winCh, isPty := sess.Pty()
	if !isPty {
		io.WriteString(sess.Stderr(), "error: no terminal allocated; use `ssh -t`\n")
		_ = sess.Exit(1)
		return
	}

	jobID := sess.User()
	l := e.l.With("component", "debugssh", "job", jobID)

	debug, err := e.OpenDebugSession(sess.Context(), jobID, ptyReq.Term, ptyReq.Window.Height, ptyReq.Window.Width)
	if err != nil {
		fmt.Fprintf(sess.Stderr(), "error: %v\n", err)
		_ = sess.Exit(1)
		return
	}
	defer debug.Close()
	l.Info("debug shell opened")

	go func() {
		for win := range winCh {
			if err := debug.Resize(win.Height, win.Width); err != nil {
				l.Debug("debug ssh resize failed", "error", err)
			}
		}
	}()

	// keyboard -> shell, runs until the client hangs up
	go func() {
		if _, err := io.Copy(debug, sess); err != nil {
			l.Debug("debug ssh input copy ended", "error", err)
		}
	}()
	// shell -> client, returns when the shell exits (Read hits EOF)
	_, outputErr := io.Copy(sess, debug)

	result := debug.result()
	code := 0
	guestError := ""
	guestReportedExit := result.exit != nil
	if guestReportedExit {
		code = int(result.exit.ExitCode)
		guestError = result.exit.Error
	}
	if guestError != "" && code == 0 {
		code = 1
	}
	if !guestReportedExit {
		code = 1
		detail := "debug shell channel closed before the guest reported an exit"
		if result.readError != "" {
			detail += ": " + result.readError
		}
		fmt.Fprintf(sess.Stderr(), "error: %s\n", detail)
	}
	if outputErr != nil {
		code = 1
		fmt.Fprintf(sess.Stderr(), "error: debug shell SSH output failed: %v\n", outputErr)
	}

	l.Info(
		"debug shell closed",
		"exitCode", code,
		"guestReportedExit", guestReportedExit,
		"guestError", guestError,
		"debugReadError", result.readError,
		"sshOutputError", outputErr,
	)
	// the user is done; let retention tear the VM down now instead of waiting
	// out the rest of the grace period
	e.releaseDebugTarget(jobID)
	_ = sess.Exit(code)
}

func (e *Engine) checkDebugAuth(ctx context.Context, jobID string, key ssh.PublicKey) bool {
	l := e.l.With("component", "debugssh", "job", jobID, "keyType", key.Type())

	target, ok := e.lookupDebugTarget(jobID)
	if !ok || target.repoDid == "" {
		l.Warn("debug ssh: no live job / unknown repo")
		return false
	}

	reqCtx, cancel := context.WithTimeout(ctx, debugAuthTimeout)
	defer cancel()
	if !e.keyMayDebug(reqCtx, target, key) {
		l.Warn("debug ssh: key not allowed to push", "repo", target.repoDid)
		return false
	}
	l.Info("debug ssh: authorized", "repo", target.repoDid)
	return true
}

// a mill-bound check wins, otherwise the spindle that owns the repo rows decides
func (e *Engine) keyMayDebug(ctx context.Context, target debugTarget, key ssh.PublicKey) bool {
	if target.keyCheck != nil {
		return target.keyCheck(ctx, key.Marshal())
	}
	authorizer := e.debugAuth.Load()
	if authorizer == nil {
		e.l.Warn("debug ssh: no authorizer wired, refusing")
		return false
	}
	did, err := syntax.ParseDID(target.repoDid)
	if err != nil {
		e.l.Error("debug ssh: bad repo did", "repo", target.repoDid, "error", err)
		return false
	}
	return (*authorizer).DebugKeyAllowed(ctx, did, key)
}