package usn

import (
	"b612.me/stario"
	"b612.me/win32api"
	"fmt"
	"os"
	"reflect"
	"runtime"
	"syscall"
	"unsafe"
)

type DiskInfo struct {
	Driver       string
	Name         string
	Format       string
	SerialNumber uint32
}

func ListDrivers() ([]string, error) {
	drivers := make([]string, 0, 26)
	buf := make([]uint16, 255)
	err := win32api.GetLogicalDriveStringsW(win32api.DWORD(len(buf)), &buf[0])
	if err != nil {
		return drivers, err
	}
	var driver []rune
	for _, v := range buf {
		if v != 0 {
			driver = append(driver, rune(uint8(v)))
			if v == 92 {
				drivers = append(drivers, string(driver))
				driver = []rune{}
			}
		}
	}
	return drivers, nil
}

func GetDiskInfo(disk string) (DiskInfo, error) {
	format := make([]rune, 0, 12)
	name := make([]rune, 0, 128)
	ptr, _ := syscall.UTF16PtrFromString(disk)
	var lpVolumeNameBuffer = make([]uint16, syscall.MAX_PATH+1)
	var nVolumeNameSize = win32api.DWORD(len(lpVolumeNameBuffer))
	var lpVolumeSerialNumber uint32
	var lpMaximumComponentLength uint32
	var lpFileSystemFlags uint32
	var lpFileSystemNameBuffer = make([]uint16, 255)
	var nFileSystemNameSize uint32 = syscall.MAX_PATH + 1
	err := win32api.GetVolumeInformationW(ptr, &lpVolumeNameBuffer[0], nVolumeNameSize, &lpVolumeSerialNumber, &lpMaximumComponentLength,
		&lpFileSystemFlags, &lpFileSystemNameBuffer[0], win32api.DWORD(nFileSystemNameSize))
	for _, v := range lpFileSystemNameBuffer {
		if v != 0 {
			format = append(format, rune(v))
		}
	}
	for _, v := range lpVolumeNameBuffer {
		if v != 0 {
			name = append(name, rune(v))
		}
	}
	return DiskInfo{
		SerialNumber: lpVolumeSerialNumber,
		Driver:       disk,
		Name:         string(name),
		Format:       string(format),
	}, err
}

func DeviceIoControl(handle syscall.Handle, controlCode uint32, in interface{}, out interface{}, done *uint32) (err error) {
	inPtr, inSize := getPointer(in)
	outPtr, outSize := getPointer(out)
	//_,err = syscall.Syscall9(procDeviceIoControl.Addr(), 8, uintptr(handle), uintptr(controlCode), inPtr, uintptr(inSize), outPtr, uintptr(outSize), uintptr(unsafe.Pointer(done)), uintptr(0), 0)
	_, err = win32api.DeviceIoControlPtr(win32api.HANDLE(handle), win32api.DWORD(controlCode), inPtr, win32api.DWORD(inSize), outPtr, win32api.DWORD(outSize), done, nil)
	return
}

func getPointer(i interface{}) (pointer, size uintptr) {
	v := reflect.ValueOf(i)
	switch k := v.Kind(); k {
	case reflect.Ptr:
		t := v.Elem().Type()
		size = t.Size()
		pointer = v.Pointer()
	case reflect.Slice:
		size = uintptr(v.Cap())
		pointer = v.Pointer()
	default:
		fmt.Println("error")
	}
	return
}

// Need a custom Open to work with backup_semantics
func CreateFile(path string, mode int, attrs uint32) (fd syscall.Handle, err error) {
	if len(path) == 0 {
		return syscall.InvalidHandle, win32api.ERROR_FILE_NOT_FOUND
	}
	pathp, err := syscall.UTF16PtrFromString(path)
	if err != nil {
		return syscall.InvalidHandle, err
	}
	var access uint32
	switch mode & (win32api.O_RDONLY | win32api.O_WRONLY | win32api.O_RDWR) {
	case win32api.O_RDONLY:
		access = win32api.GENERIC_READ
	case win32api.O_WRONLY:
		access = win32api.GENERIC_WRITE
	case win32api.O_RDWR:
		access = win32api.GENERIC_READ | win32api.GENERIC_WRITE
	}
	if mode&win32api.O_CREAT != 0 {
		access |= win32api.GENERIC_WRITE
	}
	if mode&win32api.O_APPEND != 0 {
		access &^= win32api.GENERIC_WRITE
		access |= win32api.FILE_APPEND_DATA
	}
	sharemode := uint32(win32api.FILE_SHARE_READ | win32api.FILE_SHARE_WRITE)
	var sa *syscall.SecurityAttributes
	if mode&win32api.O_CLOEXEC == 0 {
		sa = makeInheritSa()
	}
	var createmode uint32
	switch {
	case mode&(win32api.O_CREAT|win32api.O_EXCL) == (win32api.O_CREAT | win32api.O_EXCL):
		createmode = win32api.CREATE_NEW
	case mode&(win32api.O_CREAT|win32api.O_TRUNC) == (win32api.O_CREAT | win32api.O_TRUNC):
		createmode = win32api.CREATE_ALWAYS
	case mode&win32api.O_CREAT == win32api.O_CREAT:
		createmode = win32api.OPEN_ALWAYS
	case mode&win32api.O_TRUNC == win32api.O_TRUNC:
		createmode = win32api.TRUNCATE_EXISTING
	default:
		createmode = win32api.OPEN_EXISTING
	}
	h, e := syscall.CreateFile(pathp, access, sharemode, sa, createmode, attrs, 0)
	return h, e
}

