// Copyright (c) 2026 Sidero Labs, Inc.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.

// Package logreceiver provides a TCP server that receives logs from Talos.
package logreceiver

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/netip"
	"sync"

	"github.com/siderolabs/gen/containers"
	"go.uber.org/zap"

	"github.com/siderolabs/omni/client/pkg/panichandler"
)

// Server implements TCP server to receive JSON logs. It is similar to the one in Talos, except it doesn't try to parse
// the JSON and instead writes rawdata and uses newline as a separator.
//
//nolint:govet
type Server struct {
	// ctx is the root context for all connections.
	ctx context.Context //nolint:containedctx

	listener net.Listener
	handler  *ConnHandler
	logger   *zap.Logger

	closeOnce sync.Once
	// wg is used to wait for all connections to be closed.
	wg sync.WaitGroup

	m containers.ConcurrentMap[string, net.Conn]
}

// Handler handles a message or an error.
type Handler interface {
	HandleMessage(ctx context.Context, srcAddress netip.Addr, rawData []byte)
	HandleError(srcAddress netip.Addr, err error)
	HasLink(srcAddress netip.Addr) bool
}

// NewServer initializes new Server.
func NewServer(ctx context.Context, listener net.Listener, handler *ConnHandler, logger *zap.Logger) *Server {
	return &Server{
		ctx:      ctx,
		listener: listener,
		handler:  handler,
		logger:   logger,
	}
}

// Serve runs the TCP server loop.
func (srv *Server) Serve() error {
	for {
		conn, err := srv.listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return nil
			}

			return fmt.Errorf("error accepting connection: %w", err)
		}

		srcAddr, ok := conn.RemoteAddr().(*net.TCPAddr)
		if !ok {
			srv.logger.Error("error getting remote IP address")

			continue
		}

		remoteAddr, _ := netip.AddrFromSlice(srcAddr.IP)
		remoteAddress := conn.RemoteAddr().String()

		if !srv.handler.HasLink(remoteAddr) {
			srv.logger.Warn(
				"ignoring log handler connection attempt from the remote address: machine is not accepted yet",
				zap.String("addr", remoteAddr.String()),
			)

			if err = conn.Close(); err != nil {
				return err
			}

			continue
		}

		srv.wg.Add(1)
		srv.m.Set(remoteAddress, conn)

		panichandler.Go(func() {
			defer srv.wg.Done()

			srv.handler.HandleConn(srv.ctx, remoteAddr, conn)
			srv.m.Remove(remoteAddress)
		}, srv.logger)
	}
}

// Stop serving.
func (srv *Server) Stop() {
	srv.closeOnce.Do(func() {
		srv.listener.Close() //nolint:errcheck

		srv.m.ForEach(func(_ string, conn net.Conn) {
			remotrAddr := conn.RemoteAddr().String()
			srv.logger.Info("closing connection", zap.String("remote_addr", remotrAddr))

			tcpConn, ok := conn.(*net.TCPConn)
			if ok {
				// set linger to 0 to forcefully close the connection with RST.
				// This configuration is needed to handle clients that only do WRITE operation on this connection.
				// Such clients are not doing any READ operation, therefore connections can't be terminated gracefully.
				// Termination only happens when client attempts to WRITE to an already closed connection. Since
				// such a scenario depends on the client, it's safer to forcefully terminate the connection on the
				// client by sending an RST.
				if err := tcpConn.SetLinger(0); err != nil {
					srv.logger.Error("error setting linger on connection", zap.Error(err))
				}
			}

			err := conn.Close()
			if err != nil {
				srv.logger.Warn("error closing connection", zap.String("remote_addr", remotrAddr), zap.Error(err))
			}
		})

		srv.logger.Info("waiting for goroutines to finish")
		// Wait for all connections goroutines to finish.
		srv.wg.Wait()
		srv.m.Clear()
		srv.logger.Info("stopped logging server")
	})
}

// ConnHandler is called for each received connection.
type ConnHandler struct {
	msgHandler Handler
	logger     *zap.Logger
}

// NewConnHandler initializes new ConnHandler.
func NewConnHandler(msgHandler Handler, logger *zap.Logger) *ConnHandler {
	return &ConnHandler{
		msgHandler: msgHandler,
		logger:     logger,
	}
}

// HandleConn handles a connection.
func (ch *ConnHandler) HandleConn(ctx context.Context, addr netip.Addr, conn io.ReadCloser) {
	defer conn.Close() //nolint:errcheck

	bufReader := bufio.NewReader(conn)

	for {
		slice, err := bufReader.ReadSlice('\n')
		if err != nil {
			if !errors.Is(err, net.ErrClosed) && !errors.Is(err, io.EOF) && !isTimeout(err) {
				ch.logger.Error("error decoding message", zap.Error(err))
				ch.msgHandler.HandleError(addr, err)
			}

			return
		}

		ch.msgHandler.HandleMessage(ctx, addr, slice[:len(slice)-1])
	}
}

// HasLink returns true if the instance has the link with the corresponding remote address.
func (ch *ConnHandler) HasLink(addr netip.Addr) bool {
	return ch.msgHandler.HasLink(addr)
}

func isTimeout(err error) bool {
	var neterr net.Error
	if errors.As(err, &neterr) {
		return neterr != nil && neterr.Timeout()
	}

	return false
}

// MakeServer creates a listener on the given address and returns a struct which can be used to start and stop the server.
func MakeServer(ctx context.Context, address string, handler Handler, logger *zap.Logger) (*Server, error) {
	listener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", address)
	if err != nil {
		return nil, fmt.Errorf("log server: error listening on %s: %w", address, err)
	}

	return NewServer(ctx, listener, NewConnHandler(handler, logger), logger), nil
}
