package mft

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

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

func GetFileListsByMftFn(driver string, fn func(string, bool) bool) ([]MFTFile, error) {
	var result []MFTFile
	extendMftRecord := make(map[uint64][]Attribute)
	fileMap := make(map[uint64]FileEntry)
	f, size, err := 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 := ParseRecord(buf[j : j+1024])
			if err != nil {
				continue
			}
			if record.BaseRecordReference.ToUint64() != 0 {
				val := extendMftRecord[record.BaseRecordReference.ToUint64()]
				for _, v := range record.Attributes {
					if v.Type == AttributeTypeData && v.ActualSize != 0 {
						val = append(val, v)
					}
				}
				if len(val) != 0 {
					extendMftRecord[record.BaseRecordReference.ToUint64()] = val
				}
			}
			if record.Flags&RecordFlagInUse == 1 && record.Flags&RecordFlagIsIndex == 0 {
				var file MFTFile
				file.IsDir = record.Flags&RecordFlagIsDirectory != 0
				file.Node = record.FileReference.ToUint64()
				parent := uint64(0)
				for _, v := range record.Attributes {
					if v.Type == AttributeTypeData {
						file.Size = v.ActualSize
						file.Aszie = v.AllocatedSize
					}
					if v.Type == AttributeTypeStandardInformation {
						if len(v.Data) >= 48 {
							r := binutil.NewLittleEndianReader(v.Data)
							file.ModTime = ConvertFileTime(r.Uint64(0x08))
						}
					}
					if v.Type == AttributeTypeFileName {
						name := utf16.DecodeString(v.Data[66:], binary.LittleEndian)
						if len(file.Name) < len(name) && len(name) > 0 {
							if len(file.Name) > 0 && !strings.Contains(file.Name, "~") {
								continue
							}
							file.Name = name
						}
						if file.Name != "" {
							parent = binutil.NewLittleEndianReader(v.Data[:8]).Uint64(0)
						}
					}
				}

				if file.Name != "" {
					canAdd := fn(file.Name, file.IsDir)
					if canAdd {
						result = append(result, file)
					}
					if canAdd || file.IsDir {
						fileMap[uint64(file.Node)] = FileEntry{
							Name:   file.Name,
							Parent: uint64(parent),
						}
					}
				}
			}
		}
	}

	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = len(result)
	for k, v := range result {
		if attrs, ok := extendMftRecord[v.Node]; ok {
			if v.Aszie == 0 {
				for _, v := range attrs {
					if v.Type == AttributeTypeData && v.ActualSize != 0 {
						result[k].Size = v.ActualSize
						result[k].Aszie = v.AllocatedSize
					}
				}
			}
			delete(extendMftRecord, v.Node)
		}
		result[k].Path = GetFullUsnPath(driver, fileMap, uint64(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
	extendMftRecord := make(map[uint64][]Attribute)
	fileMap := make(map[uint64]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 := ParseRecord(buf[j : j+1024])
			if err != nil {
				continue
			}
			if record.BaseRecordReference.ToUint64() != 0 {
				val := extendMftRecord[record.BaseRecordReference.ToUint64()]
				for _, v := range record.Attributes {
					if v.Type == AttributeTypeData && v.ActualSize != 0 {
						val = append(val, v)
					}
				}
				if len(val) != 0 {
					extendMftRecord[record.BaseRecordReference.ToUint64()] = val
				}
			}
			if record.Flags&RecordFlagInUse == 1 && record.Flags&RecordFlagIsIndex == 0 {
				var file MFTFile
				file.IsDir = record.Flags&RecordFlagIsDirectory != 0
				file.Node = record.FileReference.ToUint64()
				parent := uint64(0)
				for _, v := range record.Attributes {
					if v.Type == AttributeTypeData {
						file.Size = v.ActualSize
						file.Aszie = v.AllocatedSize
					}
					if v.Type == AttributeTypeStandardInformation {
						if len(v.Data) >= 48 {
							r := binutil.NewLittleEndianReader(v.Data)
							file.ModTime = ConvertFileTime(r.Uint64(0x08))
						}
					}
					if v.Type == AttributeTypeFileName {
						name := utf16.DecodeString(v.Data[66:], binary.LittleEndian)
						if len(file.Name) < len(name) && len(name) > 0 {
							if len(file.Name) > 0 && !strings.Contains(file.Name, "~") {
								continue
							}
							file.Name = name
						}
						if file.Name != "" {
							parent = binutil.NewLittleEndianReader(v.Data[:8]).Uint64(0)
						}
					}
				}
				if file.Name != "" {
					canAdd := fn(file.Name, file.IsDir)
					if canAdd {
						result = append(result, file)
					}
					if canAdd || file.IsDir {
						fileMap[uint64(file.Node)] = FileEntry{
							Name:   file.Name,
							Parent: uint64(parent),
						}
					}
				}
			}
		}
	}

	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = len(result)
	for k, v := range result {
		if attrs, ok := extendMftRecord[v.Node]; ok {
			if v.Aszie == 0 {
				for _, v := range attrs {
					if v.Type == AttributeTypeData && v.ActualSize != 0 {
						result[k].Size = v.ActualSize
						result[k].Aszie = v.AllocatedSize
					}
				}
			}
			delete(extendMftRecord, v.Node)
		}
		result[k].Path = GetFullUsnPath(" ", fileMap, uint64(v.Node))
	}
	fileMap = nil
	runtime.GC()
	return result, nil
}

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

func GetFullUsnPath(diskName string, fileMap map[uint64]FileEntry, id uint64) (name string) {
	for id != 0 {
		fe := fileMap[id]
		if id == fe.Parent {
			name = "\\" + name
			break
		}
		if name == "" {
			name = fe.Name
		} else {
			name = fe.Name + "\\" + name
		}
		id = fe.Parent
	}
	name = diskName[:len(diskName)-1] + name
	return
}