#include <linux/module.h>
#include <linux/fs.h>
#include <linux/init.h>
#include <linux/pagemap.h>
#include <linux/namei.h>
#include <linux/mount.h>
#include <linux/dcache.h>
#include <linux/slab.h>
#include <linux/hashtable.h>
#include <linux/version.h>

#define FS_NAME "customfs"
#define FS_MAGIC 0x20240615
#define MODE_BOOKS 01777
#define MODE_AUTHORS 0777
#define MODE_AUTHOR_SUBDIR 01777

MODULE_LICENSE("GPL");
MODULE_AUTHOR("you");
MODULE_DESCRIPTION("Custom FS with books/authors behavior");

static struct super_block *global_sb;
static struct inode *books_inode;
DEFINE_HASHTABLE(link_table, 6);

struct link_entry {
    struct hlist_node node;
    struct dentry *link;
    struct inode *target;
};

static ssize_t dummy_read(struct file *file, char __user *buf, size_t len, loff_t *ppos) {
    return 0;
}

static ssize_t dummy_write(struct file *file, const char __user *buf, size_t len, loff_t *ppos) {
    return len;
}

static const struct file_operations custom_file_ops = {
    .read = dummy_read,
    .write = dummy_write,
    .llseek = noop_llseek,
};

static bool is_protected_dentry(struct dentry *d) {
    return strcmp(d->d_name.name, "books") == 0 || strcmp(d->d_name.name, "authors") == 0;
}

static struct inode *get_inode(struct super_block *sb, umode_t mode) {
    struct inode *inode = new_inode(sb);
    if (!inode)
        return NULL;

    inode_init_owner(&init_user_ns, inode, NULL, mode);
    inode->i_ino = get_next_ino();
    inode->i_sb = sb;
    inode->i_atime = inode->i_mtime = inode->i_ctime = current_time(inode);

    if (S_ISDIR(mode)) {
        inode->i_op = &simple_dir_inode_operations;
        inode->i_fop = &simple_dir_operations;
        inc_nlink(inode);
    } else if (S_ISREG(mode)) {
        inode->i_fop = &custom_file_ops;
    }
    return inode;
}

static int books_create(struct user_namespace *ns, struct inode *dir,
                        struct dentry *dentry, umode_t mode, bool excl) {
    struct inode *inode = get_inode(dir->i_sb, S_IFREG | mode);
    if (!inode)
        return -ENOMEM;

    inode->i_uid = current_fsuid();
    d_add(dentry, inode);
    return 0;
}

static const struct inode_operations books_dir_iops = {
    .lookup = simple_lookup,
    .create = books_create,
};

static void remove_links_to_inode(struct inode *inode) {
    int bkt;
    struct link_entry *entry;
    struct hlist_node *tmp;
    hash_for_each_safe(link_table, bkt, tmp, entry, node) {
        if (entry->target == inode) {
            d_delete(entry->link);
            dput(entry->link);
            hash_del(&entry->node);
            kfree(entry);
            printk(KERN_INFO "customfs: removed link to inode %lu\n", inode->i_ino);
        }
    }
}

static int custom_mkdir(struct user_namespace *ns, struct inode *dir, struct dentry *dentry, umode_t mode) {
    if (dir == global_sb->s_root->d_inode)
        return -EPERM;

    if (S_ISDIR(dir->i_mode)) {
        umode_t mode_only = dir->i_mode & 07777;
        if ((mode_only == MODE_AUTHORS || mode_only == MODE_AUTHOR_SUBDIR) &&
            !uid_eq(current_fsuid(), dir->i_uid))
            return -EACCES;
    }

    struct inode *inode = get_inode(dir->i_sb, S_IFDIR | MODE_AUTHOR_SUBDIR);
    if (!inode)
        return -ENOMEM;

    inode->i_uid = current_fsuid();
    d_add(dentry, inode);
    inc_nlink(dir);
    return 0;
}

static int custom_unlink(struct inode *dir, struct dentry *dentry) {
    if (dir == global_sb->s_root->d_inode && is_protected_dentry(dentry))
        return -EPERM;
    return simple_unlink(dir, dentry);
}

