package wincmd

import (
	"b612.me/win32api"
	"b612.me/wincmd/ntfs/binutil"
	"b612.me/wincmd/ntfs/mft"
	"b612.me/wincmd/ntfs/utf16"
	"encoding/binary"
	"errors"
	"io"
	"os"
	"reflect"
	"runtime"
	"time"
	"unsafe"
)

type MFTFile struct {
	Name    string
	Path    string
	ModTime time.Time
	Size    uint64
	Aszie   uint64
	IsDir   bool
	Node    uint64
}

func GetFileListsByMftFn(driver string, fn func(string, bool) bool) ([]MFTFile, error) {
	var result []MFTFile
	fileMap := make(map[win32api.DWORDLONG]FileEntry)
	f, size, err := mft.GetMFTFile(driver)
	if err != nil {
		return []MFTFile{}, err
	}
	recordSize := int64(1024)
	alreadyGot := int64(0)
	maxRecordSize := size / recordSize
	if maxRecordSize > 1024 {
		maxRecordSize = 1024
	}
	for {
		for {
			if (size - alreadyGot) < maxRecordSize*recordSize {
				maxRecordSize--
			} else {
				break
			}
		}
		if maxRecordSize < 10 {
			maxRecordSize = 1
		}
		buf := make([]byte, maxRecordSize*recordSize)
		got, err := io.ReadFull(f, buf)
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return []MFTFile{}, err
		}
		alreadyGot += int64(got)
		for j := int64(0); j < 1024*maxRecordSize; j += 1024 {
			record, err := mft.ParseRecord(buf[j : j+1024])
			if err != nil {
				continue
			}
			if record.Flags&mft.RecordFlagInUse == 1 && record.Flags&mft.RecordFlagIsIndex == 0 {
				var file MFTFile
				file.IsDir = record.Flags&mft.RecordFlagIsDirectory != 0
				file.Node = record.FileReference.ToUint64()
				parent := uint64(0)
				for _, v := range record.Attributes {
					if v.Type == mft.AttributeTypeFileName {
						name := utf16.DecodeString(v.Data[66:], binary.LittleEndian)
						if len(file.Name) < len(name) && len(name) > 0 {
							if len(name) > 2 && name[len(name)-2] == '~' {
								continue
							}
							file.Name = name
						}
						if file.Name != "" {
							parent = binutil.NewLittleEndianReader(v.Data[:8]).Uint64(0)
						}
					}
					if v.Type == mft.AttributeTypeData {
						file.Size = v.ActualSize
						file.Aszie = v.AllocatedSize
					}
					if v.Type == mft.AttributeTypeStandardInformation {
						if len(v.Data) < 48 {
							continue
						}
						r := binutil.NewLittleEndianReader(v.Data)
						file.ModTime = mft.ConvertFileTime(r.Uint64(0x08))
					}
				}
				if file.Name != "" {
					canAdd := fn(file.Name, file.IsDir)
					if canAdd {
						result = append(result, file)
					}
					if canAdd || file.IsDir {
						fileMap[win32api.DWORDLONG(file.Node)] = FileEntry{
							Name:   file.Name,
							Parent: win32api.DWORDLONG(parent),
							Type:   0,
						}
					}
				}
			}
		}
	}

	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = len(result)
	for k, v := range result {
		result[k].Path = GetFullUsnPath(driver, fileMap, win32api.DWORDLONG(v.Node))
	}
	fileMap = nil
	runtime.GC()
	return result, nil
}

func GetFileListsByMft(driver string) ([]MFTFile, error) {
	return GetFileListsByMftFn(driver, func(string, bool) bool { return true })
}

func GetFileListsFromMftFileFn(filepath string, fn func(string, bool) bool) ([]MFTFile, error) {
	var result []MFTFile
	fileMap := make(map[win32api.DWORDLONG]FileEntry)
	f, err := os.Open(filepath)
	if err != nil {
		return []MFTFile{}, err
	}
	stat, err := f.Stat()
	if err != nil {
		return []MFTFile{}, err
	}
	size := stat.Size()
	recordSize := int64(1024)
	alreadyGot := int64(0)
	maxRecordSize := size / recordSize
	if maxRecordSize > 1024 {
		maxRecordSize = 1024
	}
	for {
		for {
			if (size - alreadyGot) < maxRecordSize*recordSize {
				maxRecordSize--
			} else {
				break
			}
		}
		if maxRecordSize < 10 {
			maxRecordSize = 1
		}
		buf := make([]byte, maxRecordSize*recordSize)
		got, err := io.ReadFull(f, buf)
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return []MFTFile{}, err
		}
		alreadyGot += int64(got)
		for j := int64(0); j < 1024*maxRecordSize; j += 1024 {
			record, err := mft.ParseRecord(buf[j : j+1024])
			if err != nil {
				continue
			}
			if record.Flags&mft.RecordFlagInUse == 1 && record.Flags&mft.RecordFlagIsIndex == 0 {
				var file MFTFile
				file.IsDir = record.Flags&mft.RecordFlagIsDirectory != 0
				file.Node = record.FileReference.ToUint64()
				parent := uint64(0)
				for _, v := range record.Attributes {
					if v.Type == mft.AttributeTypeFileName {
						name := utf16.DecodeString(v.Data[66:], binary.LittleEndian)
						if len(file.Name) < len(name) && len(name) > 0 {
							if len(name) > 2 && name[len(name)-2] == '~' {
								continue
							}
							file.Name = name
						}
						if file.Name != "" {
							parent = binutil.NewLittleEndianReader(v.Data[:8]).Uint64(0)
						}
					}
					if v.Type == mft.AttributeTypeData {
						file.Size = v.ActualSize
						file.Aszie = v.AllocatedSize
					}
					if v.Type == mft.AttributeTypeStandardInformation {
						if len(v.Data) < 48 {
							continue
						}
						r := binutil.NewLittleEndianReader(v.Data)
						file.ModTime = mft.ConvertFileTime(r.Uint64(0x08))
					}
				}
				if file.Name != "" {
					canAdd := fn(file.Name, file.IsDir)
					if canAdd {
						result = append(result, file)
					}
					if canAdd || file.IsDir {
						fileMap[win32api.DWORDLONG(file.Node)] = FileEntry{
							Name:   file.Name,
							Parent: win32api.DWORDLONG(parent),
							Type:   0,
						}
					}
				}
			}
		}
	}

	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = len(result)
	for k, v := range result {
		result[k].Path = GetFullUsnPath(" ", fileMap, win32api.DWORDLONG(v.Node))
	}
	fileMap = nil
	runtime.GC()
	return result, nil
}

func GetFileListsFromMftFile(filepath string) ([]MFTFile, error) {
	return GetFileListsFromMftFileFn(filepath, func(string, bool) bool { return true })
}