func makeInheritSa() *syscall.SecurityAttributes {
	var sa syscall.SecurityAttributes
	sa.Length = uint32(unsafe.Sizeof(sa))
	sa.InheritHandle = 1
	return &sa
}

// Query usn journal data
func queryUsnJournal(fd syscall.Handle) (ujd win32api.USN_JOURNAL_DATA, done uint32, err error) {
	err = DeviceIoControl(fd, win32api.FSCTL_QUERY_USN_JOURNAL, []byte{}, &ujd, &done)
	return
}

func readUsnJournal(fd syscall.Handle, rujd *win32api.READ_USN_JOURNAL_DATA) (data []byte, done uint32, err error) {
	data = make([]byte, 0x1000)
	err = DeviceIoControl(fd, win32api.FSCTL_READ_USN_JOURNAL, rujd, data, &done)
	return
}

func enumUsnData(fd syscall.Handle, med *win32api.MFT_ENUM_DATA) (data []byte, done uint32, err error) {
	data = make([]byte, 0x10000)
	err = DeviceIoControl(fd, win32api.FSCTL_ENUM_USN_DATA, med, data, &done)
	return
}

type FileEntry struct {
	Name   string
	Parent win32api.DWORDLONG
	Type   uint8
}

type FileMonitor struct {
	Name   string
	Self   win32api.DWORDLONG
	Parent win32api.DWORDLONG
	Type   uint8
	Reason string
}

func ListUsnFile(driver string) (map[win32api.DWORDLONG]FileEntry, error) {
	fileMap := make(map[win32api.DWORDLONG]FileEntry)
	pDriver := "\\\\.\\" + driver[:len(driver)-1]
	fd, err := CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return fileMap, err
	}
	ujd, _, err := queryUsnJournal(fd)
	if err != nil {
		return fileMap, err
	}
	med := win32api.MFT_ENUM_DATA{0, 0, ujd.NextUsn}
	for {
		data, done, err := enumUsnData(fd, &med)
		if err != nil && done != 0 {
			return fileMap, err
		}
		if done == 0 {
			return fileMap, nil
		}

		var usn win32api.USN = *(*win32api.USN)(unsafe.Pointer(&data[0]))
		// fmt.Println("usn", usn)

		var ur *win32api.USN_RECORD
		for i := unsafe.Sizeof(usn); i < uintptr(done); i += uintptr(ur.RecordLength) {
			ur = (*win32api.USN_RECORD)(unsafe.Pointer(&data[i]))
			nameLength := uintptr(ur.FileNameLength) / unsafe.Sizeof(ur.FileName[0])
			fnp := unsafe.Pointer(&data[i+uintptr(ur.FileNameOffset)])
			fnUtf := (*[10000]uint16)(fnp)[:nameLength]
			fn := syscall.UTF16ToString(fnUtf)
			(*reflect.SliceHeader)(unsafe.Pointer(&fn)).Cap = int(nameLength)
			typed := uint8(0)
			if ur.FileAttributes&win32api.FILE_ATTRIBUTE_DIRECTORY != 0 {
				typed = 1
			}
			// fmt.Println("len", ur.FileNameLength, ur.FileNameOffset, "fn", fn)
			fileMap[ur.FileReferenceNumber] = FileEntry{Name: fn, Parent: ur.ParentFileReferenceNumber, Type: typed}
		}
		med.StartFileReferenceNumber = win32api.DWORDLONG(usn)
	}
}

