//go:build linux || freebsd
// +build linux freebsd

package sftp

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"testing"

	"codeberg.org/petrbalvin/goget/internal/core"
	"golang.org/x/crypto/ssh"
)

// startTestServer starts a minimal SSH+SFTP server for testing.
func startTestServer(t *testing.T) (addr string, cleanup func()) {
	t.Helper()

	config := &ssh.ServerConfig{
		NoClientAuth: true,
	}
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	signer, err := ssh.NewSignerFromKey(privateKey)
	if err != nil {
		t.Fatalf("new signer: %v", err)
	}
	config.AddHostKey(signer)

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	done := make(chan struct{})
	go func() {
		defer close(done)
		// Accept connections in a loop so the parallel-download tests
		// (which open N SFTP sessions) can run. Each connection is
		// served in its own goroutine, so the N pool connections can be
		// alive concurrently as long as the test goroutine keeps
		// accepting.
		for {
			conn, err := listener.Accept()
			if err != nil {
				return
			}

			go func(conn net.Conn) {
				defer conn.Close()
				sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
				if err != nil {
					return
				}
				defer sshConn.Close()
				go ssh.DiscardRequests(reqs)

				for newChannel := range chans {
					if newChannel.ChannelType() != "session" {
						newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
						continue
					}
					channel, requests, err := newChannel.Accept()
					if err != nil {
						return
					}
					go func(ch ssh.Channel, reqs <-chan *ssh.Request) {
						for req := range reqs {
							ok := false
							if req.Type == "subsystem" && len(req.Payload) > 4 {
								l := binary.BigEndian.Uint32(req.Payload[:4])
								name := string(req.Payload[4 : 4+l])
								if name == "sftp" {
									ok = true
									if req.WantReply {
										req.Reply(ok, nil)
									}
									handleSFTPSession(ch)
									return
								}
							}
							if req.WantReply {
								req.Reply(ok, nil)
							}
						}
						ch.Close()
					}(channel, requests)
				}
			}(conn)
		}
	}()

	return listener.Addr().String(), func() {
		listener.Close()
		<-done
	}
}