static int custom_rmdir(struct inode *dir, struct dentry *dentry) {
    if (dir == global_sb->s_root->d_inode && is_protected_dentry(dentry))
        return -EPERM;
    return simple_rmdir(dir, dentry);
}

static int custom_rename(struct user_namespace *ns,
                         struct inode *old_dir, struct dentry *old_dentry,
                         struct inode *new_dir, struct dentry *new_dentry,
                         unsigned int flags) {
    if ((old_dir == global_sb->s_root->d_inode || new_dir == global_sb->s_root->d_inode) &&
        (is_protected_dentry(old_dentry) || is_protected_dentry(new_dentry)))
        return -EPERM;
    return simple_rename(ns, old_dir, old_dentry, new_dir, new_dentry, flags);
}

static int custom_link(struct dentry *old_dentry, struct inode *dir, struct dentry *new_dentry) {
    if (dir->i_sb != old_dentry->d_inode->i_sb)
        return -EXDEV;
    if (!S_ISREG(old_dentry->d_inode->i_mode))
        return -EPERM;

    int r = simple_link(old_dentry, dir, new_dentry);
    if (r == 0) {
        struct link_entry *e = kmalloc(sizeof(*e), GFP_KERNEL);
        if (!e)
            return -ENOMEM;
        e->target = d_inode(old_dentry);
        e->link = dget(new_dentry);
        hash_add(link_table, &e->node, e->target->i_ino);
    }
    return r;
}

static const struct inode_operations fs_dir_iops = {
    .lookup = simple_lookup,
    .mkdir = custom_mkdir,
    .unlink = custom_unlink,
    .rmdir = custom_rmdir,
    .rename = custom_rename,
    .link = custom_link,
};

static void evict_inode(struct inode *inode) {
    if (inode->i_sb == global_sb && inode->i_ino && S_ISREG(inode->i_mode))
        remove_links_to_inode(inode);
    clear_inode(inode);
}

static const struct super_operations fs_sops = {
    .statfs = simple_statfs,
    .drop_inode = generic_delete_inode,
    .evict_inode = evict_inode,
};

static int fill_super(struct super_block *sb, void *data, int silent) {
    struct inode *root_inode = get_inode(sb, S_IFDIR | 0755);
    if (!root_inode)
        return -ENOMEM;

    struct dentry *root_dentry = d_make_root(root_inode);
    if (!root_dentry)
        return -ENOMEM;

    sb->s_root = root_dentry;
    sb->s_magic = FS_MAGIC;
    sb->s_op = &fs_sops;
    global_sb = sb;

    // books
    books_inode = get_inode(sb, S_IFDIR | MODE_BOOKS);
    if (!books_inode)
        return -ENOMEM;
    books_inode->i_op = &books_dir_iops;
    books_inode->i_fop = &simple_dir_operations;
    struct dentry *books = d_alloc_name(root_dentry, "books");
    d_add(books, books_inode);

    // authors
    struct inode *authors_inode = get_inode(sb, S_IFDIR | MODE_AUTHORS);
    if (!authors_inode)
        return -ENOMEM;
    authors_inode->i_op = &fs_dir_iops;
    authors_inode->i_fop = &simple_dir_operations;
    struct dentry *authors = d_alloc_name(root_dentry, "authors");
    d_add(authors, authors_inode);

    return 0;
}

static struct dentry *mount_fs(struct file_system_type *type, int flags, const char *dev, void *data) {
    return mount_nodev(type, flags, data, fill_super);
}

static struct file_system_type fs_type = {
    .owner = THIS_MODULE,
    .name = FS_NAME,
    .mount = mount_fs,
    .kill_sb = kill_litter_super,
};

static int __init fs_init(void) {
    hash_init(link_table);
    return register_filesystem(&fs_type);
}

static void __exit fs_exit(void) {
    unregister_filesystem(&fs_type);
}

module_init(fs_init);
module_exit(fs_exit);