package main

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

// readBlockAt reads exactly min(len(buf), totalSize-offset) logical bytes at
// offset, treating a fully-satisfied read as success even if the underlying
// implementation also reports io.EOF for it. When align > 1 (a raw disk
// handle that only accepts sector-aligned I/O) the physical read is rounded
// up to the next multiple of align, but the returned count is still the
// logical size — the caller only ever looks at buf[:n].
func readBlockAt(f *os.File, buf []byte, offset, totalSize, align int64) (int, error) {
	remaining := totalSize - offset
	if remaining <= 0 {
		return 0, io.EOF
	}
	want := int64(len(buf))
	if remaining < want {
		want = remaining
	}
	readLen := want
	if align > 1 && want%align != 0 {
		readLen = roundUp(want, align)
		if readLen > int64(len(buf)) {
			readLen = int64(len(buf))
		}
	}
	n, err := f.ReadAt(buf[:readLen], offset)
	if err != nil && !(err == io.EOF && int64(n) >= want) {
		return n, err
	}
	if int64(n) > want {
		n = int(want)
	}
	return n, nil
}

// writeBlockAt writes block at offset. When align > 1 and the block is not a
// whole number of sectors (only ever the final block of a device sync), the
// enclosing aligned span is read first and the block overlaid onto it, so a
// raw disk handle that rejects sub-sector writes still gets an aligned write
// and the bytes past the sync size are preserved.
func writeBlockAt(f *os.File, block []byte, offset, align int64) error {
	if align <= 1 || int64(len(block))%align == 0 {
		_, err := f.WriteAt(block, offset)
		return err
	}
	padded := roundUp(int64(len(block)), align)
	tmp := make([]byte, padded)
	if _, err := f.ReadAt(tmp, offset); err != nil && err != io.EOF {
		return fmt.Errorf("read-modify-write tail at %d: %w", offset, err)
	}
	copy(tmp, block)
	_, err := f.WriteAt(tmp, offset)
	return err
}

// checkBlockAlign rejects a block size that a raw disk handle's sector
// alignment can't satisfy (every block offset is a multiple of the block
// size, so the block size itself must be a whole number of sectors).
func checkBlockAlign(blockSize, align int64) error {
	if align > 1 && blockSize%align != 0 {
		return fmt.Errorf("block size %d is not a multiple of the device's %d-byte sector size; pass --block-size divisible by %d",
			blockSize, align, align)
	}
	return nil
}

// blockHash pairs a block index with the SHA-256 of the destination's
// current content for that block. The destination side streams these in
// ascending index order as it scans; the source side consumes them in the
// same order, so the two scans overlap instead of running back to back.
type blockHash struct {
	index uint64
	hash  [32]byte
}

// streamHashBlocks reads f in blockSize-byte blocks over the window
// [base, base+size), hashing each one and handing it to onHash the moment it
// is computed (this is what lets the peer start comparing before the whole
// side has been scanned). It keeps no hash state of its own. onProgress, if
// non-nil, is called with (blocksHashed, totalBlocks) before the first block
// and after each one. base is 0 for a whole-file sync; a non-zero base
// fingerprints just a byte window within a larger handle.
func streamHashBlocks(f *os.File, base, size, blockSize, align int64, onHash func(bh blockHash) error, onProgress func(done, total int64)) error {
	if err := checkBlockAlign(blockSize, align); err != nil {
		return err
	}
	blockCount := (size + blockSize - 1) / blockSize
	buf := make([]byte, blockSize)
	if onProgress != nil {
		onProgress(0, blockCount)
	}
	for i := int64(0); i < blockCount; i++ {
		n, err := readBlockAt(f, buf, base+i*blockSize, base+size, align)
		if err != nil {
			return fmt.Errorf("hash %s at block %d: %w", f.Name(), i, err)
		}
		if err := onHash(blockHash{index: uint64(i), hash: sha256.Sum256(buf[:n])}); err != nil {
			return err
		}
		if onProgress != nil {
			onProgress(i+1, blockCount)
		}
	}
	return nil
}

// scanProgressEmitter returns an onProgress callback for streamHashBlocks
// that forwards "scan" phase progress onto out, throttled to ~3/second.
func scanProgressEmitter(out *FrameWriter) func(done, total int64) {
	var last time.Time
	return func(done, total int64) {
		if done != 0 && done != total && time.Since(last) < 350*time.Millisecond {
			return
		}
		last = time.Now()
		_ = out.WriteJSON(CtrlMsg{Type: msgProgress, Phase: "scan", Hashed: done, TotalBlocks: total})
	}
}

// sourceLoopParams drives the single read -> hash -> compare -> maybe-send
// pass over a source path. The hash of each destination block arrives on
// Hashes (in ascending index order); the block read for the comparison is
// the same buffer that gets sent on if it differs — a differing block is
// never read twice.
type sourceLoopParams struct {
	File      *os.File // already opened for reading by the caller
	Base      int64    // byte offset the window starts at (0 for a whole-file sync)
	Size      int64    // window length in bytes
	BlockSize int64
	Align     int64
	Hashes    <-chan blockHash
	Out       *FrameWriter
	// OnRead is called immediately after a block has been read and hashed
	// for comparison — whether it turns out to match or differ — so
	// callers can track true source read throughput independent of the
	// send/ack timing tracked by OnSkip/OnSend below.
	OnRead func(index uint64)
	OnSkip func(index uint64)
	// OnSend is called for a differing block immediately before its DATA
	// frame goes on the wire — never after. The push driver relies on this
	// ordering to record the block as awaiting confirmation before the peer
	// can possibly ACK it (over a fast local pipe the ACK really can arrive
	// first), and to apply in-flight backpressure before the send.
	OnSend func(index uint64, hash [32]byte)
}

func runSourceLoop(p sourceLoopParams) error {
	if err := checkBlockAlign(p.BlockSize, p.Align); err != nil {
		return err
	}
	f := p.File
	buf := make([]byte, p.BlockSize)
	for bh := range p.Hashes {
		offset := p.Base + int64(bh.index)*p.BlockSize
		n, err := readBlockAt(f, buf, offset, p.Base+p.Size, p.Align)
		if err != nil {
			return fmt.Errorf("read %s at block %d: %w", f.Name(), bh.index, err)
		}
		hash := sha256.Sum256(buf[:n])
		if p.OnRead != nil {
			p.OnRead(bh.index)
		}
		if hash == bh.hash {
			if p.OnSkip != nil {
				p.OnSkip(bh.index)
			}
			continue
		}
		if p.OnSend != nil {
			p.OnSend(bh.index, hash)
		}
		if err := p.Out.WriteFrame(frameData, encodeDataFrame(bh.index, hash, buf[:n])); err != nil {
			return fmt.Errorf("send block %d: %w", bh.index, err)
		}
	}
	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
	Base      int64    // byte offset the window starts at (0 for a whole-file sync)
	BlockSize int64
	Align     int64 // sector alignment for a raw disk handle, else 0/1
	In        *FrameReader
	AckOut    *FrameWriter // optional
	// 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)
	// OnWrite is called right after a block is durably written — the true
	// destination write-throughput measurement point, as opposed to
	// whatever "copied" count a driver derives from the send/ack protocol.
	OnWrite func(index uint64)
}

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 := writeBlockAt(f, block, p.Base+int64(index)*p.BlockSize, p.Align); 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.OnWrite != nil {
				p.OnWrite(index)
			}
			if p.AckOut != nil {
				if err := p.AckOut.WriteFrame(frameAck, encodeIndexFrame(index)); err != nil {
					return err
				}
			}
		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)
		}
	}
}