/*
   Copyright The containerd Authors.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package service

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync"
	"time"

	"github.com/containerd/stargz-snapshotter/analyzer/fanotify/conn"
	"golang.org/x/sys/unix"
)

// Serve starts fanotify at target directory and notifies all accessed files.
// Passed io.Reader an io.Writer are used for communicating with the client.
func Serve(target string, r io.Reader, w io.Writer) error {
	sConn := conn.NewService(r, w, 5*time.Second)

	fd, err := unix.FanotifyInit(unix.FAN_CLASS_NOTIF, unix.O_RDONLY)
	if err != nil {
		return fmt.Errorf("fanotify_init: %w", err)
	}

	// This blocks until the client tells us to start monitoring the target mountpoint.
	if err := sConn.WaitStart(); err != nil {
		return fmt.Errorf("waiting for start inst: %w", err)
	}

	// Start monitoring the target mountpoint.
	if err := unix.FanotifyMark(fd,
		unix.FAN_MARK_ADD|unix.FAN_MARK_MOUNT,
		unix.FAN_ACCESS|unix.FAN_OPEN,
		unix.AT_FDCWD,
		target,
	); err != nil {
		return fmt.Errorf("fanotify_mark: %w", err)
	}

	// Notify "started" state to the client.
	if err := sConn.SendStarted(); err != nil {
		return fmt.Errorf("failed to send started message: %w", err)
	}

	nr := bufio.NewReader(os.NewFile(uintptr(fd), ""))
	for {
		event := &unix.FanotifyEventMetadata{}
		if err := binary.Read(nr, binary.LittleEndian, event); err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("read fanotify fd: %w", err)
		}
		if event.Vers != unix.FANOTIFY_METADATA_VERSION {
			return fmt.Errorf("fanotify version mismatch %d(got) != %d(want)",
				event.Vers, unix.FANOTIFY_METADATA_VERSION)
		}
		if event.Fd < 0 {
			// queue overflow
			// TODO: do we need something special?
			fmt.Fprintf(os.Stderr, "Warn: queue overflow")
			continue
		}

		// Notify file descriptor.
		// NOTE: There is no guarantee that we have /proc in this mount namespace
		// (the target container's rootfs is mounted on "/") so we send file
		// descriptor and let the client resolve the path of this file using /proc of
		// this process.
		if err := sConn.SendFd(int(event.Fd)); err != nil {
			return fmt.Errorf("failed to send fd %d to client: %w", fd, err)
		}
		if err := unix.Close(int(event.Fd)); err != nil {
			return fmt.Errorf("Close(fd): %w", err)
		}

		continue
	}

	return nil
}

type PreContainerMonitor struct {
	targetDir string
	file      *os.File
	paths     map[string]struct{}
	mu        sync.RWMutex
	done      chan struct{}
	closed    bool
	closeMu   sync.Mutex
}

func NewPreContainerMonitor(targetDir string) *PreContainerMonitor {
	return &PreContainerMonitor{
		targetDir: targetDir,
		paths:     make(map[string]struct{}),
		done:      make(chan struct{}),
	}
}

func (m *PreContainerMonitor) Start() error {
	fd, err := unix.FanotifyInit(unix.FAN_CLASS_NOTIF, unix.O_RDONLY)
	if err != nil {
		return fmt.Errorf("fanotify_init: %w", err)
	}
	m.file = os.NewFile(uintptr(fd), "fanotify")

	if err := unix.FanotifyMark(fd,
		unix.FAN_MARK_ADD|unix.FAN_MARK_FILESYSTEM,
		unix.FAN_ACCESS|unix.FAN_OPEN,
		unix.AT_FDCWD,
		m.targetDir,
	); err != nil {
		m.file.Close()
		return fmt.Errorf("fanotify_mark: %w", err)
	}

	return nil
}

func (m *PreContainerMonitor) Monitor() error {
	if m.file == nil {
		return fmt.Errorf("monitor not started")
	}

	nr := bufio.NewReader(m.file)

	for {
		event := &unix.FanotifyEventMetadata{}
		if err := binary.Read(nr, binary.LittleEndian, event); err != nil {
			if err == io.EOF {
				return nil
			}
			select {
			case <-m.done:
				return nil
			default:
			}
			return fmt.Errorf("read fanotify fd: %w", err)
		}

		if event.Vers != unix.FANOTIFY_METADATA_VERSION {
			return fmt.Errorf("fanotify version mismatch %d(got) != %d(want)",
				event.Vers, unix.FANOTIFY_METADATA_VERSION)
		}

		if event.Fd < 0 {
			fmt.Fprintf(os.Stderr, "Warn: queue overflow")
			continue
		}

		path, err := os.Readlink(fmt.Sprintf("/proc/self/fd/%d", event.Fd))
		if err != nil {
			unix.Close(int(event.Fd))
			continue
		}

		m.mu.Lock()
		m.paths[path] = struct{}{}
		m.mu.Unlock()

		if err := unix.Close(int(event.Fd)); err != nil {
			fmt.Fprintf(os.Stderr, "Warn: failed to close fd %d: %v", event.Fd, err)
		}
	}
}

func (m *PreContainerMonitor) GetPaths() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	paths := make([]string, 0, len(m.paths))
	for path := range m.paths {
		paths = append(paths, path)
	}
	return paths
}

func (m *PreContainerMonitor) Close() error {
	m.closeMu.Lock()
	defer m.closeMu.Unlock()

	if m.closed {
		return nil
	}
	m.closed = true

	close(m.done)

	if m.file != nil {
		if err := m.file.Close(); err != nil {
			return fmt.Errorf("failed to close fanotify fd: %w", err)
		}
		m.file = nil
	}

	return nil
}