// SPDX-FileCopyrightText: 2025 Romain Maneschi <romain@gitroot.dev>
//
// SPDX-License-Identifier: EUPL-1.2

package main

import (
	"bytes"
	"compress/zlib"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/anmitsu/go-shlex"
	securejoin "github.com/cyphar/filepath-securejoin"
	"github.com/go-git/go-billy/v5/osfs"
	"github.com/go-git/go-git/v5/plumbing/format/pktline"
	"github.com/go-git/go-git/v5/plumbing/object"
	"github.com/go-git/go-git/v5/plumbing/protocol/packp"
	"github.com/go-git/go-git/v5/plumbing/protocol/packp/capability"
	"github.com/go-git/go-git/v5/plumbing/server"
	"github.com/go-git/go-git/v5/plumbing/transport"
	"github.com/samber/oops"
	"gitroot.dev/server/background"
	"gitroot.dev/server/configuration"
	"gitroot.dev/server/logger"
	"gitroot.dev/server/plugin"
	"gitroot.dev/server/repository"
	"gitroot.dev/server/user"
	"golang.org/x/crypto/ssh"
)

const (
	SSH_EXTENSIONS_KEY_PUBKEY_FP = "pubkey-fp"
)

type sshServer struct {
	keys              map[string]ssh.PublicKey
	logger            *logger.Logger
	conf              *configuration.Configuration
	pluginManager     *plugin.Manager
	repoManager       *repository.Manager
	userManager       *user.Manager
	backgroundManager *background.Manager
}

func NewServerSsh(conf *configuration.Configuration, repoManager *repository.Manager, userManager *user.Manager, pluginManager *plugin.Manager, backgroundManager *background.Manager) *sshServer {
	return &sshServer{
		keys:              make(map[string]ssh.PublicKey),
		logger:            logger.NewLoggerCtx(logger.SSH_SERVER_LOGGER_NAME, context.Background()),
		conf:              conf,
		pluginManager:     pluginManager,
		repoManager:       repoManager,
		userManager:       userManager,
		backgroundManager: backgroundManager,
	}
}

func (srv *sshServer) ListenAndServe() error {
	config := &ssh.ServerConfig{
		NoClientAuth: false,
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			sha256fp := ssh.FingerprintSHA256(pubKey)
			srv.keys[sha256fp] = pubKey
			return &ssh.Permissions{
				// Record the public key used for authentication.
				Extensions: map[string]string{
					SSH_EXTENSIONS_KEY_PUBKEY_FP: sha256fp,
				},
			}, nil
		},
	}

	config.AddHostKey(srv.userManager.RootCommiter().Signer.Signer())

	srv.logger.Warn("starting SSH server on", logger.NewLoggerPair("addr", srv.conf.SshAddr))

	lis, err := net.Listen("tcp", srv.conf.SshAddr)
	if err != nil {
		return err
	}
	defer lis.Close()
	for {
		conn, err := lis.Accept()
		srv.logger.PrintMemUsage()
		if err != nil {
			return err
		}

		go func(conn net.Conn) {
			defer conn.Close()

			sshConn, newChanChan, newReq, err := ssh.NewServerConn(conn, config)
			if err != nil {
				srv.logger.Error("error when creating server conn", err)
				return
			}
			defer sshConn.Close()
			srv.logger.Info("new ssh connexion", logger.NewLoggerPair("user", sshConn.Conn.User()), logger.NewLoggerPair("key", sshConn.Permissions.Extensions[SSH_EXTENSIONS_KEY_PUBKEY_FP]))
			go ssh.DiscardRequests(newReq)
			for newChan := range newChanChan {
				if newChan.ChannelType() == "session" {
					ch, reqc, err := newChan.Accept()
					if err != nil {
						srv.logger.Error("error when accepting session", err)
						return
					}
					srv.newSession(ch, reqc, sshConn.Conn.User(), sshConn.Permissions.Extensions[SSH_EXTENSIONS_KEY_PUBKEY_FP]).handle()
				} else {
					srv.logger.Error("unknown channel type", errors.New("channel unknown"), logger.NewLoggerPair("pair", newChan.ChannelType()))
				}
			}
		}(conn)
	}
}

