/* ext2.c - Read-only ext2 filesystem driver for Pactor64 */
#include "ext2.h"
#include "ata.h"

/* === Partition-aware block read wrapper === */
static uint32_t g_part_offset = 0;  /* Partition start LBA */

static int part_read(uint64_t lba, uint32_t count, void *buf) {
    return ata_block_read(lba + g_part_offset, count, buf);
}

/* Mount ext2 from ATA disk at given partition LBA offset */
int ext2_mount(ext2_fs_t *fs, uint32_t part_lba) {
    if (!ata_is_present()) {
        kprintf("[ext2] No ATA drive available\n");
        return -1;
    }
    g_part_offset = part_lba;
    kprintf("[ext2] Mounting ext2 from ATA (partition LBA %d)...\n", part_lba);
    return ext2_init(fs, part_read);
}

/* Temporary buffers for block reads */
static uint8_t block_buf[4096];
static uint8_t block_buf2[4096];

/* Read a single block into buf */
static int read_block_raw(ext2_fs_t *fs, uint32_t block, void *buf) {
    if (block == 0) {
        kmemset(buf, 0, fs->block_size);
        return 0;
    }
    uint32_t sectors_per_block = fs->block_size / 512;
    return fs->read_block((uint64_t)block * sectors_per_block, sectors_per_block, buf);
}

/* Read the superblock (always at byte 1024, sector 2) */
int ext2_init(ext2_fs_t *fs, block_read_fn read_fn) {
    kmemset(fs, 0, sizeof(ext2_fs_t));
    fs->read_block = read_fn;

    /* Read the superblock - it's at byte offset 1024 (sector 2) */
    uint8_t sb_buf[1024];
    if (read_fn(2, 2, sb_buf) != 0) {
        kprintf("[ext2] Failed to read superblock\n");
        return -1;
    }

    kmemcpy(&fs->sb, sb_buf, sizeof(ext2_superblock_t));

    if (fs->sb.s_magic != EXT2_SUPER_MAGIC) {
        kprintf("[ext2] Bad magic: 0x%x (expected 0xEF53)\n", fs->sb.s_magic);
        return -1;
    }

    /* Calculate block size: 1024 << s_log_block_size */
    fs->block_size = 1024U << fs->sb.s_log_block_size;
    fs->inodes_per_group = fs->sb.s_inodes_per_group;
    fs->inode_size = (fs->sb.s_rev_level >= 1) ? fs->sb.s_inode_size : 128;
    fs->groups_count = (fs->sb.s_blocks_count + fs->sb.s_blocks_per_group - 1)
                       / fs->sb.s_blocks_per_group;

    kprintf("[ext2] Magic OK, block_size=%d, inodes=%d, groups=%d\n",
            fs->block_size, fs->sb.s_inodes_count, fs->groups_count);
    kprintf("[ext2] Volume: %.16s\n", fs->sb.s_volume_name);

    /* Read block group descriptor table */
    static ext2_bgd_t bgdt_buf[64];
    fs->bgdt = bgdt_buf;

    uint32_t bgdt_block = (fs->block_size == 1024) ? 2 : 1;
    uint32_t bgdt_bytes = fs->groups_count * sizeof(ext2_bgd_t);
    uint32_t bgdt_blocks = (bgdt_bytes + fs->block_size - 1) / fs->block_size;

    uint8_t *bgdt_ptr = (uint8_t *)fs->bgdt;
    for (uint32_t i = 0; i < bgdt_blocks && i * fs->block_size < sizeof(bgdt_buf); i++) {
        if (read_block_raw(fs, bgdt_block + i, bgdt_ptr + i * fs->block_size) != 0) {
            kprintf("[ext2] Failed to read BGDT block %d\n", bgdt_block + i);
            return -1;
        }
    }

    kprintf("[ext2] Initialized successfully\n");
    return 0;
}

/* Read an inode from the inode table */
int ext2_read_inode(ext2_fs_t *fs, uint32_t inode_num, ext2_inode_t *inode) {
    if (inode_num == 0) return -1;

    uint32_t group = (inode_num - 1) / fs->inodes_per_group;
    uint32_t index = (inode_num - 1) % fs->inodes_per_group;
    if (group >= fs->groups_count) return -1;

    uint32_t offset = index * fs->inode_size;
    uint32_t block_offset = offset / fs->block_size;
    uint32_t byte_offset = offset % fs->block_size;
    uint32_t table_block = fs->bgdt[group].bg_inode_table + block_offset;

    if (read_block_raw(fs, table_block, block_buf) != 0) {
        kprintf("[ext2] Failed to read inode table block %d\n", table_block);
        return -1;
    }

    kmemcpy(inode, block_buf + byte_offset, sizeof(ext2_inode_t));
    return 0;
}

/* Get physical block number from file block index */
static uint32_t get_block_num(ext2_fs_t *fs, ext2_inode_t *inode, uint32_t file_block) {
    uint32_t ptrs_per_block = fs->block_size / 4;

    /* Direct blocks (0-11) */
    if (file_block < 12) return inode->i_block[file_block];
    file_block -= 12;

    /* Single indirect */
    if (file_block < ptrs_per_block) {
        uint32_t indirect = inode->i_block[12];
        if (indirect == 0) return 0;
        read_block_raw(fs, indirect, block_buf);
        return ((uint32_t *)block_buf)[file_block];
    }
    file_block -= ptrs_per_block;

    /* Double indirect */
    if (file_block < ptrs_per_block * ptrs_per_block) {
        uint32_t dindirect = inode->i_block[13];
        if (dindirect == 0) return 0;
        read_block_raw(fs, dindirect, block_buf);
        uint32_t *ptrs1 = (uint32_t *)block_buf;
        uint32_t idx1 = file_block / ptrs_per_block;
        uint32_t idx2 = file_block % ptrs_per_block;
        if (ptrs1[idx1] == 0) return 0;
        read_block_raw(fs, ptrs1[idx1], block_buf2);
        return ((uint32_t *)block_buf2)[idx2];
    }

    return 0; /* Triple indirect not implemented */
}

