package probe

import (
	"context"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"net"
	"strings"
	"time"

	"tea.karlcom.de/karl/bucketeer/internal/model"
)

// Deep, unauthenticated, read-only protocol probes. They determine whether a
// database endpoint actually answers a read without credentials. No password,
// token or credential of any kind is ever sent, and no state is mutated.

func (p *Prober) deepDB(ctx context.Context, c model.Candidate) (Result, bool) {
	switch c.Service {
	case model.ServiceRedis, model.ServiceElastiCache, model.ServiceMemoryDB:
		return p.probeRedis(ctx, c.Hostname, dbPort(c.Service))
	case model.ServiceRDS, model.ServicePostgres, model.ServiceRedshift:
		return p.probePostgres(ctx, c.Hostname, 5432)
	case model.ServiceDocumentDB, model.ServiceCosmosDB:
		return p.probeMongo(ctx, c.Hostname, 27017)
	case model.ServiceOpenSearch:
		return p.probeHTTPReadable(ctx, "https://"+c.Hostname+"/_cat/indices?format=json&h=index")
	case model.ServiceNeptune:
		return p.probeHTTPReadable(ctx, "https://"+c.Hostname+"/status")
	}
	return Result{}, false
}

func dbPort(s model.Service) int {
	switch s {
	case model.ServiceMySQL:
		return 3306
	case model.ServicePostgres, model.ServiceRDS:
		return 5432
	case model.ServiceRedis, model.ServiceElastiCache, model.ServiceMemoryDB:
		return 6379
	case model.ServiceDocumentDB, model.ServiceCosmosDB:
		return 27017
	case model.ServiceRedshift:
		return 5439
	case model.ServiceNeptune:
		return 8182
	case model.ServiceDatabase:
		return 1433
	default:
		return 5432
	}
}

func (p *Prober) dial(ctx context.Context, host string, port int) (net.Conn, error) {
	d := net.Dialer{Timeout: p.cfg.Timeout}
	return d.DialContext(ctx, "tcp", net.JoinHostPort(host, fmt.Sprint(port)))
}

func (p *Prober) dialTLS(ctx context.Context, host string, port int) (net.Conn, error) {
	conn, err := p.dial(ctx, host, port)
	if err != nil {
		return nil, err
	}
	tconn := tls.Client(conn, &tls.Config{InsecureSkipVerify: true, ServerName: host}) // observation only
	_ = tconn.SetDeadline(time.Now().Add(p.cfg.Timeout))
	if err := tconn.HandshakeContext(ctx); err != nil {
		conn.Close()
		return nil, err
	}
	return tconn, nil
}

// probeRedis sends PING. Redis answers +PONG when no authentication is required.
func (p *Prober) probeRedis(ctx context.Context, host string, port int) (Result, bool) {
	var lastErr error
	for _, tlsOn := range []bool{false, true} {
		var (
			conn net.Conn
			err  error
		)
		if tlsOn {
			conn, err = p.dialTLS(ctx, host, port)
		} else {
			conn, err = p.dial(ctx, host, port)
		}
		if err != nil {
			lastErr = err
			continue
		}
		_ = conn.SetDeadline(time.Now().Add(p.cfg.Timeout))
		_, _ = conn.Write([]byte("PING\r\n"))
		buf := make([]byte, 256)
		n, _ := conn.Read(buf)
		conn.Close()
		resp := string(buf[:n])
		switch {
		case strings.HasPrefix(resp, "+PONG"):
			return listRead("redis PING answered without auth"), true
		case strings.Contains(resp, "NOAUTH"), strings.HasPrefix(resp, "-ERR"), strings.Contains(resp, "DENIED"), strings.Contains(resp, "protected mode"):
			return wrap(model.ExposureAuthRequired, "redis requires authentication"), true
		}
		// no/invalid response: try the other transport
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("redis: no response")
	}
	st, detail := classifyError(lastErr)
	return wrap(st, detail), true
}