// handleSFTPSession handles a minimal SFTP v3 session.
func handleSFTPSession(ch ssh.Channel) {
	defer ch.Close()

	// Expect INIT
	pkt, err := readTestPacket(ch)
	if err != nil || len(pkt) < 5 || pkt[0] != sshFxpInit {
		return
	}
	// Respond VERSION
	var resp bytes.Buffer
	resp.WriteByte(sshFxpVersion)
	writeUint32(&resp, 3)
	writeTestPacket(ch, resp.Bytes())

	handles := make(map[string]string)
	dirHandles := make(map[string][]string)
	dirHandlePaths := make(map[string]string)
	handleCounter := 0

	for {
		pkt, err := readTestPacket(ch)
		if err != nil {
			return
		}
		if len(pkt) < 1 {
			continue
		}
		reqID, _ := readUint32(pkt, 1)

		switch pkt[0] {
		case sshFxpRealpath:
			path, _ := readString(pkt, 5)
			if path == "." {
				path = "/home/test"
			}
			var r bytes.Buffer
			r.WriteByte(sshFxpName)
			writeUint32(&r, reqID)
			writeUint32(&r, 1)
			writeString(&r, path)
			writeString(&r, path)
			writeUint32(&r, 0) // attrs
			writeTestPacket(ch, r.Bytes())

		case sshFxpStat:
			path, _ := readString(pkt, 5)
			var r bytes.Buffer
			r.WriteByte(sshFxpAttrs)
			writeUint32(&r, reqID)
			fs := testFSLookup(path)
			if fs != nil && fs.isDir {
				// mode flag only
				writeUint32(&r, 0x4)
				writeUint32(&r, 0x4000) // S_IFDIR
				writeTestPacket(ch, r.Bytes())
				continue
			}
			if fs != nil && !fs.isDir {
				// size flag only
				writeUint32(&r, 0x1)
				writeUint64(&r, uint64(len(fs.content)))
				writeTestPacket(ch, r.Bytes())
				continue
			}
			// legacy single-file mock behaviour
			writeUint32(&r, 0x1) // size flag
			size := uint64(len(testFileContent))
			if testOverrideAttrSize != nil {
				size = *testOverrideAttrSize
			}
			writeUint64(&r, size)
			writeTestPacket(ch, r.Bytes())

		case sshFxpOpendir:
			path, _ := readString(pkt, 5)
			fs := testFSLookup(path)
			if fs == nil || !fs.isDir {
				status(ch, reqID, sshFxNoSuchFile, "no such directory")
				continue
			}
			handleCounter++
			h := fmt.Sprintf("dhandle%d", handleCounter)
			dirHandles[h] = fs.children
			dirHandlePaths[h] = path
			var r bytes.Buffer
			r.WriteByte(sshFxpHandle)
			writeUint32(&r, reqID)
			writeString(&r, h)
			writeTestPacket(ch, r.Bytes())

		case sshFxpReaddir:
			handle, _ := readString(pkt, 5)
			children, ok := dirHandles[handle]
			if !ok {
				status(ch, reqID, sshFxFailure, "no such dir handle")
				continue
			}
			parentPath := dirHandlePaths[handle]
			var r bytes.Buffer
			r.WriteByte(sshFxpName)
			writeUint32(&r, reqID)
			writeUint32(&r, uint32(len(children)))
			for _, name := range children {
				writeString(&r, name)
				writeString(&r, name)
				childPath := parentPath + "/" + name
				childFS := testFSLookup(childPath)
				isDir := childFS != nil && childFS.isDir
				// attrs: size flag (0x1) and mode flag (0x4) so the
				// client can both classify the entry (dir vs file) and
				// read the file size for parallel accounting.
				if isDir {
					writeUint32(&r, 0x4)
					writeUint32(&r, 0x4000) // S_IFDIR
				} else {
					writeUint32(&r, 0x5) // size + mode
					var sz uint64
					if childFS != nil {
						sz = uint64(len(childFS.content))
					} else {
						sz = uint64(len(testFileContent))
					}
					writeUint64(&r, sz)
					writeUint32(&r, 0x8000) // S_IFREG
				}
			}
			writeTestPacket(ch, r.Bytes())

		case sshFxpOpen:
			path, off := readString(pkt, 5)
			pflags, _ := readUint32(pkt, off)
			handleCounter++
			h := fmt.Sprintf("handle%d", handleCounter)
			handles[h] = path
			var r bytes.Buffer
			if pflags&sshFxfRead != 0 {
				r.WriteByte(sshFxpHandle)
				writeUint32(&r, reqID)
				writeString(&r, h)
				writeTestPacket(ch, r.Bytes())
			} else if pflags&sshFxfWrite != 0 {
				r.WriteByte(sshFxpHandle)
				writeUint32(&r, reqID)
				writeString(&r, h)
				writeTestPacket(ch, r.Bytes())
			} else {
				status(ch, reqID, sshFxFailure, "unsupported flags")
			}

		case sshFxpRead:
			handle, off := readString(pkt, 5)
			offset, off := readUint64(pkt, off)
			length, _ := readUint32(pkt, off)
			_ = length
			path := handles[handle]
			if path == "" {
				status(ch, reqID, sshFxNoSuchFile, "no such handle")
				continue
			}
			var data []byte
			if fs := testFSLookup(path); fs != nil && !fs.isDir {
				data = fs.content
			} else {
				data = testFileContent
			}
			if int(offset) >= len(data) {
				status(ch, reqID, sshFxEOF, "")
				continue
			}
			end := int(offset) + int(length)
			if end > len(data) {
				end = len(data)
			}
			chunk := data[offset:end]
			var r bytes.Buffer
			r.WriteByte(sshFxpData)
			writeUint32(&r, reqID)
			writeUint32(&r, uint32(len(chunk)))
			r.Write(chunk)
			writeTestPacket(ch, r.Bytes())

		case sshFxpClose:
			handle, _ := readString(pkt, 5)
			delete(handles, handle)
			delete(dirHandles, handle)
			var r bytes.Buffer
			r.WriteByte(sshFxpStatus)
			writeUint32(&r, reqID)
			writeUint32(&r, sshFxOk)
			writeString(&r, "")
			writeString(&r, "")
			writeTestPacket(ch, r.Bytes())

		default:
			status(ch, reqID, sshFxOpUnsupported, "unsupported")
		}
	}
}

// testFSLookup is a tiny helper that hides the nil-map check so the
// mock switch arms above can stay readable.
func testFSLookup(path string) *testFSEntry {
	if testFS == nil {
		return nil
	}
	return testFS[path]
}

func status(ch ssh.Channel, reqID, code uint32, msg string) {
	var r bytes.Buffer
	r.WriteByte(sshFxpStatus)
	writeUint32(&r, reqID)
	writeUint32(&r, code)
	writeString(&r, msg)
	writeString(&r, "")
	writeTestPacket(ch, r.Bytes())
}

func readTestPacket(r io.Reader) ([]byte, error) {
	var lenBuf [4]byte
	if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(lenBuf[:])
	if length > 256*1024 {
		return nil, fmt.Errorf("packet too large")
	}
	pkt := make([]byte, length)
	if _, err := io.ReadFull(r, pkt); err != nil {
		return nil, err
	}
	return pkt, nil
}

func writeTestPacket(w io.Writer, pkt []byte) error {
	var lenBuf [4]byte
	binary.BigEndian.PutUint32(lenBuf[:], uint32(len(pkt)))
	if _, err := w.Write(lenBuf[:]); err != nil {
		return err
	}
	_, err := w.Write(pkt)
	return err
}

