// vfs/disk/disk.go
package disk

import (
	"container/list"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"s1d3sw1ped/SteamCache2/steamcache/logger"
	"s1d3sw1ped/SteamCache2/vfs"
	"s1d3sw1ped/SteamCache2/vfs/vfserror"
	"strings"
	"sync"
	"time"

	"github.com/docker/go-units"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
)

var (
	diskCapacityBytes = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "disk_cache_capacity_bytes",
			Help: "Total capacity of the disk cache in bytes",
		},
	)

	diskSizeBytes = promauto.NewGauge(
		prometheus.GaugeOpts{
			Name: "disk_cache_size_bytes",
			Help: "Total size of the disk cache in bytes",
		},
	)

	diskReadBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "disk_cache_read_bytes_total",
			Help: "Total number of bytes read from the disk cache",
		},
	)

	diskWriteBytes = promauto.NewCounter(
		prometheus.CounterOpts{
			Name: "disk_cache_write_bytes_total",
			Help: "Total number of bytes written to the disk cache",
		},
	)
)

// Ensure DiskFS implements VFS.
var _ vfs.VFS = (*DiskFS)(nil)

// DiskFS is a virtual file system that stores files on disk.
type DiskFS struct {
	root string

	info         map[string]*vfs.FileInfo
	capacity     int64
	size         int64
	mu           sync.RWMutex
	keyLocks     sync.Map // map[string]*sync.RWMutex
	LRU          *lruList
	accessCounts map[string]int64 // key: filename, value: access count
	fileMeta     map[string]struct {
		AccessCount int64
		CTime       int64
	} // key: filename
}

// lruList for LRU eviction
type lruList struct {
	list *list.List
	elem map[string]*list.Element
}

func newLruList() *lruList {
	return &lruList{
		list: list.New(),
		elem: make(map[string]*list.Element),
	}
}

func (l *lruList) MoveToFront(key string) {
	if e, ok := l.elem[key]; ok {
		l.list.MoveToFront(e)
	}
}

func (l *lruList) Add(key string, fi *vfs.FileInfo) *list.Element {
	e := l.list.PushFront(fi)
	l.elem[key] = e
	return e
}

func (l *lruList) Remove(key string) {
	if e, ok := l.elem[key]; ok {
		l.list.Remove(e)
		delete(l.elem, key)
	}
}

func (l *lruList) Back() *vfs.FileInfo {
	if e := l.list.Back(); e != nil {
		return e.Value.(*vfs.FileInfo)
	}
	return nil
}

// New creates a new DiskFS.
func new(root string, capacity int64, skipinit bool) *DiskFS {
	if capacity <= 0 {
		panic("disk capacity must be greater than 0") // panic if the capacity is less than or equal to 0
	}

	if root == "" {
		panic("disk root must not be empty") // panic if the root is empty
	}

	fi, err := os.Stat(root)
	if err != nil {
		if !os.IsNotExist(err) {
			panic(err) // panic if the error is something other than not found
		}
		os.Mkdir(root, 0755)    // create the root directory if it does not exist
		fi, err = os.Stat(root) // re-stat to get the file info
		if err != nil {
			panic(err) // panic if the re-stat fails
		}
	}
	if !fi.IsDir() {
		panic("disk root must be a directory") // panic if the root is not a directory
	}

	dfs := &DiskFS{
		root:         root,
		info:         make(map[string]*vfs.FileInfo),
		capacity:     capacity,
		mu:           sync.RWMutex{},
		keyLocks:     sync.Map{},
		LRU:          newLruList(),
		accessCounts: make(map[string]int64),
		fileMeta: make(map[string]struct {
			AccessCount int64
			CTime       int64
		}),
	}

	os.MkdirAll(dfs.root, 0755)

	diskCapacityBytes.Set(float64(dfs.capacity))

	if !skipinit {
		dfs.init()
		diskSizeBytes.Set(float64(dfs.Size()))
	}

	return dfs
}

func New(root string, capacity int64) *DiskFS {
	return new(root, capacity, false)
}

func NewSkipInit(root string, capacity int64) *DiskFS {
	return new(root, capacity, true)
}

