package mft

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"
)

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

type FileEntry struct {
	Name   string
	Parent uint64
}

const (
	defaultMFTRecordSize = int64(1024)
	maxMFTBatchRecords   = int64(1024)
)

func GetFileListsByMftFn(driver string, fn func(string, bool) bool) ([]MFTFile, error) {
	reader, size, recordSize, err := openMFTFile(driver)
	if err != nil {
		return []MFTFile{}, err
	}
	defer reader.Close()

	return collectMFTFiles(driver, reader, size, recordSize, fn)
}

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) {
	f, err := os.Open(filepath)
	if err != nil {
		return []MFTFile{}, err
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return []MFTFile{}, err
	}

	return collectMFTFiles(" ", f, stat.Size(), defaultMFTRecordSize, fn)
}

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

// WalkRecordsByMFT walks parsed MFT records from a live NTFS volume.
func WalkRecordsByMFT(driver string, fn func(Record) error) error {
	reader, size, recordSize, err := openMFTFile(driver)
	if err != nil {
		return err
	}
	defer reader.Close()

	return walkRecords(reader, size, recordSize, ParseRecord, fn)
}

// WalkRecordsFromMFTFile walks parsed MFT records from a dumped $MFT file.
func WalkRecordsFromMFTFile(filepath string, fn func(Record) error) error {
	f, err := os.Open(filepath)
	if err != nil {
		return err
	}
	defer f.Close()

	stat, err := f.Stat()
	if err != nil {
		return err
	}

	return walkRecords(f, stat.Size(), defaultMFTRecordSize, ParseRecord, fn)
}

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
}

type extendedData struct {
	Size          uint64
	AllocatedSize uint64
}

func collectMFTFiles(diskName string, reader io.Reader, size int64, recordSize int64, fn func(string, bool) bool) ([]MFTFile, error) {
	if fn == nil {
		fn = func(string, bool) bool { return true }
	}

	extendMFTRecord := make(map[uint64]extendedData)
	fileMap := make(map[uint64]FileEntry)
	result := make([]MFTFile, 0)

	err := walkRecords(reader, size, recordSize, ParseRecord, func(record Record) error {
		appendExtendedData(extendMFTRecord, record)

		file, ok := FileFromRecord(record)
		if !ok {
			return nil
		}

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

	for i := range result {
		if attrs, ok := extendMFTRecord[result[i].Node]; ok {
			if result[i].Aszie == 0 {
				applyExtendedData(&result[i], attrs)
			}
			delete(extendMFTRecord, result[i].Node)
		}
		result[i].Path = GetFullUsnPath(diskName, fileMap, result[i].Node)
	}

	return result, nil
}

func walkRecords(reader io.Reader, size int64, recordSize int64, parser func([]byte) (Record, error), visit func(Record) error) error {
	if recordSize <= 0 {
		return fmt.Errorf("invalid MFT record size %d", recordSize)
	}
	if recordSize > maxInt {
		return fmt.Errorf("MFT record size %d overflows maximum int value %d", recordSize, maxInt)
	}
	if parser == nil {
		return fmt.Errorf("nil MFT record parser")
	}
	if visit == nil {
		return fmt.Errorf("nil MFT record visitor")
	}

	chunkSize := recordSize * maxMFTBatchRecords
	if chunkSize <= 0 {
		chunkSize = recordSize
	}
	if size > 0 && chunkSize > size {
		chunkSize = size
	}
	if chunkSize <= 0 {
		chunkSize = recordSize
	}

	intRecordSize := int(recordSize)
	buf := make([]byte, int(chunkSize))
	for {
		got, err := io.ReadFull(reader, buf)
		if err != nil {
			if errors.Is(err, io.EOF) && got == 0 {
				return nil
			}
			if !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
				return err
			}
		}
		if got == 0 {
			return nil
		}

		usable := got - got%intRecordSize
		for offset := 0; offset < usable; offset += intRecordSize {
			record, err := parser(buf[offset : offset+intRecordSize])
			if err != nil {
				continue
			}
			if err := visit(record); err != nil {
				return err
			}
		}

		if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) {
			return nil
		}
	}
}

func appendExtendedData(extended map[uint64]extendedData, record Record) {
	baseRecord := record.BaseRecordReference.ToUint64()
	if baseRecord == 0 {
		return
	}

	for _, attr := range record.Attributes {
		if attr.Type == AttributeTypeData && attr.ActualSize != 0 {
			extended[baseRecord] = extendedData{
				Size:          attr.ActualSize,
				AllocatedSize: attr.AllocatedSize,
			}
		}
	}
}

// FileFromRecord extracts a high-level file entry from a parsed MFT record.
func FileFromRecord(record Record) (MFTFile, bool) {
	if record.Flags&RecordFlagInUse == 0 || record.Flags&RecordFlagIsIndex != 0 {
		return MFTFile{}, false
	}

	file := MFTFile{
		IsDir: record.Flags&RecordFlagIsDirectory != 0,
		Node:  record.FileReference.ToUint64(),
	}
	bestNamespace := FileNameNamespace(0)

	for _, attr := range record.Attributes {
		switch attr.Type {
		case AttributeTypeData:
			file.Size = attr.ActualSize
			file.Aszie = attr.AllocatedSize
		case AttributeTypeStandardInformation:
			info, err := ParseStandardInformation(attr.Data)
			if err == nil {
				file.ModTime = info.FileLastModified
			}
		case AttributeTypeFileName:
			name, nameParent, namespace, ok := bestFileName(file.Name, bestNamespace, attr.Data)
			if ok {
				file.Name = name
				file.Parent = nameParent
				bestNamespace = namespace
			}
		}
	}

	if file.Name == "" {
		return MFTFile{}, false
	}

	return file, true
}

func bestFileName(current string, currentNamespace FileNameNamespace, data []byte) (string, uint64, FileNameNamespace, bool) {
	fileName, err := ParseFileName(data)
	if err != nil || fileName.Name == "" {
		return current, 0, currentNamespace, false
	}
	if !shouldPreferFileNameWithNamespace(current, currentNamespace, fileName.Name, fileName.Namespace) {
		return current, 0, currentNamespace, false
	}
	return fileName.Name, fileName.ParentFileReference.ToUint64(), fileName.Namespace, true
}

func shouldPreferFileName(current string, candidate string) bool {
	return shouldPreferFileNameWithNamespace(current, 0, candidate, 0)
}

func shouldPreferFileNameWithNamespace(current string, currentNamespace FileNameNamespace, candidate string, candidateNamespace FileNameNamespace) bool {
	if candidate == "" {
		return false
	}
	if current == "" {
		return true
	}
	currentRank := fileNameNamespaceRank(currentNamespace)
	candidateRank := fileNameNamespaceRank(candidateNamespace)
	if currentRank != candidateRank {
		return candidateRank > currentRank
	}

	currentShort := strings.Contains(current, "~")
	candidateShort := strings.Contains(candidate, "~")
	if currentShort != candidateShort {
		return currentShort && !candidateShort
	}

	return len(candidate) > len(current)
}

func fileNameNamespaceRank(namespace FileNameNamespace) int {
	switch namespace {
	case FileNameNamespaceWin32, FileNameNamespaceWin32Dos:
		return 3
	case FileNameNamespacePosix:
		return 2
	case FileNameNamespaceDos:
		return 1
	default:
		return 0
	}
}

func applyExtendedData(file *MFTFile, data extendedData) {
	file.Size = data.Size
	file.Aszie = data.AllocatedSize
}