package ssh

import (
	"farga/db"
	"farga/db/types"
	"farga/lib"
	"farga/lib/git"
	"fmt"
	"io"
	"log"
	"net"
	"os/exec"
	"path/filepath"
	"strings"

	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"os"

	"github.com/google/uuid"
	"golang.org/x/crypto/ssh"
)

// Tries to load the host's keys. If there are none, it makes new keys.
func loadOrGenerateHostKey(path string) (ssh.Signer, error) {
	// Is there one already?
	if data, err := os.ReadFile(path); err == nil {
		signer, err := ssh.ParsePrivateKey(data)
		if err != nil {
			return nil, fmt.Errorf("parse host key: %w", err)
		}
		return signer, nil
	}

	// Generate a new ed25519 priv key.
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate private key: %w", err)
	}

	// Marshal to OpenSSH PEM format.
	privPemBlock, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		return nil, fmt.Errorf("marshal key: %w", err)
	}

	// Save key.
	if err := os.WriteFile(path, pem.EncodeToMemory(privPemBlock), 0600); err != nil {
		return nil, fmt.Errorf("write host key: %w", err)
	}

	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		return nil, fmt.Errorf("signer: %w", err)
	}
	return signer, nil
}

func SSHConfig(d *db.DB) *ssh.ServerConfig {
	config := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return AuthenticatePublicKey(d, conn, key)
		},
	}

	keyPath, err := filepath.Abs(lib.SSHHostKeyPath)
	if err != nil {
		log.Fatal("Maybe the SSH host key path is improperly configured.")
	}

	hostKey, err := loadOrGenerateHostKey(keyPath) // todo
	if err != nil {
		log.Fatalf("host key: %v", err)
	}
	config.AddHostKey(hostKey)

	return config
}

func ServeSSH(d *db.DB, port int) {
	config := SSHConfig(d)
	listener, err := net.Listen("tcp", fmt.Sprint("0.0.0.0:", port))
	if err != nil {
		log.Fatalf("listen: %v", err)
	}
	fmt.Println(fmt.Sprint("Listening for SSH on 0.0.0.0:", lib.SSHPort))

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Printf("accept: %v", err)
			continue
		}
		go handleConn(d, conn, config)
	}
}

// Function commented since I might forget how it works
func handleConn(d *db.DB, conn net.Conn, config *ssh.ServerConfig) {
	// NewServerConn upgrades the connection from raw TCP to SSH.
	// PublicKeyCallback will be fired and the client will be authenticated
	//
	// sshConn: the connection
	// chans: SSH Channels that can request stuff?.
	// reqs: Connection-level requests. Stuff such as `keepalive`.
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
	if err != nil {
		log.Printf("handshake failed: %v", err)
		return
	}

	// SSH global requests do not drain themselves. Avoids deadlocks.
	go ssh.DiscardRequests(reqs)

	// Git clients only use the "session" channel. Therefore, all other channels must be rejected.
	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unsupported channel type")
			continue
		}
		go handleSessionChannel(d, newChannel, sshConn.User(), sshConn.Permissions)
	}

	// sshConn is used just to get the username, in this case.
	defer sshConn.Close()
}

// Also commented. Handles the individual channel.
func handleSessionChannel(d *db.DB, newChannel ssh.NewChannel, _ string, sshPerms *ssh.Permissions) {
	// Accept the channel since it IS the git channel.
	//
	// channel: bi-direcional stream (stdin stdout stderr etc.)
	// requests: channel-level requests, such as "kill channel"
	channel, requests, err := newChannel.Accept()
	if err != nil {
		log.Printf("accept channel: %v", err)
		return
	}
	defer channel.Close()

	// For every request sent,
	for req := range requests {
		switch req.Type {
		// Is it a command?
		case "exec":
			// YEAH, RUN IT!
			handleExec(d, channel, req, sshPerms.Extensions)
			return
		// No?
		default:
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}
}

func handleExec(d *db.DB, channel ssh.Channel, req *ssh.Request, extensions map[string]string) {
	// Unmarshal the SSH into a string
	var payload struct{ Command string }
	if err := ssh.Unmarshal(req.Payload, &payload); err != nil {
		req.Reply(false, nil)
		return
	}

	// Is the string valid?
	command, err := parseGitCommand(payload.Command)
	if err != nil {
		closeChannel(channel, req, fmt.Sprintf("invalid command: %v\n", err))
		return
	}

	// Retrieve the user ID we previously got.
	userID, _ := uuid.Parse(extensions[userIDKey])

	repoUser, repoName, err := getUserAndRepo(d, command.Args)
	if err != nil {
		closeChannel(channel, req, "No such user.\n")
		return
	}

	repoMetadata := types.Repo{
		UserID: &repoUser.ID,
		Name:   repoName,
	}

	repoMetadata, err = d.GetRepo(repoMetadata)

	repoPath := git.GetRepoPath(&repoUser.ID, &repoMetadata.ID)

	if err != nil {
		req.Reply(false, nil)
		fmt.Fprintf(channel.Stderr(), "repo not found: %v\n", err)
		channel.CloseWrite()
		sendExitStatus(channel, 1)
		return
	}

	// Authorization: does sshUser have access to repoPath?
	if (repoMetadata.IsPrivate || strings.Contains(command.Command, "upload")) && *repoMetadata.UserID != userID {
		fmt.Println("unauthorized")
		fmt.Println(repoMetadata.UserID)
		fmt.Println(userID)
		req.Reply(false, nil)
		fmt.Fprintln(channel.Stderr(), "access denied")
		channel.CloseWrite()
		sendExitStatus(channel, 1)
		return
	}

	req.Reply(true, nil)

	cmd := exec.Command(command.Command, repoPath)
	cmd.Stdin = channel
	runGitCommand(cmd, channel)
}

func runGitCommand(cmd *exec.Cmd, channel ssh.Channel) {
	// stdin, _ := cmd.StdinPipe()
	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()

	if err := cmd.Start(); err != nil {
		fmt.Fprintf(channel.Stderr(), "failed to start: %v\n", err)
		sendExitStatus(channel, 1)
		return
	}

	go io.Copy(channel, stdout)
	go io.Copy(channel.Stderr(), stderr)

	err := cmd.Wait()
	channel.CloseWrite()

	exitCode := 0
	if exitErr, ok := err.(*exec.ExitError); ok {
		exitCode = exitErr.ExitCode()
	} else if err != nil {
		exitCode = 1
	}

	sendExitStatus(channel, exitCode)
}