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 := GetMFTFileReader(volume)
	if err != nil {
		return nil, err
	}
	defer reader.Close()

	bfio := bytes.NewBuffer(make([]byte, 0, length))
	written, err := copyBytes(bfio, reader, length)
	if err != nil {
		return nil, err
	}
	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 := GetMFTFileReader(volume)
	if err != nil {
		return err
	}
	defer reader.Close()

	out, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer out.Close()
	written, err := copyFiles(out, reader, length, fn)
	if err != nil {
		return err
	}
	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) {
	reader, length, err := GetMFTFileReader(volume)
	if err != nil {
		return nil, 0, err
	}
	return reader, length, nil
}

func GetMFTFileReader(volume string) (io.ReadCloser, int64, error) {
	reader, length, _, err := openMFTFile(volume)
	if err != nil {
		return nil, 0, err
	}
	return reader, length, nil
}

func openMFTFile(volume string) (io.ReadCloser, int64, int64, error) {
	if isWin {
		volume = `\\.\` + volume[:len(volume)-1]
	}
	in, err := os.Open(volume)
	if err != nil {
		return nil, 0, 0, err
	}
	success := false
	defer func() {
		if !success {
			in.Close()
		}
	}()

	bootSectorData := make([]byte, 512)
	_, err = io.ReadFull(in, bootSectorData)
	if err != nil {
		return nil, 0, 0, fmt.Errorf("Unable to read boot sector: %v", err)
	}

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

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

	bytesPerCluster := bootSector.BytesPerSector * bootSector.SectorsPerCluster
	if bytesPerCluster <= 0 {
		return nil, 0, 0, fmt.Errorf("Invalid bytes per cluster %d", bytesPerCluster)
	}
	mftPosInBytes := int64(bootSector.MftClusterNumber) * int64(bytesPerCluster)

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

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

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

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

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

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

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

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

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

	success = true
	return fragment.NewReader(in, fragments), totalLength, int64(mftSizeInBytes), nil
}

func copyBytes(dst io.Writer, src io.Reader, totalLength int64) (written int64, err error) {
	return copyWithProgress(dst, src, totalLength, nil)
}

func copyFiles(dst io.Writer, src io.Reader, totalLength int64, fn func(int64, int64, float64)) (written int64, err error) {
	return copyWithProgress(dst, src, totalLength, fn)
}

func copyWithProgress(dst io.Writer, src io.Reader, totalLength int64, fn func(int64, int64, float64)) (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 {
		reportCopyProgress(fn, written, totalLength)
		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
		}
	}
	reportCopyProgress(fn, written, totalLength)
	return written, err
}

func reportCopyProgress(fn func(int64, int64, float64), written int64, totalLength int64) {
	if fn == nil {
		return
	}
	if totalLength <= 0 {
		fn(written, totalLength, 100)
		return
	}
	fn(written, totalLength, float64(written)/float64(totalLength)*100)
}