package main

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"os"
)

// readBlockAt reads exactly min(len(buf), totalSize-offset) bytes at offset,
// treating a fully-satisfied read as success even if the underlying
// implementation also reports io.EOF for it.
func readBlockAt(f *os.File, buf []byte, offset, totalSize int64) (int, error) {
	remaining := totalSize - offset
	if remaining <= 0 {
		return 0, io.EOF
	}
	want := int64(len(buf))
	if remaining < want {
		want = remaining
	}
	n, err := f.ReadAt(buf[:want], offset)
	if err != nil && !(err == io.EOF && int64(n) == want) {
		return n, err
	}
	return n, nil
}

// sourceLoopParams drives the single read -> hash -> compare -> maybe-send
// pass over a source path. It is used both by the standalone
// "source-stream" role (writing to its own stdout) and by a source control
// agent's push driver (writing into a spawned ssh subprocess's stdin).
type sourceLoopParams struct {
	File      *os.File // already opened for reading by the caller
	Size      int64
	BlockSize int64
	Hashes    [][32]byte // previous known hashes, len == block count
	Out       *FrameWriter
	OnSkip    func(index uint64)
	OnSend    func(index uint64, hash [32]byte) // called after the DATA frame is written
}

func runSourceLoop(p sourceLoopParams) error {
	f := p.File
	buf := make([]byte, p.BlockSize)
	blockCount := uint64(len(p.Hashes))
	for index := uint64(0); index < blockCount; index++ {
		offset := int64(index) * p.BlockSize
		n, err := readBlockAt(f, buf, offset, p.Size)
		if err != nil {
			return fmt.Errorf("read %s at block %d: %w", f.Name(), index, err)
		}
		hash := sha256.Sum256(buf[:n])
		if hash == p.Hashes[index] {
			if p.OnSkip != nil {
				p.OnSkip(index)
			}
			continue
		}
		if err := p.Out.WriteFrame(frameData, encodeDataFrame(index, hash, buf[:n])); err != nil {
			return fmt.Errorf("send block %d: %w", index, err)
		}
		if p.OnSend != nil {
			p.OnSend(index, hash)
		}
	}
	return p.Out.WriteFrame(frameDone, nil)
}

// destLoopParams drives the receive -> verify -> pwrite pass on the
// destination path. It is used both by the standalone "sink" role (reading
// from its own stdin, and required to ack/err back over AckOut) and by a
// dest control agent's pull driver (reading from a spawned ssh subprocess's
// stdout; no AckOut needed since the write is confirmed locally, in the
// same process, before OnWritten is called).
type destLoopParams struct {
	File      *os.File // already opened for read/write by the caller
	BlockSize int64
	In        *FrameReader
	AckOut    *FrameWriter // optional
	OnWritten func(index uint64, hash [32]byte)
	// OnCtrlMsg handles an interleaved frameCtrlJSON frame (used only in
	// pull mode, where the remote source-stream has no other channel back
	// to the manager for progress updates). Sink-role callers leave it nil.
	OnCtrlMsg func(CtrlMsg)
}

func runDestLoop(p destLoopParams) error {
	f := p.File
	for {
		typ, payload, err := p.In.ReadFrame()
		if err != nil {
			return fmt.Errorf("read frame: %w", err)
		}
		switch typ {
		case frameDone:
			return f.Sync()
		case frameData:
			index, hash, block, err := decodeDataFrame(payload)
			if err != nil {
				return err
			}
			if sha256.Sum256(block) != hash {
				if p.AckOut != nil {
					_ = p.AckOut.WriteFrame(frameErr, encodeErrFrame(index, "hash mismatch after transfer"))
				}
				return fmt.Errorf("block %d: hash mismatch after transfer", index)
			}
			if _, err := f.WriteAt(block, int64(index)*p.BlockSize); err != nil {
				if p.AckOut != nil {
					_ = p.AckOut.WriteFrame(frameErr, encodeErrFrame(index, err.Error()))
				}
				return fmt.Errorf("write block %d: %w", index, err)
			}
			if p.AckOut != nil {
				if err := p.AckOut.WriteFrame(frameAck, encodeIndexFrame(index)); err != nil {
					return err
				}
			}
			if p.OnWritten != nil {
				p.OnWritten(index, hash)
			}
		case frameCtrlJSON:
			if p.OnCtrlMsg != nil {
				var m CtrlMsg
				if err := json.Unmarshal(payload, &m); err == nil {
					p.OnCtrlMsg(m)
				}
			}
		default:
			return fmt.Errorf("unexpected frame type %d on data channel", typ)
		}
	}
}