// probePostgres negotiates TLS if offered, then sends a StartupMessage with no
// password. The server's AuthenticationRequest tells us whether it trusts the
// connection (AuthenticationOk) or requires a password.
func (p *Prober) probePostgres(ctx context.Context, host string, port int) (Result, bool) {
	conn, err := p.dial(ctx, host, port)
	if err != nil {
		return wrap(classifyError(err)), true
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(p.cfg.Timeout))

	// SSLRequest (length 8, code 80877103).
	req := []byte{0, 0, 0, 8, 4, 210, 22, 47}
	if _, err := conn.Write(req); err == nil {
		one := make([]byte, 1)
		if _, err := io.ReadFull(conn, one); err == nil {
			switch one[0] {
			case 'S': // proceed with TLS
				tconn := tls.Client(conn, &tls.Config{InsecureSkipVerify: true, ServerName: host})
				if err := tconn.HandshakeContext(ctx); err == nil {
					conn = tconn
				}
			case 'N': // plaintext only
			case 'E': // server refused; fall through and try startup anyway
			}
		}
	}

	// StartupMessage: user=postgres database=postgres
	payload := []byte{0, 3, 0, 0}
	payload = append(payload, "user\x00postgres\x00database\x00postgres\x00\x00"...)
	msg := make([]byte, 4+len(payload))
	binary.BigEndian.PutUint32(msg, uint32(4+len(payload)))
	copy(msg[4:], payload)
	if _, err := conn.Write(msg); err != nil {
		return wrap(classifyError(err)), true
	}

	head := make([]byte, 5)
	if _, err := io.ReadFull(conn, head); err != nil {
		return wrap(classifyError(err)), true
	}
	kind := head[0]
	length := int(binary.BigEndian.Uint32(head[1:5]))
	if length > 4 {
		body := make([]byte, length-4)
		_, _ = io.ReadFull(conn, body)
		if kind == 'R' && len(body) >= 4 {
			authType := binary.BigEndian.Uint32(body[:4])
			if authType == 0 {
				return listRead("postgres AuthenticationOk (trust, no password required)"), true
			}
			return wrap(model.ExposureAuthRequired, fmt.Sprintf("postgres requires password (auth type %d)", authType)), true
		}
		if kind == 'E' {
			return wrap(model.ExposureAuthRequired, "postgres rejected startup (pg_hba/host not permitted)"), true
		}
	}
	return wrap(model.ExposureReachable, "postgres port open; auth state unclear"), true
}

// probeMongo performs the OP_MSG handshake then a read-only listDatabases. A
// successful, non-empty reply means the deployment has no authentication.
func (p *Prober) probeMongo(ctx context.Context, host string, port int) (Result, bool) {
	var (
		conn net.Conn
		err  error
	)
	// Managed MongoDB (DocumentDB, Cosmos) requires TLS; try it, then plain.
	conn, err = p.dialTLS(ctx, host, port)
	if err != nil {
		conn, err = p.dial(ctx, host, port)
		if err != nil {
			return wrap(classifyError(err)), true
		}
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(p.cfg.Timeout))

	doc := bsonDoc{"listDatabases": int32(1), "$db": "admin"}
	if err := writeOPMsg(conn, doc); err != nil {
		return wrap(classifyError(err)), true
	}
	raw, err := readOPMsg(conn)
	if err != nil {
		return wrap(classifyError(err)), true
	}
	ok, errmsg := parseBSONFields(raw)
	switch {
	case ok == 1 && errmsg == "":
		return listRead("mongodb listDatabases succeeded without auth"), true
	case strings.Contains(strings.ToLower(errmsg), "auth") || strings.Contains(strings.ToLower(errmsg), "unauthorized"):
		return wrap(model.ExposureAuthRequired, "mongodb requires authentication"), true
	default:
		return wrap(model.ExposureReachable, "mongodb replied but auth state unclear"), true
	}
}

func (p *Prober) probeHTTPReadable(ctx context.Context, endpoint string) (Result, bool) {
	status, _, err := p.do(ctx, endpoint)
	if err != nil {
		return wrap(classifyError(err)), true
	}
	return wrap(classifyStatus(status, "")), true
}

// --- minimal BSON/OP_MSG helpers (read-only) ---

type bsonDoc map[string]any

func appendCString(b []byte, s string) []byte {
	b = append(b, s...)
	return append(b, 0)
}