/* Read bytes from an inode at a given offset */
static int read_inode_data(ext2_fs_t *fs, ext2_inode_t *inode,
                           uint32_t offset, uint32_t size, void *buf) {
    uint32_t file_size = inode->i_size;
    uint8_t *dst = (uint8_t *)buf;
    uint32_t bs = fs->block_size;

    while (size > 0 && offset < file_size) {
        uint32_t file_block = offset / bs;
        uint32_t block_off = offset % bs;
        uint32_t to_read = bs - block_off;
        if (to_read > size) to_read = size;
        if (offset + to_read > file_size) to_read = file_size - offset;

        uint32_t phys_block = get_block_num(fs, inode, file_block);
        if (phys_block == 0) {
            kmemset(dst, 0, to_read);
        } else {
            if (read_block_raw(fs, phys_block, block_buf) != 0) return -1;
            kmemcpy(dst, block_buf + block_off, to_read);
        }
        dst += to_read;
        offset += to_read;
        size -= to_read;
    }
    return (int)(dst - (uint8_t *)buf);
}

/* Find a directory entry by name */
static int find_dirent(ext2_fs_t *fs, ext2_inode_t *dir_inode,
                       const char *name, uint32_t name_len,
                       uint32_t *out_inode, uint8_t *out_type) {
    uint32_t offset = 0;
    uint32_t dir_size = dir_inode->i_size;

    while (offset < dir_size) {
        ext2_dirent_t ent;
        if (read_inode_data(fs, dir_inode, offset, sizeof(ext2_dirent_t), &ent) < (int)sizeof(ext2_dirent_t))
            return -1;
        if (ent.rec_len == 0) return -1;

        if (ent.inode != 0 && ent.name_len == name_len) {
            char ent_name[EXT2_NAME_LEN + 1];
            if (read_inode_data(fs, dir_inode, offset + sizeof(ext2_dirent_t),
                                ent.name_len, ent_name) < (int)ent.name_len)
                return -1;
            ent_name[ent.name_len] = '\0';
            if (kstrncmp(ent_name, name, name_len) == 0) {
                *out_inode = ent.inode;
                *out_type = ent.file_type;
                return 0;
            }
        }
        offset += ent.rec_len;
    }
    return -1;
}

/* Parse a path and open the file */
int ext2_open(ext2_fs_t *fs, const char *path, ext2_file_t *file) {
    kmemset(file, 0, sizeof(ext2_file_t));
    file->fs = fs;

    uint32_t cur_ino = EXT2_ROOT_INO;
    ext2_inode_t cur_inode;

    if (ext2_read_inode(fs, cur_ino, &cur_inode) != 0) {
        kprintf("[ext2] Failed to read root inode\n");
        return -1;
    }

    const char *p = path;
    if (*p == '/') p++;

    while (*p) {
        const char *comp_start = p;
        uint32_t comp_len = 0;
        while (*p && *p != '/') { p++; comp_len++; }
        if (comp_len == 0) continue;

        if (!(cur_inode.i_mode & EXT2_S_IFDIR)) return -1;

        uint32_t child_ino;
        uint8_t child_type;
        if (find_dirent(fs, &cur_inode, comp_start, comp_len, &child_ino, &child_type) != 0) {
            kprintf("[ext2] Not found: %.*s\n", comp_len, comp_start);
            return -1;
        }

        cur_ino = child_ino;
        if (ext2_read_inode(fs, cur_ino, &cur_inode) != 0) return -1;
        if (*p == '/') p++;
    }

    file->inode = cur_inode;
    file->inode_num = cur_ino;
    file->size = cur_inode.i_size;
    file->position = 0;

    kprintf("[ext2] Opened inode %d, size=%d bytes\n", cur_ino, file->size);
    return 0;
}

/* Read from an open file */
int ext2_read(ext2_file_t *file, void *buf, uint32_t size) {
    if (file->position >= file->size) return 0;
    if (file->position + size > file->size) size = file->size - file->position;

    int bytes = read_inode_data(file->fs, &file->inode, file->position, size, buf);
    if (bytes > 0) file->position += bytes;
    return bytes;
}

void ext2_close(ext2_file_t *file) {
    file->position = 0;
    file->size = 0;
}

/* List directory entries via callback */
int ext2_readdir(ext2_fs_t *fs, uint32_t dir_inode, ext2_dirent_cb cb, void *ctx) {
    ext2_inode_t inode;
    if (ext2_read_inode(fs, dir_inode, &inode) != 0) return -1;
    if (!(inode.i_mode & EXT2_S_IFDIR)) return -1;

    uint32_t offset = 0;
    uint32_t dir_size = inode.i_size;

    while (offset < dir_size) {
        ext2_dirent_t ent;
        if (read_inode_data(fs, &inode, offset, sizeof(ext2_dirent_t), &ent) < (int)sizeof(ext2_dirent_t))
            return -1;
        if (ent.rec_len == 0) return -1;

        if (ent.inode != 0) {
            char name[EXT2_NAME_LEN + 1];
            if (read_inode_data(fs, &inode, offset + sizeof(ext2_dirent_t),
                                ent.name_len, name) < (int)ent.name_len)
                return -1;
            name[ent.name_len] = '\0';
            cb(name, ent.name_len, ent.inode, ent.file_type, ctx);
        }
        offset += ent.rec_len;
    }
    return 0;
}