func (srv *sshServer) newSession(ch ssh.Channel, reqc <-chan *ssh.Request, pseudo string, sshKeyFp string) *sshSession {
	ctx, cnlCtx := context.WithCancel(context.Background())
	return &sshSession{
		srv:        srv,
		logger:     *srv.logger.NewSubLoggerCtx("SshSession", ctx),
		ch:         ch,
		reqc:       reqc,
		ctx:        ctx,
		cnlCtx:     cnlCtx,
		simpleUser: user.SimpleUser{Pseudo: pseudo, Ssh: string(bytes.TrimSuffix(ssh.MarshalAuthorizedKey(srv.keys[sshKeyFp]), []byte("\n")))},
	}
}

type sshSession struct {
	srv        *sshServer
	logger     logger.Logger
	ch         ssh.Channel
	reqc       <-chan *ssh.Request
	ctx        context.Context
	cnlCtx     context.CancelFunc
	simpleUser user.SimpleUser
}

func (session *sshSession) handle() {
	var exitCode uint32 = 0
	defer func() {
		b := ssh.Marshal(struct{ Value uint32 }{exitCode})
		_, err := session.ch.SendRequest("exit-status", false, b)
		if err != nil {
			session.logger.Error("SendRequest exit-status error", err, logger.NewLoggerPair("exitCode", exitCode))
		}
		time.Sleep(10 * time.Millisecond)
		session.ch.Close()
		session.cnlCtx()
	}()

	envs := make(map[string]string)
	for req := range session.reqc {
		switch req.Type {
		case "env":
			payload := struct{ Key, Value string }{}
			ssh.Unmarshal(req.Payload, &payload)
			envs[payload.Key] = payload.Value
			req.Reply(true, nil)
		case "exec":
			payload := struct{ Value string }{}
			ssh.Unmarshal(req.Payload, &payload)
			args, err := shlex.Split(payload.Value, true)
			if err != nil {
				session.logger.Error("shlex args", err)
				exitCode = 1
				return
			}

			cmd := args[0]
			name := strings.TrimPrefix(args[1], "/")
			if name == "" {
				name = session.srv.conf.ForgeConfigName()
			}
			dir, err := securejoin.SecureJoin(session.srv.conf.PathRepositories(), name)
			if err != nil {
				session.logger.Error("invalid repo upload pack", err, logger.NewLoggerPair("arg", name))
				exitCode = 1
				return
			}

			session.logger.Info("ssh request", logger.NewLoggerPair("cmd", cmd), logger.NewLoggerPair("dir", dir), logger.NewLoggerPair("name", name))

			switch cmd {
			case "git-upload-pack": // read
				// if gp := envs["GIT_PROTOCOL"]; gp != "version=2" {
				// 	log.Println("unhandled GIT_PROTOCOL", gp)
				// 	exitCode = 1
				// 	return
				// }
				err = session.handleUploadPack(dir, name)
				if err != nil {
					session.logger.Error("handle upload pack error", err)
					pktline.WriteError(session.ch, err)
					pktline.WriteFlush(session.ch)
					exitCode = 1
					return
				}

				session.logger.Info("finish upload plack", logger.NewLoggerPair("path", name))
				if err := req.Reply(true, nil); err != nil {
					session.logger.Error("req reply error upload pack", err)
				}
				return
			case "git-receive-pack": // write
				session.logger.Info("start receive plack")

				err = session.handleReceivePack(dir, name)
				if err != nil {
					session.logger.Error("handle receive pack error", err)
					exitCode = 1
					return
				}

				session.logger.Info("finish receive plack", logger.NewLoggerPair("path", name))
				if err := req.Reply(true, nil); err != nil {
					session.logger.Error("req reply error receive pack", err)
				}
				return
			default:
				session.logger.Error("unhandled cmd", errors.New("unknown cmd"), logger.NewLoggerPair("cmd", cmd))
				req.Reply(false, nil)
				exitCode = 1
				return
			}
		case "auth-agent-req@openssh.com":
			if req.WantReply {
				req.Reply(true, nil)
			}
		default:
			session.logger.Error("unhandled req type", errors.New("unknown req type"), logger.NewLoggerPair("type", req.Type))
			req.Reply(false, nil)
			exitCode = 1
			return
		}
	}
}