func encodeBSON(doc bsonDoc) []byte {
	var body []byte
	for k, v := range doc {
		switch val := v.(type) {
		case int32:
			body = append(body, 0x10)
			body = appendCString(body, k)
			var tmp [4]byte
			binary.LittleEndian.PutUint32(tmp[:], uint32(val))
			body = append(body, tmp[:]...)
		case string:
			body = append(body, 0x02)
			body = appendCString(body, k)
			var tmp [4]byte
			binary.LittleEndian.PutUint32(tmp[:], uint32(len(val)+1))
			body = append(body, tmp[:]...)
			body = appendCString(body, val)
		}
	}
	out := make([]byte, 4, 4+len(body)+1)
	binary.LittleEndian.PutUint32(out, uint32(4+len(body)+1))
	out = append(out, body...)
	out = append(out, 0)
	return out
}

func writeOPMsg(w io.Writer, doc bsonDoc) error {
	bson := encodeBSON(doc)
	msg := make([]byte, 0, 16+4+1+len(bson))
	var hdr [16]byte
	binary.LittleEndian.PutUint32(hdr[0:], uint32(16+4+1+len(bson))) // messageLength
	binary.LittleEndian.PutUint32(hdr[4:], 1)                        // requestID
	binary.LittleEndian.PutUint32(hdr[8:], 0)                        // responseTo
	binary.LittleEndian.PutUint32(hdr[12:], 2013)                    // opCode OP_MSG
	msg = append(msg, hdr[:]...)
	msg = append(msg, 0, 0, 0, 0) // flagBits
	msg = append(msg, 0)          // section kind 0
	msg = append(msg, bson...)
	_, err := w.Write(msg)
	return err
}

func readOPMsg(r io.Reader) ([]byte, error) {
	hdr := make([]byte, 16)
	if _, err := io.ReadFull(r, hdr); err != nil {
		return nil, err
	}
	length := int(binary.LittleEndian.Uint32(hdr[0:]))
	if length < 21 || length > 8<<20 {
		return nil, fmt.Errorf("mongo: bad message length %d", length)
	}
	body := make([]byte, length-16)
	if _, err := io.ReadFull(r, body); err != nil {
		return nil, err
	}
	return body, nil
}

// parseBSONFields extracts the top-level "ok" (double/int) and "errmsg" fields.
func parseBSONFields(b []byte) (float64, string) {
	if len(b) < 5 {
		return 0, ""
	}
	// body: flagBits(4) sectionKind(1) [doc] or [kind 1 + size + id + doc]
	p := 4
	if p < len(b) && b[p] == 1 && p+5 <= len(b) {
		p++ // section kind 1
		sz := int(binary.LittleEndian.Uint32(b[p:]))
		p += 4 + sz // skip identifier
	} else if p < len(b) && b[p] == 0 {
		p++
	}
	if p >= len(b) {
		return 0, ""
	}
	docLen := int(int32(binary.LittleEndian.Uint32(b[p:])))
	if p+docLen > len(b) {
		docLen = len(b) - p
	}
	doc := b[p : p+docLen]
	var ok float64
	var errmsg string
	i := 4
	for i < len(doc)-1 {
		typ := doc[i]
		i++
		start := i
		for i < len(doc) && doc[i] != 0 {
			i++
		}
		name := string(doc[start:i])
		i++ // skip null
		switch typ {
		case 0x01: // double
			if i+8 <= len(doc) {
				bits := binary.LittleEndian.Uint64(doc[i:])
				if name == "ok" {
					ok = math.Float64frombits(bits)
				}
			}
			i += 8
		case 0x10: // int32
			if i+4 <= len(doc) {
				v := int32(binary.LittleEndian.Uint32(doc[i:]))
				if name == "ok" {
					ok = float64(v)
				}
			}
			i += 4
		case 0x02: // string
			if i+4 <= len(doc) {
				n := int(int32(binary.LittleEndian.Uint32(doc[i:])))
				i += 4
				if n > 0 && i+n <= len(doc) {
					if name == "errmsg" {
						errmsg = string(doc[i : i+n-1])
					}
					i += n
				}
			}
		case 0x07: // objectId
			i += 12
		case 0x08: // bool
			i++
		case 0x03, 0x04: // document/array
			if i+4 <= len(doc) {
				n := int(int32(binary.LittleEndian.Uint32(doc[i:])))
				i += n
			}
		case 0x09: // datetime
			i += 8
		case 0x0A: // null
		case 0x11: // timestamp
			i += 8
		default:
			// Unknown type: abandon parsing to avoid mis-stepping.
			return ok, errmsg
		}
	}
	return ok, errmsg
}