package mft

import (
	"b612.me/wincmd/ntfs/bootsect"
	"b612.me/wincmd/ntfs/fragment"
	"bytes"
	"fmt"
	"io"
	"os"
	"runtime"
)

const supportedOemId = "NTFS    "

const isWin = runtime.GOOS == "windows"

func GetMFTFileBytes(volume string) ([]byte, error) {
	reader, length, err := GetMFTFile(volume)
	if err != nil {
		return nil, err
	}
	buf := make([]byte, length)
	bfio := bytes.NewBuffer(buf)
	written, err := copyBytes(bfio, reader, length)
	if written != length {
		return nil, fmt.Errorf("Write Not Ok,Should %d got %d", length, written)
	}
	return bfio.Bytes(), nil
}

func DumpMFTFile(volume, filepath string, fn func(int64, int64, float64)) error {
	reader, length, err := GetMFTFile(volume)
	if err != nil {
		return err
	}
	out, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer out.Close()
	written, err := copyFiles(out, reader, length, fn)
	if written != length {
		return fmt.Errorf("Write Not Ok,Should %d got %d", length, written)
	}
	return nil
}

func GetMFTFile(volume string) (io.Reader, int64, error) {
	if isWin {
		volume = `\\.\` + volume[:len(volume)-1]
	}
	in, err := os.Open(volume)
	if err != nil {
		return nil, 0, err
	}
	bootSectorData := make([]byte, 512)
	_, err = io.ReadFull(in, bootSectorData)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to read boot sector: %v\n", err)
	}

	bootSector, err := bootsect.Parse(bootSectorData)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to parse boot sector data: %v\n", err)
	}

	if bootSector.OemId != supportedOemId {
		return nil, 0, fmt.Errorf("Unknown OemId (file system type) %q (expected %q)\n", bootSector.OemId, supportedOemId)
	}

	bytesPerCluster := bootSector.BytesPerSector * bootSector.SectorsPerCluster
	mftPosInBytes := int64(bootSector.MftClusterNumber) * int64(bytesPerCluster)

	_, err = in.Seek(mftPosInBytes, 0)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to seek to MFT position: %v\n", err)
	}

	mftSizeInBytes := bootSector.FileRecordSegmentSizeInBytes
	mftData := make([]byte, mftSizeInBytes)
	_, err = io.ReadFull(in, mftData)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to read $MFT record: %v\n", err)
	}

	record, err := ParseRecord(mftData)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to parse $MFT record: %v\n", err)
	}

	dataAttributes := record.FindAttributes(AttributeTypeData)
	if len(dataAttributes) == 0 {
		return nil, 0, fmt.Errorf("No $DATA attribute found in $MFT record\n")
	}

	if len(dataAttributes) > 1 {
		return nil, 0, fmt.Errorf("More than 1 $DATA attribute found in $MFT record\n")
	}

	dataAttribute := dataAttributes[0]
	if dataAttribute.Resident {
		return nil, 0, fmt.Errorf("Don't know how to handle resident $DATA attribute in $MFT record\n")
	}

	dataRuns, err := ParseDataRuns(dataAttribute.Data)
	if err != nil {
		return nil, 0, fmt.Errorf("Unable to parse dataruns in $MFT $DATA record: %v\n", err)
	}

	if len(dataRuns) == 0 {
		return nil, 0, fmt.Errorf("No dataruns found in $MFT $DATA record\n")
	}

	fragments := DataRunsToFragments(dataRuns, bytesPerCluster)
	totalLength := int64(0)
	for _, frag := range fragments {
		totalLength += int64(frag.Length)
	}

	return fragment.NewReader(in, fragments), totalLength, nil
}

func copyBytes(dst io.Writer, src io.Reader, totalLength int64) (written int64, err error) {
	buf := make([]byte, 1024*1024)

	// Below copied from io.copyBuffer (https://golang.org/src/io/io.go?s=12796:12856#L380)
	for {

		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
	}
	return written, err
}

func copyFiles(dst io.Writer, src io.Reader, totalLength int64, fn func(int64, int64, float64)) (written int64, err error) {
	buf := make([]byte, 1024*1024)
	onePercent := float64(written) / float64(totalLength) * float64(100.0)

	// Below copied from io.copyBuffer (https://golang.org/src/io/io.go?s=12796:12856#L380)
	for {
		fn(written, totalLength, onePercent)
		nr, er := src.Read(buf)
		if nr > 0 {
			nw, ew := dst.Write(buf[0:nr])
			if nw > 0 {
				written += int64(nw)
			}
			if ew != nil {
				err = ew
				break
			}
			if nr != nw {
				err = io.ErrShortWrite
				break
			}
		}
		if er != nil {
			if er != io.EOF {
				err = er
			}
			break
		}
	}
	fn(written, totalLength, onePercent)
	return written, err
}