func (session *sshSession) handleReceivePack(dir string, repoName string) error {
	errHandler := oops.In("sshSession").Code("handleReceivePack").With("session", session.simpleUser.Pseudo)

	ep, err := transport.NewEndpoint(dir)
	if err != nil {
		return errHandler.Wrapf(err, "create transport endpoint")
	}
	ld, repo, writer, err := repository.NewGitRootFsLoader(session.ctx, repoName, session.srv.repoManager)
	if err != nil {
		return errHandler.Wrapf(err, "repo not found")
	}
	writer.Reject() //by default we reject == in case of errror we don't take new code
	defer ld.Close()
	svr := server.NewServer(ld)
	session.logger.Info("NewReceivePackSession")
	sess, err := svr.NewReceivePackSession(ep, nil)
	if err != nil {
		return errHandler.Wrapf(err, "create receive-pack session")
	}

	ar, err := sess.AdvertisedReferencesContext(session.ctx)
	if err != nil {
		return errHandler.Wrapf(err, "get advertised references")
	}
	if err := ar.Encode(session.ch); err != nil {
		return errHandler.Wrapf(err, "encode advertised references")
	}

	rur := packp.NewReferenceUpdateRequest()
	defer func() {
		if rur.Packfile != nil {
			rur.Packfile.Close()
		}
	}()
	if err := rur.Decode(io.NopCloser(session.ch)); err != nil {
		if errors.Is(err, packp.ErrEmptyCommands) || err.Error() == "capabilities delimiter not found" {
			return pktline.WriteResponseEnd(session.ch)
		}
		return errHandler.Wrapf(err, "decode reference-update request")
	}

	onlyDelete := true
	for _, cmd := range rur.Commands {
		onlyDelete = onlyDelete && cmd.Action() == "delete"
	}
	if onlyDelete {
		rur.Packfile.Close()
		rur.Packfile = nil
	}

	res, err := sess.ReceivePack(session.ctx, rur)
	if errors.Is(err, packp.ErrEmpty) {
		//nothing pushed == nothing todo
		return nil
	} else if err != nil {
		return errHandler.Wrapf(err, "create receive-pack response")
	}

	repoConfiguration, err := repo.Configuration()
	if err != nil {
		return errHandler.Wrapf(err, "repo configuration")
	}

	for _, c := range rur.Commands {
		session.logger.Info("Command", logger.NewLoggerPair("branch", c.Name.Short()), logger.NewLoggerPair("from", c.Old.String()), logger.NewLoggerPair("to", c.New.String()))
		canWrite, currentUser, err := repo.CanWrite(session.simpleUser.Ssh, c.Name.Short())
		if err != nil {
			errUser := errors.New("error in finding right")
			pktline.WriteError(session.ch, errUser)
			return oops.Wrapf(err, "error in finding right")
		}
		if !canWrite {
			err := fmt.Errorf("you can't write in %s", c.Name.Short())
			pktline.WriteError(session.ch, err)
			session.logger.Info("user can't write in branch", logger.NewLoggerPair("branch", c.Name))
			return nil
		}

		if c.Action() == packp.Update {
			isForcePush, err := writer.IsForcePush(c.Name, c.Old)
			session.logger.Info("ForcePush??", logger.NewLoggerPair("isForcePush", isForcePush), logger.NewLoggerPair("err", err))
			if isForcePush && repoConfiguration.IsNoPushBranch(c.Name) {
				err := errors.New("you can't force-push on main")
				pktline.WriteError(session.ch, err)
				return nil //don't return error
			}
			if isForcePush {
				com, err := writer.GetLastCommit(c.New)
				if err != nil {
					errUser := errors.New("not found previous commit")
					pktline.WriteError(session.ch, errUser)
					return oops.Wrapf(err, "error force-push previous commit")
				}
				p, err := com.Commit.Parents().Next() //TODO need to find all commits of force-pushed branch, only one in tests
				c.Old = p.Hash
			}
		}

		err = repo.WalkCommit(c.Old, c.New, func(commit *object.Commit) error {
			if commit.PGPSignature == "" {
				err := errors.New("you need to sign your commit")
				pktline.WriteError(session.ch, err)
				return oops.With("no signed commit", commit.Hash.String()).Wrap(errors.New("commit not signed"))
			}

			if err = Verify(session.simpleUser.Ssh, commit); err != nil {
				errUser := errors.New("ssh/pgp signature not good")
				pktline.WriteError(session.ch, errUser)
				return oops.With("commit", commit.Hash).Wrapf(err, "error in key and signature")
			}

			return nil
		})
		if err != nil {
			return errHandler.With("from", c.Old.String()).With("to", c.New.String()).Wrapf(err, "can't iter over commits")
		}

		if currentUser.Group == nil && c.Action() != packp.Delete { // user push new branch
			session.logger.Info("Add user into branch", logger.NewLoggerPair("branch", c.Name.Short()))
			if err := writer.AddUserInfo(currentUser, c.Name.Short()); err != nil {
				session.logger.Error("can't AddUserInfo", err)
			}
		}
	}

	if err := res.Encode(session.ch); err != nil {
		return errHandler.Wrapf(err, "encode receive-pack response")
	}

	writer.Accept() //after all if no error we accept changes

	session.srv.backgroundManager.PostPush(session.simpleUser, repoName, rur.Commands)
	return nil
}