var testFileContent = []byte("hello sftp world")

// testOverrideAttrSize, when non-nil, replaces the value reported by the
// mock server's sshFxpStat handler. Tests use it to simulate a server
// that returns a wrong (typically too large) size attribute so we can
// verify the client does not blindly trust attr.Size for
// BytesDownloaded accounting.
var testOverrideAttrSize *uint64

// testFSEntry describes a single in-memory file or directory used by
// the recursive-download tests. isDir==true entries carry a children
// slice; isDir==false entries carry the file content.
type testFSEntry struct {
	isDir    bool
	content  []byte
	children []string
}

// testFS is an optional in-memory filesystem consulted by the mock
// server before falling back to the default testFileContent behaviour.
// Tests that exercise recursive downloads set it up and clear it via
// t.Cleanup. Keys are absolute paths; paths NOT present in the map
// keep the legacy single-file mock semantics so existing
// non-recursive tests continue to pass.
var testFS map[string]*testFSEntry

func TestDownload(t *testing.T) {
	addr, cleanup := startTestServer(t)
	defer cleanup()

	u, _ := url.Parse("sftp://user@" + addr + "/testfile")
	outDir := t.TempDir()
	outputPath := filepath.Join(outDir, "out")

	req := &core.DownloadRequest{
		URL:         u,
		Output:      outputPath,
		Ctx:         context.Background(),
		SSHInsecure: true,
	}

	result, err := download(context.Background(), req)
	if err != nil {
		t.Fatalf("download failed: %v", err)
	}

	data, err := os.ReadFile(outputPath)
	if err != nil {
		t.Fatalf("read output: %v", err)
	}
	if !bytes.Equal(data, testFileContent) {
		t.Errorf("got %q, want %q", data, testFileContent)
	}
	if result.BytesDownloaded != int64(len(testFileContent)) {
		t.Errorf("bytes=%d, want %d", result.BytesDownloaded, len(testFileContent))
	}
	if result.Protocol != "sftp" {
		t.Errorf("protocol=%q, want sftp", result.Protocol)
	}
}

// TestDownloadWithStaleAttrSize is a regression guard for the BACKLOG
// entry "SFTP BytesDownloaded unconditionally overwritten". The previous
// implementation set result.BytesDownloaded = int64(attr.Size) when
// attr.Size > 0, ignoring the locally-tracked `total` counter. If the
// server reported a stale (typically too large) size, the result would
// over-report the bytes actually transferred — which is exactly what
// this test exercises.
func TestDownloadWithStaleAttrSize(t *testing.T) {
	// Simulate a server that reports 10 KiB even though the file is
	// only len(testFileContent) bytes long. After the fix, the result
	// must reflect the actual bytes received, not the stale attribute.
	wrong := uint64(10 * 1024)
	testOverrideAttrSize = &wrong
	t.Cleanup(func() { testOverrideAttrSize = nil })

	addr, cleanup := startTestServer(t)
	defer cleanup()

	u, _ := url.Parse("sftp://user@" + addr + "/testfile")
	outDir := t.TempDir()
	outputPath := filepath.Join(outDir, "out")

	req := &core.DownloadRequest{
		URL:         u,
		Output:      outputPath,
		Ctx:         context.Background(),
		SSHInsecure: true,
	}

	result, err := download(context.Background(), req)
	if err != nil {
		t.Fatalf("download failed: %v", err)
	}

	// The downloaded file must still hold the real (small) payload.
	data, err := os.ReadFile(outputPath)
	if err != nil {
		t.Fatalf("read output: %v", err)
	}
	if !bytes.Equal(data, testFileContent) {
		t.Errorf("got %q, want %q", data, testFileContent)
	}

	// And BytesDownloaded must reflect the actual bytes received, not
	// the stale attr.Size the server returned.
	if result.BytesDownloaded != int64(len(testFileContent)) {
		t.Errorf("BytesDownloaded = %d, want %d (real payload size, not stale attr.Size=%d)",
			result.BytesDownloaded, len(testFileContent), wrong)
	}
}

func TestParseHeader(t *testing.T) {
	var buf bytes.Buffer
	writeUint32(&buf, 0x05) // size + mode flags
	writeUint64(&buf, 60)
	writeUint32(&buf, 0644)

	attr, off := readAttrs(buf.Bytes(), 0)
	if attr.Size != 60 {
		t.Errorf("size=%d, want 60", attr.Size)
	}
	if off != 16 {
		t.Errorf("off=%d, want 16", off)
	}
}

func TestIsEOF(t *testing.T) {
	if !isEOF(fmt.Errorf("sftp status %d", sshFxEOF)) {
		t.Error("expected true")
	}
	if isEOF(fmt.Errorf("sftp status %d", sshFxFailure)) {
		t.Error("expected false")
	}
}