func ListUsnFileFn(driver string, searchFn func(string, bool) bool) (map[win32api.DWORDLONG]FileEntry, error) {
	fileMap := make(map[win32api.DWORDLONG]FileEntry)
	pDriver := "\\\\.\\" + driver[:len(driver)-1]
	fd, err := CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return fileMap, err
	}
	ujd, _, err := queryUsnJournal(fd)
	if err != nil {
		return fileMap, err
	}
	med := win32api.MFT_ENUM_DATA{0, 0, ujd.NextUsn}
	for {
		data, done, err := enumUsnData(fd, &med)
		if err != nil && done != 0 {
			return fileMap, err
		}
		if done == 0 {
			return fileMap, nil
		}

		var usn win32api.USN = *(*win32api.USN)(unsafe.Pointer(&data[0]))
		// fmt.Println("usn", usn)

		var ur *win32api.USN_RECORD
		for i := unsafe.Sizeof(usn); i < uintptr(done); i += uintptr(ur.RecordLength) {
			ur = (*win32api.USN_RECORD)(unsafe.Pointer(&data[i]))
			nameLength := uintptr(ur.FileNameLength) / unsafe.Sizeof(ur.FileName[0])
			fnp := unsafe.Pointer(&data[i+uintptr(ur.FileNameOffset)])
			fnUtf := (*[10000]uint16)(fnp)[:nameLength]
			fn := syscall.UTF16ToString(fnUtf)
			(*reflect.SliceHeader)(unsafe.Pointer(&fn)).Cap = int(nameLength)
			typed := uint8(0)
			if ur.FileAttributes&win32api.FILE_ATTRIBUTE_DIRECTORY != 0 {
				typed = 1
			}
			if typed == 1 || searchFn(fn, typed == 1) {
				// fmt.Println("len", ur.FileNameLength, ur.FileNameOffset, "fn", fn)
				fileMap[ur.FileReferenceNumber] = FileEntry{Name: fn, Parent: ur.ParentFileReferenceNumber, Type: typed}
			}
		}
		med.StartFileReferenceNumber = win32api.DWORDLONG(usn)
	}
}

func GetFullUsnPath(diskName string, fileMap map[win32api.DWORDLONG]FileEntry, id win32api.DWORDLONG) (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
}

func GetFullUsnPathEntry(diskName string, fileMap map[win32api.DWORDLONG]FileEntry, en FileMonitor) (name string) {
	fileMap[en.Self] = FileEntry{
		Name:   en.Name,
		Parent: en.Parent,
		Type:   en.Type,
	}
	id := en.Self
	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
}

const (
	ALL_FILES = iota
	ONLY_FOLDER
	NO_FOLDER
)

func ListNTFSUsnDriverFilesFn(diskName string, searchFn func(string, bool) bool) ([]string, error) {
	var result []string
	data, err := ListUsnFileFn(diskName, searchFn)
	if err != nil {
		return result, err
	}
	return listNTFSUsnDriverFiles(diskName, searchFn, data)
}

func ListNTFSUsnDriverFiles(diskName string, folder uint8) ([]string, error) {
	var result []string
	data, err := ListUsnFile(diskName)
	if err != nil {
		return result, err
	}
	return listNTFSUsnDriverFiles(diskName, func(name string, tp bool) bool {
		if !tp && folder == ONLY_FOLDER {
			return false
		}
		if tp && folder == NO_FOLDER {
			return false
		}
		return true
	}, data)
}

func listNTFSUsnDriverFiles(diskName string, fn func(string, bool) bool, data map[win32api.DWORDLONG]FileEntry) ([]string, error) {
	result := make([]string, len(data))
	i := 0
	for k, v := range data {
		if !fn(v.Name, v.Type == 1) {
			continue
		}
		name := GetFullUsnPath(diskName, data, k)
		result[i] = name
		i++
	}
	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = i
	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Len = i
	data = nil
	data = make(map[win32api.DWORDLONG]FileEntry, 0)
	runtime.GC()
	return result, nil
}

func ListNTFSUsnDriverInfoFn(diskName string, searchFn func(string, bool) bool) ([]FileStat, error) {
	data, err := ListUsnFileFn(diskName, searchFn)
	if err != nil {
		return nil, err
	}
	return listNTFSUsnDriverInfo(diskName, searchFn, data)
}