func (d *DiskFS) init() {
	tstart := time.Now()

	err := filepath.Walk(d.root, func(npath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		d.mu.Lock()
		k := strings.ReplaceAll(npath[len(d.root)+1:], "\\", "/")
		fi := vfs.NewFileInfoFromOS(info, k)
		d.info[k] = fi
		d.LRU.Add(k, fi)
		d.size += info.Size()
		d.mu.Unlock()

		return nil
	})
	if err != nil {
		logger.Logger.Error().Err(err).Msg("Walk failed")
	}

	logger.Logger.Info().
		Str("name", d.Name()).
		Str("root", d.root).
		Str("capacity", units.HumanSize(float64(d.capacity))).
		Str("size", units.HumanSize(float64(d.Size()))).
		Str("files", fmt.Sprint(len(d.info))).
		Str("duration", time.Since(tstart).String()).
		Msg("init")
}

func (d *DiskFS) Capacity() int64 {
	return d.capacity
}

func (d *DiskFS) Name() string {
	return "DiskFS"
}

func (d *DiskFS) Size() int64 {
	d.mu.RLock()
	defer d.mu.RUnlock()
	return d.size
}

func (d *DiskFS) getKeyLock(key string) *sync.RWMutex {
	mu, _ := d.keyLocks.LoadOrStore(key, &sync.RWMutex{})
	return mu.(*sync.RWMutex)
}

func (d *DiskFS) Create(key string, size int64) (io.WriteCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	key = filepath.Clean(key)
	key = strings.ReplaceAll(key, "\\", "/") // Ensure forward slashes for consistency
	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	d.mu.RLock()
	if d.capacity > 0 {
		if d.size+size > d.capacity {
			d.mu.RUnlock()
			return nil, vfserror.ErrDiskFull
		}
	}
	d.mu.RUnlock()

	keyMu := d.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	// Check again after lock
	d.mu.Lock()
	if fi, exists := d.info[key]; exists {
		d.size -= fi.Size()
		d.LRU.Remove(key)
		delete(d.info, key)
		path := filepath.Join(d.root, key)
		os.Remove(path) // Ignore error, as file might not exist or other issues
	}
	d.mu.Unlock()

	path := filepath.Join(d.root, key)
	path = strings.ReplaceAll(path, "\\", "/") // Ensure forward slashes for consistency
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}

	file, err := os.Create(path)
	if err != nil {
		return nil, err
	}

	return &diskWriteCloser{
		Writer: file,
		onClose: func(n int64) error {
			fi, err := os.Stat(path)
			if err != nil {
				os.Remove(path)
				return err
			}

			d.mu.Lock()
			finfo := vfs.NewFileInfoFromOS(fi, key)
			d.info[key] = finfo
			d.LRU.Add(key, finfo)
			d.size += n
			d.mu.Unlock()

			diskWriteBytes.Add(float64(n))
			diskSizeBytes.Set(float64(d.Size()))

			// On new file, set access count to 1
			finfo.AccessCount = 1
			finfo.CTime = time.Now()
			d.fileMeta[key] = struct {
				AccessCount int64
				CTime       int64
			}{1, finfo.CTime.Unix()}
			flushFileMeta(d.root, d.fileMeta)

			return nil
		},
		key:  key,
		file: file,
	}, nil
}

type diskWriteCloser struct {
	io.Writer
	onClose func(int64) error
	n       int64
	key     string
	file    *os.File
}

func (wc *diskWriteCloser) Write(p []byte) (int, error) {
	n, err := wc.Writer.Write(p)
	wc.n += int64(n)
	return n, err
}

func (wc *diskWriteCloser) Close() error {
	err := wc.file.Close()
	if e := wc.onClose(wc.n); e != nil {
		os.Remove(wc.file.Name())
		return e
	}
	return err
}