func (session *sshSession) handleUploadPack(dir string, name string) error {
	errHandler := oops.In("sshSession").Code("handleUploadPack").With("session", session.simpleUser.Pseudo).With("dir", dir)

	if dir == "/" {
		dir = session.srv.conf.RootRepositoryName
	}

	if name == session.srv.conf.ForgeConfigName() {
		if err := session.srv.repoManager.ForgeRepoNeedOwner(session.ctx, session.simpleUser); err != nil {
			return errHandler.Wrapf(err, "NeedOwner")
		}
	}

	repo, err := session.srv.repoManager.Open(logger.AddCaller(session.ctx, "handleUploadPack"), name)
	if err != nil {
		return errHandler.With("name", name).Wrapf(err, "can't open repo")
	}
	defer repo.Close()

	ep, err := transport.NewEndpoint("/")
	if err != nil {
		return errHandler.Wrapf(err, "create transport endpoint")
	}
	bfs := osfs.New(dir)
	ld := server.NewFilesystemLoader(bfs)
	svr := server.NewServer(ld)
	sess, err := svr.NewUploadPackSession(ep, nil)
	if err != nil {
		return errHandler.Wrapf(err, "create upload-pack session")
	}
	defer sess.Close()

	ar, err := sess.AdvertisedReferencesContext(session.ctx)
	if err != nil {
		return errHandler.Wrapf(err, "get advertised references")
	}
	if err := ar.Capabilities.Add(capability.ThinPack); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.OFSDelta); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.MultiACK); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.MultiACKDetailed); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.Shallow); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.DeepenRelative); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.DeepenSince); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	if err := ar.Capabilities.Add(capability.DeepenNot); err != nil {
		return errHandler.Wrapf(err, "set advertised capabilities")
	}
	// if err := ar.Capabilities.Add(capability.Sideband64k); err != nil {
	// 	return errHandler.Wrapf(err, "set advertised capabilities")
	// }
	err = ar.Encode(session.ch)
	if err != nil {
		return errHandler.Wrapf(err, "encode advertised references")
	}

	capa := capability.NewList()
	capa.Add(capability.Shallow)
	upr := packp.NewUploadPackRequestFromCapabilities(capa)
	err = upr.Decode(session.ch)
	if err != nil {
		return errHandler.Wrapf(err, "decode upload-pack request")
	}
	session.logger.Info("upr", logger.NewLoggerPair("upr", upr))

	res, err := sess.UploadPack(session.ctx, upr)
	if err != nil && !errors.Is(err, zlib.ErrHeader) {
		if errors.Is(err, transport.ErrEmptyUploadPackRequest) {
			return nil
		}
		return errHandler.Wrapf(err, "create upload-pack response")
	}
	session.logger.Info("res", logger.NewLoggerPair("res", res))
	err = res.Encode(session.ch)
	if err != nil {
		return errHandler.Wrapf(err, "encode upload-pack response")
	}

	return res.Close()
}