func ListNTFSUsnDriverInfo(diskName string, folder uint8) ([]FileStat, error) {
	data, err := ListUsnFile(diskName)
	if err != nil {
		return nil, err
	}
	return listNTFSUsnDriverInfo(diskName, func(name string, tp bool) bool {
		if !tp && folder == ONLY_FOLDER {
			return false
		}
		if tp && folder == NO_FOLDER {
			return false
		}
		return true
	}, data)
}

func listNTFSUsnDriverInfo(diskName string, fn func(string, bool) bool, data map[win32api.DWORDLONG]FileEntry) ([]FileStat, error) {
	//fmt.Println("finished 1")
	pDriver := "\\\\.\\" + diskName[:len(diskName)-1]
	fd, err := CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return nil, err
	}
	defer syscall.Close(fd)
	result := make([]FileStat, len(data))
	i := int(0)
	wg := stario.NewWaitGroup(100)
	for k, v := range data {
		if !fn(v.Name, v.Type == 1) {
			continue
		}
		wg.Add(1)
		go func(k win32api.DWORDLONG, v FileEntry, i int) {
			defer wg.Done()
			//now := time.Now().UnixNano()
			/*
				fd2, err := OpenFileByIdWithfd(fd, k, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
				if err != nil {
					return
				}
				//fmt.Println("cost", float64((time.Now().UnixNano()-now)/1000000))
				var info syscall.ByHandleFileInformation
				err = syscall.GetFileInformationByHandle(fd2, &info)
				syscall.Close(fd2)
				//fmt.Println("cost", float64((time.Now().UnixNano()-now)/1000000))
				if err != nil {
					return
				}

			*/
			path := GetFullUsnPath(diskName, data, k)
			fileInfo, err := os.Stat(path)
			if err != nil {
				return
			}
			fs := fileInfo.Sys().(*syscall.Win32FileAttributeData)
			stat := FileStat{
				FileAttributes: fs.FileAttributes,
				CreationTime:   fs.CreationTime,
				LastAccessTime: fs.LastAccessTime,
				LastWriteTime:  fs.LastWriteTime,
				FileSizeHigh:   fs.FileSizeHigh,
				FileSizeLow:    fs.FileSizeLow,
			}
			stat.name = v.Name
			stat.path = path
			return
			result[i] = stat
			//result[i] = newFileStatFromInformation(&info, v.Name, path)
		}(k, v, i)
		i++
	}
	wg.Wait()
	//fmt.Println("finished 2")
	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Cap = i
	(*reflect.SliceHeader)(unsafe.Pointer(&result)).Len = i
	data = nil
	//data = make(map[win32api.DWORDLONG]FileEntry, 0)
	runtime.GC()
	return result, nil
}

func getUsnJournalReasonString(reason win32api.DWORD) (s string) {
	var reasons = []string{
		"DataOverwrite",       // 0x00000001
		"DataExtend",          // 0x00000002
		"DataTruncation",      // 0x00000004
		"0x00000008",          // 0x00000008
		"NamedDataOverwrite",  // 0x00000010
		"NamedDataExtend",     // 0x00000020
		"NamedDataTruncation", // 0x00000040
		"0x00000080",          // 0x00000080
		"FileCreate",          // 0x00000100
		"FileDelete",          // 0x00000200
		"PropertyChange",      // 0x00000400
		"SecurityChange",      // 0x00000800
		"RenameOldName",       // 0x00001000
		"RenameNewName",       // 0x00002000
		"IndexableChange",     // 0x00004000
		"BasicInfoChange",     // 0x00008000
		"HardLinkChange",      // 0x00010000
		"CompressionChange",   // 0x00020000
		"EncryptionChange",    // 0x00040000
		"ObjectIdChange",      // 0x00080000
		"ReparsePointChange",  // 0x00100000
		"StreamChange",        // 0x00200000
		"0x00400000",          // 0x00400000
		"0x00800000",          // 0x00800000
		"0x01000000",          // 0x01000000
		"0x02000000",          // 0x02000000
		"0x04000000",          // 0x04000000
		"0x08000000",          // 0x08000000
		"0x10000000",          // 0x10000000
		"0x20000000",          // 0x20000000
		"0x40000000",          // 0x40000000
		"*Close*",             // 0x80000000
	}
	for i := 0; reason != 0; {
		if reason&1 == 1 {
			s = s + ", " + reasons[i]
		}
		reason >>= 1
		i++
	}
	return
}

