// Package sshserver implements an SSH server for git clone/push/pull operations.
// It authenticates users via their stored SSH public keys and executes
// git-upload-pack / git-receive-pack as subprocesses.
package sshserver

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/ssh"
	"xorm.io/xorm"

	"github.com/forgeo/forgebucket/internal/config"
)

// Server is the SSH git server.
type Server struct {
	db  *xorm.Engine
	cfg *config.Config
}

func New(db *xorm.Engine, cfg *config.Config) *Server {
	return &Server{db: db, cfg: cfg}
}

// ListenAndServe binds to cfg.SSHPort, loads or generates a host key, and accepts
// connections until ctx is cancelled. Returns nil when the context is done.
func (s *Server) ListenAndServe(ctx context.Context) error {
	hostKey, err := s.loadOrGenerateHostKey()
	if err != nil {
		return fmt.Errorf("sshserver: host key: %w", err)
	}

	srvConfig := &ssh.ServerConfig{
		PublicKeyCallback: s.lookupKey,
	}
	srvConfig.AddHostKey(hostKey)

	addr := ":" + s.cfg.SSHPort
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		log.Printf("sshserver: cannot bind %s — SSH transport disabled: %v", addr, err)
		return nil
	}
	log.Printf("sshserver: listening on %s", addr)

	go func() {
		<-ctx.Done()
		ln.Close()
	}()

	for {
		conn, err := ln.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return nil
			default:
				log.Printf("sshserver: accept: %v", err)
				continue
			}
		}
		go s.handleConn(conn, srvConfig)
	}
}

func (s *Server) handleConn(netConn net.Conn, srvConfig *ssh.ServerConfig) {
	defer netConn.Close()

	sshConn, chans, reqs, err := ssh.NewServerConn(netConn, srvConfig)
	if err != nil {
		return
	}
	defer sshConn.Close()

	go ssh.DiscardRequests(reqs)

	username, _ := sshConn.Permissions.Extensions["username"]

	for newChan := range chans {
		if newChan.ChannelType() != "session" {
			newChan.Reject(ssh.UnknownChannelType, "unknown channel type") //nolint:errcheck
			continue
		}
		ch, requests, err := newChan.Accept()
		if err != nil {
			return
		}
		go s.handleSession(ch, requests, username)
	}
}

// loadOrGenerateHostKey loads the host key from SSHHostKeyPath if set,
// otherwise generates an ephemeral RSA-4096 key (lost on restart).
func (s *Server) loadOrGenerateHostKey() (ssh.Signer, error) {
	if s.cfg.SSHHostKeyPath != "" {
		data, err := os.ReadFile(s.cfg.SSHHostKeyPath)
		if err != nil {
			return nil, fmt.Errorf("read host key %s: %w", s.cfg.SSHHostKeyPath, err)
		}
		return ssh.ParsePrivateKey(data)
	}

	log.Printf("sshserver: SSH_HOST_KEY_PATH not set — generating ephemeral host key (host key changes on restart)")
	privKey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, fmt.Errorf("generate host key: %w", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privKey),
	})
	return ssh.ParsePrivateKey(keyPEM)
}