package protocol

import (
	"encoding/binary"
	"log"
)

func init() {
	registerPostgreSQL()
}

func registerPostgreSQL() {
	Register(Server, "R\x00???", detectPostgreSQLServer)             // Authentication request
	Register(Server, "K\x00???", detectPostgreSQLServer)             // BackendKeyData
	Register(Server, "S\x00???", detectPostgreSQLServer)             // ParameterStatus
	Register(Server, "Z\x00???", detectPostgreSQLServer)             // ReadyForQuery
	Register(Server, "E\x00???", detectPostgreSQLServer)             // ErrorResponse
	Register(Server, "N\x00???", detectPostgreSQLServer)             // NoticeResponse
	Register(Client, "????\x00\x02\x00\x00", detectPostgreSQLClient) // Startup packet, protocol 2.0
	Register(Client, "????\x00\x03\x00\x00", detectPostgreSQLClient) // Startup packet, protocol 3.0
}

func detectPostgreSQLClient(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
	// A client startup message needs at least 8 bytes (length + protocol version).
	if len(data) < 8 {
		return nil, 0
	}

	length := int(binary.BigEndian.Uint32(data[0:]))
	if len(data) != length {
		log.Printf("not postgres %q: %d != %d", data, len(data), length)
		return nil, 0
	}

	if dstPort == 5432 {
		confidence = .1
	}

	major := int(binary.BigEndian.Uint16(data[4:]))
	minor := int(binary.BigEndian.Uint16(data[6:]))
	if major == 2 || major == 3 {
		return &Protocol{
			Name: ProtocolPostgreSQL,
			Version: Version{
				Major: major,
				Minor: minor,
				Patch: -1,
			},
		}, confidence + .75
	}
	return nil, 0
}

func detectPostgreSQLServer(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
	// A server message needs at least 5 bytes (type + length).
	if len(data) < 5 {
		return nil, 0
	}

	if srcPort == 5432 {
		confidence = .1
	}

	// All server messages (and subsequent client messages) are tagged with a single-byte type.
	firstByte := data[0]
	switch firstByte {
	case 'R', // Authentication request
		'K', // BackendKeyData
		'S', // ParameterStatus
		'Z', // ReadyForQuery
		'E', // ErrorResponse
		'N': // NoticeResponse
		return &Protocol{Name: ProtocolPostgreSQL}, confidence + .65

	default:
		return nil, 0
	}
}