func TestDialTimeout(t *testing.T) {
	u, _ := url.Parse("sftp://user@127.0.0.1:1/")
	_, err := dial(context.Background(), u, "", true)
	if err == nil {
		t.Fatal("expected error")
	}
}

func TestSFTPRecursiveDownloadParallel(t *testing.T) {
	addr, cleanup := startTestServer(t)
	defer cleanup()

	setupRecursiveFS := func(t *testing.T) {
		t.Helper()
		testFS = map[string]*testFSEntry{
			"/home/test/dir":           {isDir: true, children: []string{"a.txt", "b.txt", "c.txt", "sub"}},
			"/home/test/dir/a.txt":     {isDir: false, content: []byte("alpha")},
			"/home/test/dir/b.txt":     {isDir: false, content: []byte("bravo-bravo")},
			"/home/test/dir/c.txt":     {isDir: false, content: []byte("charlie")},
			"/home/test/dir/sub":       {isDir: true, children: []string{"x.txt", "y.txt"}},
			"/home/test/dir/sub/x.txt": {isDir: false, content: []byte("x-ray")},
			"/home/test/dir/sub/y.txt": {isDir: false, content: []byte("yankee")},
		}
		t.Cleanup(func() { testFS = nil })
	}

	t.Run("parallel=3 downloads all files", func(t *testing.T) {
		setupRecursiveFS(t)
		u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
		outDir := t.TempDir()
		req := &core.DownloadRequest{
			URL:               u,
			Output:            outDir,
			Recursive:         true,
			RecursiveParallel: 3,
			SSHInsecure:       true,
		}
		result, err := download(context.Background(), req)
		if err != nil {
			t.Fatalf("recursive parallel download failed: %v", err)
		}
		if result.ChunksCount != 5 {
			t.Errorf("ChunksCount = %d, want 5 (3 top-level files + 2 in subdir)", result.ChunksCount)
		}
		// 5 + 11 + 7 + 5 + 6 = 34 bytes total content
		const wantBytes = int64(5 + 11 + 7 + 5 + 6)
		if result.BytesDownloaded != wantBytes {
			t.Errorf("BytesDownloaded = %d, want %d", result.BytesDownloaded, wantBytes)
		}
		// Verify all 5 files actually exist with correct content.
		wantFiles := map[string]string{
			"a.txt":     "alpha",
			"b.txt":     "bravo-bravo",
			"c.txt":     "charlie",
			"sub/x.txt": "x-ray",
			"sub/y.txt": "yankee",
		}
		for name, want := range wantFiles {
			got, err := os.ReadFile(filepath.Join(outDir, name))
			if err != nil {
				t.Errorf("file %s: %v", name, err)
				continue
			}
			if string(got) != want {
				t.Errorf("file %s = %q, want %q", name, string(got), want)
			}
		}
	})

	t.Run("parallel=0 falls back to sequential", func(t *testing.T) {
		setupRecursiveFS(t)
		u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
		outDir := t.TempDir()
		req := &core.DownloadRequest{
			URL:               u,
			Output:            outDir,
			Recursive:         true,
			RecursiveParallel: 0,
			SSHInsecure:       true,
		}
		result, err := download(context.Background(), req)
		if err != nil {
			t.Fatalf("recursive sequential download failed: %v", err)
		}
		if result.ChunksCount != 5 {
			t.Errorf("ChunksCount = %d, want 5", result.ChunksCount)
		}
	})

	t.Run("parallel=1 falls back to sequential", func(t *testing.T) {
		setupRecursiveFS(t)
		u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
		outDir := t.TempDir()
		req := &core.DownloadRequest{
			URL:               u,
			Output:            outDir,
			Recursive:         true,
			RecursiveParallel: 1,
			SSHInsecure:       true,
		}
		result, err := download(context.Background(), req)
		if err != nil {
			t.Fatalf("download with parallel=1 failed: %v", err)
		}
		if result.ChunksCount != 5 {
			t.Errorf("ChunksCount = %d, want 5", result.ChunksCount)
		}
	})

	t.Run("subdir propagation with parallel", func(t *testing.T) {
		setupRecursiveFS(t)
		u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
		outDir := t.TempDir()
		req := &core.DownloadRequest{
			URL:               u,
			Output:            outDir,
			Recursive:         true,
			RecursiveParallel: 4,
			SSHInsecure:       true,
		}
		_, err := download(context.Background(), req)
		if err != nil {
			t.Fatalf("download failed: %v", err)
		}
		// Both subdir files should be present and correct.
		for _, name := range []string{"sub/x.txt", "sub/y.txt"} {
			path := filepath.Join(outDir, name)
			if _, err := os.Stat(path); err != nil {
				t.Errorf("expected file %s to exist: %v", name, err)
			}
		}
	})
}