// Delete deletes the value of key.
func (d *DiskFS) Delete(key string) error {
	if key == "" {
		return vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	key = filepath.Clean(key)
	key = strings.ReplaceAll(key, "\\", "/") // Ensure forward slashes for consistency
	if strings.Contains(key, "..") {
		return vfserror.ErrInvalidKey
	}

	keyMu := d.getKeyLock(key)
	keyMu.Lock()
	defer keyMu.Unlock()

	d.mu.Lock()
	fi, exists := d.info[key]
	if !exists {
		d.mu.Unlock()
		return vfserror.ErrNotFound
	}
	d.size -= fi.Size()
	d.LRU.Remove(key)
	delete(d.info, key)
	d.mu.Unlock()

	path := filepath.Join(d.root, key)
	path = strings.ReplaceAll(path, "\\", "/") // Ensure forward slashes for consistency
	if err := os.Remove(path); err != nil {
		return err
	}

	diskSizeBytes.Set(float64(d.Size()))

	delete(d.accessCounts, key)
	delete(d.fileMeta, key)
	flushFileMeta(d.root, d.fileMeta)

	return nil
}

// Open opens the file at key and returns it.
func (d *DiskFS) Open(key string) (io.ReadCloser, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	key = filepath.Clean(key)
	key = strings.ReplaceAll(key, "\\", "/") // Ensure forward slashes for consistency
	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := d.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	d.mu.Lock()
	fi, exists := d.info[key]
	if !exists {
		d.mu.Unlock()
		return nil, vfserror.ErrNotFound
	}
	fi.ATime = time.Now()
	fi.AccessCount++ // Increment access count
	d.LRU.MoveToFront(key)
	d.mu.Unlock()

	fi.AccessCount++
	d.fileMeta[key] = struct {
		AccessCount int64
		CTime       int64
	}{fi.AccessCount, fi.CTime.Unix()}
	flushFileMeta(d.root, d.fileMeta)

	path := filepath.Join(d.root, key)
	path = strings.ReplaceAll(path, "\\", "/") // Ensure forward slashes for consistency
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	// Update metrics on close
	return &readCloser{
		ReadCloser: file,
		onClose: func(n int64) {
			diskReadBytes.Add(float64(n))
		},
	}, nil
}

type readCloser struct {
	io.ReadCloser
	onClose func(int64)
	n       int64
}

func (rc *readCloser) Read(p []byte) (int, error) {
	n, err := rc.ReadCloser.Read(p)
	rc.n += int64(n)
	return n, err
}

func (rc *readCloser) Close() error {
	err := rc.ReadCloser.Close()
	rc.onClose(rc.n)
	return err
}

// Stat returns the FileInfo of key. If key is not found in the cache, it will stat the file on disk. If the file is not found on disk, it will return vfs.ErrNotFound.
func (d *DiskFS) Stat(key string) (*vfs.FileInfo, error) {
	if key == "" {
		return nil, vfserror.ErrInvalidKey
	}
	if key[0] == '/' {
		return nil, vfserror.ErrInvalidKey
	}

	// Sanitize key to prevent path traversal
	key = filepath.Clean(key)
	key = strings.ReplaceAll(key, "\\", "/") // Ensure forward slashes for consistency
	if strings.Contains(key, "..") {
		return nil, vfserror.ErrInvalidKey
	}

	keyMu := d.getKeyLock(key)
	keyMu.RLock()
	defer keyMu.RUnlock()

	d.mu.RLock()
	defer d.mu.RUnlock()

	if fi, ok := d.info[key]; !ok {
		return nil, vfserror.ErrNotFound
	} else {
		return fi, nil
	}
}

func (d *DiskFS) StatAll() []*vfs.FileInfo {
	d.mu.RLock()
	defer d.mu.RUnlock()

	// hard copy the file info to prevent modification of the original file info or the other way around
	files := make([]*vfs.FileInfo, 0, len(d.info))
	for _, v := range d.info {
		fi := *v
		files = append(files, &fi)
	}

	return files
}

func flushAccessCounts(root string, counts map[string]int64) {
	path := filepath.Join(root, "access_counts.json")
	data, _ := json.MarshalIndent(counts, "", "  ")
	_ = ioutil.WriteFile(path, data, 0644)
}

func flushFileMeta(root string, meta map[string]struct {
	AccessCount int64
	CTime       int64
}) {
	path := filepath.Join(root, "filemeta.json")
	if len(meta) == 0 {
		_ = os.Remove(path)
		return
	}
	data, _ := json.MarshalIndent(meta, "", "  ")
	_ = ioutil.WriteFile(path, data, 0644)
}