// Copyright (C) MongoDB, Inc. 2023-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0

package mnet

import (
	"context"
	"io"
	"sync/atomic"

	"go.mongodb.org/mongo-driver/v2/mongo/address"
	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/description"
)

// ReadWriteCloser represents a Connection where server operations
// can read from, written to, and closed.
type ReadWriteCloser interface {
	Read(ctx context.Context) ([]byte, error)
	Write(ctx context.Context, wm []byte) error
	io.Closer
}

// Describer represents a Connection that can be described.
type Describer interface {
	Description() description.Server
	ID() string
	ServerConnectionID() *int64
	DriverConnectionID() int64
	Address() address.Address
	Stale() bool
	OIDCTokenGenID() uint64
	SetOIDCTokenGenID(uint64)
}

// Streamer represents a Connection that supports streaming wire protocol
// messages using the moreToCome and exhaustAllowed flags.
//
// The SetStreaming and CurrentlyStreaming functions correspond to the
// moreToCome flag on server responses. If a response has moreToCome set,
// SetStreaming(true) will be called and CurrentlyStreaming() should return
// true.
//
// CanStream corresponds to the exhaustAllowed flag. The operations layer will
// set exhaustAllowed on outgoing wire messages to inform the server that the
// driver supports streaming.
type Streamer interface {
	SetStreaming(bool)
	CurrentlyStreaming() bool
	SupportsStreaming() bool
}

// Compressor is an interface used to compress wire messages. If a Connection
// supports compression it should implement this interface as well. The
// CompressWireMessage method will be called during the execution of an
// operation if the wire message is allowed to be compressed.
type Compressor interface {
	CompressWireMessage(src, dst []byte) ([]byte, error)
}

// Pinner represents a Connection that can be pinned by one or more cursors or
// transactions. Implementations of this interface should maintain the following
// invariants:
//
//  1. Each Pin* call should increment the number of references for the
//     connection.
//  2. Each Unpin* call should decrement the number of references for the
//     connection.
//  3. Calls to Close() should be ignored until all resources have unpinned the
//     connection.
type Pinner interface {
	PinToCursor() error
	PinToTransaction() error
	UnpinFromCursor() error
	UnpinFromTransaction() error
}

// Connection represents a connection to a MongoDB server.
type Connection struct {
	// closed is a shadow of the underlying ReadWriteCloser's state. It only
	// reflects intentional closes performed via Connection.Close(); it does
	// not catch silent underlying death (network drop, server-side close,
	// pool eviction, etc.).
	//
	// TODO(GODRIVER-3905): remove this field along with Close()/Closed()
	// once IsAlive() lands on the topology connection.
	closed uint32

	ReadWriteCloser
	Describer
	Streamer
	Compressor
	Pinner
}

// Close closes the connection.
//
// TODO(GODRIVER-3905): remove this override. Once GODRIVER-3603 lands
// the bufio.Reader on topology.connection, the "is this connection
// still usable?" question (the only reason this override exists, via
// the Closed flag) will be answered by a new IsAlive method that runs
// a peek-based liveness probe. Callers then get Close from the embedded
// ReadWriteCloser directly via interface promotion. See
// change_stream_deployment.Connection for the consumer that drives this.
func (c *Connection) Close() error {
	// Logically mark the connection as closed.
	atomic.StoreUint32(&c.closed, 1)

	return c.ReadWriteCloser.Close()
}

// Closed returns true if the connection has been logically closed.
//
// TODO(GODRIVER-3905): remove in favor of an IsAlive method on the
// topology connection. Closed answers "did our code call
// Connection.Close()?": a strict subset of "is this connection
// usable?". A bufio.Reader-backed Peek probe with a short read deadline
// catches the silent-death cases (peer FIN, RST, TLS abort, network
// drop already noticed by the kernel) and returns the broader answer.
func (c *Connection) Closed() bool {
	return atomic.LoadUint32(&c.closed) == 1
}

// NewConnection creates a new Connection with the provided component. This
// constructor returns a component that is already a Connection to avoid
// mis-asserting the composite interfaces.
func NewConnection(component interface {
	ReadWriteCloser
	Describer
},
) *Connection {
	if _, ok := component.(*Connection); ok {
		return component.(*Connection)
	}

	conn := &Connection{
		ReadWriteCloser: component,
	}

	if describer, ok := component.(Describer); ok {
		conn.Describer = describer
	}

	if streamer, ok := component.(Streamer); ok {
		conn.Streamer = streamer
	}

	if compressor, ok := component.(Compressor); ok {
		conn.Compressor = compressor
	}

	if pinner, ok := component.(Pinner); ok {
		conn.Pinner = pinner
	}

	return conn
}