func MonitorUsnChange(driver string, rec chan FileMonitor) error {
	pDriver := "\\\\.\\" + driver[:len(driver)-1]
	fd, err := CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return err
	}
	ujd, _, err := queryUsnJournal(fd)
	if err != nil {
		return err
	}

	rujd := win32api.READ_USN_JOURNAL_DATA{ujd.NextUsn, 0xFFFFFFFF, 0, 0, 1, ujd.UsnJournalID}

	for {
		var usn win32api.USN
		data, done, err := readUsnJournal(fd, &rujd)
		if err != nil || done <= uint32(unsafe.Sizeof(usn)) {
			return err
		}

		usn = *(*win32api.USN)(unsafe.Pointer(&data[0]))

		var ur *win32api.USN_RECORD
		for i := unsafe.Sizeof(usn); i < uintptr(done); i += uintptr(ur.RecordLength) {
			ur = (*win32api.USN_RECORD)(unsafe.Pointer(&data[i]))
			nameLength := uintptr(ur.FileNameLength) / unsafe.Sizeof(ur.FileName[0])
			fnp := unsafe.Pointer(&data[i+uintptr(ur.FileNameOffset)])
			fn := syscall.UTF16ToString((*[10000]uint16)(fnp)[:nameLength])
			(*reflect.SliceHeader)(unsafe.Pointer(&fn)).Cap = int(nameLength)
			// fmt.Println("len", ur.FileNameLength, ur.FileNameOffset, "fn", getFullPath(folders, ur.ParentFileReferenceNumber), syscall.UTF16ToString(fn), getUsnJournalReasonString(ur.Reason))
			typed := uint8(0)
			if ur.FileAttributes&win32api.FILE_ATTRIBUTE_DIRECTORY != 0 {
				typed = 1
			}
			// fmt.Println("len", ur.FileNameLength, ur.FileNameOffset, "fn", fn)
			rec <- FileMonitor{Name: fn, Parent: ur.ParentFileReferenceNumber, Type: typed, Self: ur.FileReferenceNumber, Reason: getUsnJournalReasonString(ur.Reason)}
		}
		rujd.StartUsn = usn
		if usn == 0 {
			return nil
		}
	}
}

func GetUsnFileInfo(diskName string, fileMap map[win32api.DWORDLONG]FileEntry, id win32api.DWORDLONG) (FileStat, error) {
	name := fileMap[id].Name
	path := GetFullUsnPath(diskName, fileMap, id)
	fd, err := OpenFileById(diskName, id, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return FileStat{}, err
	}
	var info syscall.ByHandleFileInformation
	err = syscall.GetFileInformationByHandle(fd, &info)
	return newFileStatFromInformation(&info, name, path), err
}

// Need a custom Open to work with backup_semantics
func OpenFileById(diskName string, id win32api.DWORDLONG, mode int, attrs uint32) (syscall.Handle, error) {
	pDriver := "\\\\.\\" + diskName[:len(diskName)-1]
	fd, err := CreateFile(pDriver, syscall.O_RDONLY, win32api.FILE_ATTRIBUTE_NORMAL)
	if err != nil {
		return syscall.InvalidHandle, err
	}
	defer syscall.Close(fd)
	return OpenFileByIdWithfd(fd, id, mode, attrs)
}

func OpenFileByIdWithfd(fd syscall.Handle, id win32api.DWORDLONG, mode int, attrs uint32) (syscall.Handle, error) {
	var access uint32
	switch mode & (win32api.O_RDONLY | win32api.O_WRONLY | win32api.O_RDWR) {
	case win32api.O_RDONLY:
		access = win32api.GENERIC_READ
	case win32api.O_WRONLY:
		access = win32api.GENERIC_WRITE
	case win32api.O_RDWR:
		access = win32api.GENERIC_READ | win32api.GENERIC_WRITE
	}
	if mode&win32api.O_CREAT != 0 {
		access |= win32api.GENERIC_WRITE
	}
	if mode&win32api.O_APPEND != 0 {
		access &^= win32api.GENERIC_WRITE
		access |= win32api.FILE_APPEND_DATA
	}
	sharemode := uint32(win32api.FILE_SHARE_READ | win32api.FILE_SHARE_WRITE)
	var sa *syscall.SecurityAttributes
	if mode&win32api.O_CLOEXEC == 0 {
		sa = makeInheritSa()
	}
	fid := win32api.FILE_ID_DESCRIPTOR{
		DwSize: 16,
		Type:   0,
		FileId: id,
	}
	fid.DwSize = win32api.DWORD(unsafe.Sizeof(fid))
	h, e := win32api.OpenFileById(win32api.HANDLE(fd), &fid, win32api.DWORD(access),
		win32api.DWORD(sharemode), sa, win32api.DWORD(attrs))
	return syscall.Handle(h), e
}