//See Install Documentation.
//Study the account examples: Config Strs, Update Funs

#include "BootMHD"
#include "BootMHD2"

#help_index "Install"

#define BOOT_DIR                "/Boot"
//Stage 2 of master boot loader
#define BOOT_DIR_BOOTMHD2_BIN BOOT_DIR "/BootMHD2.BIN"
//Old master boot record
#define BOOT_DIR_OLDMBR_BIN     BOOT_DIR "/OldMBR.BIN"

public U0 BootMHDOldRead(U8 src_drive, U8 dst_drive, I64 size=1)
{//Reads MBR from disk drive containing src partition.
//Writes MBR file to dst BOOT_DIR.
// Takes optional arg 'size': count of 512-byte blocks.
    CBlkDev     *bd = Letter2BlkDev(src_drive);
    CDrive      *drive;
    U8          *mbr = MAlloc(BLK_SIZE * size);

    Drive(dst_drive);
    drive = Fs->cur_dv;

    if (drive->fs_type != FSt_REDSEA && drive->fs_type != FSt_FAT32)
        PrintErr("File System Not Supported\n");
    else
    {
//Bypass partition bounds-checking
        BlkDevLock(bd);
        AHCIAtaBlksRead(bd, mbr, 0, size);
        BlkDevUnlock(bd);

        Drive(dst_drive);
        DirMake(BOOT_DIR);
        FileWrite(BOOT_DIR_OLDMBR_BIN, mbr, size * BLK_SIZE);
    }
}

public U0 BootMHDOldWrite(U8 src_drive, U8 dst_drive)
{//Reads OldMBR from src drive BOOT_DIR.
//writes it to the MBR of the drive with dst partition.
// If OldMBR.BIN size > BLK_SIZE, writes to post-MBR gap.
    CBlkDev     *bd = Letter2BlkDev(dst_drive);
    CMasterBoot *mbr;
    I64          size;

    Drive(src_drive);

    if (mbr = FileRead(BOOT_DIR_OLDMBR_BIN, &size))
    {
//Bypass partition bounds-checking
        BlkDevLock(bd);
        AHCIAtaBlksWrite(bd, mbr, 0, size / BLK_SIZE);
        BlkDevUnlock(bd);
    }
    Free(mbr);
}

public U0 BootMHDZero(U8 dst_drive)
{//Set MBR of disk with dst partition to zero.

    //This is dangerous!!
    //The ZealOS partitioner doesn't play well
    //with other operating systems at this time and you need
    //to do this on a drive partitioned by ZealOS
    //if you wish to partition with another operating system.
    CBlkDev     *bd = Letter2BlkDev(dst_drive);
    CMasterBoot  mbr;

    MemSet(&mbr, 0, BLK_SIZE);
//Bypass partition bounds-checking
    BlkDevLock(bd);
    AHCIAtaBlksWrite(bd, &mbr, 0, 1);
    BlkDevUnlock(bd);
}

public Bool BootMHDIns(U8 drv_let, U8 *drv_list=NULL)
{//Create new MBR on the disk that has drv_let as a partition.
//Puts stage 2 in BOOT_DIR of drv_let.
    CBlkDev     *bd, *bd1;
    CDrive      *drive, *p1;
    CMasterBoot  mbr;
    CDirEntry    de;
    I64          i, j, size, *_q;
    U8          *menu_ptr, *ptr, ch, buf[STR_LEN];
    Bool         res = FALSE;

    try
    {
        if (drv_list)
        {
            StrCopy(buf, drv_list);
            StrUtil(buf, SUF_TO_UPPER);
        }
        else
        {
            j = 0;
            for (i = 'A'; i <= 'Z'; i++)
                buf[j++] = i;
            buf[j++] = 0;
        }

        Drive(drv_let);
        drive = Fs->cur_dv;

        if (drive->fs_type != FSt_REDSEA && drive->fs_type != FSt_FAT32)
            PrintErr("File System Not Supported\n");
        else
        {
            bd = drive->bd;

            if (!FileFind(BOOT_DIR_OLDMBR_BIN,, FUF_JUST_FILES))
                BootMHDOldRead(drv_let, drv_let);

            _q = BMHD2_BLK_ARRAY;
            MemSet(_q, 0, sizeof(I64) * 8);
            menu_ptr = BMHD2_BOOT_MESSAGE;
            StrPrint(menu_ptr, "\n\r\n\rZealOS Boot Loader\n\r\n\r");
            j = 0;

            "\nDetected boot drives:\n";
            "(Choice %16s)\n", "drv_offset";
            if (FileFind(BOOT_DIR_OLDMBR_BIN, &de, FUF_JUST_FILES))
            {
                Free(de.full_name);
                *_q++ = Clus2Blk(drive, de.clus);
                "Old Boot Record\n";
                CatPrint(menu_ptr, "0. Old Boot Record\n\r");
                j++;
            }

            ptr = buf;
            while (ch = *ptr++)
            {
                if ((p1 = Letter2Drive(ch, FALSE)) && (bd1 = p1->bd) && bd1 == bd)
                {
                    *_q = p1->drv_offset;
                    "Drive %C:%16X\n", Drive2Letter(p1), *_q;
                    CatPrint(menu_ptr, "%d. Drive %C\n\r", j++, Drive2Letter(p1));
                    _q++;
                }
            }
            CatPrint(menu_ptr, "\n\rSelection:");

            size = BMHD2_END - BMHD2_START;
            FileWrite(BOOT_DIR_BOOTMHD2_BIN, BMHD2_START, size);

            if (!FileFind(BOOT_DIR_BOOTMHD2_BIN, &de, FUF_JUST_FILES))
                "No Boot Loader Image\n";
            else
            {
                Free(de.full_name);

                *BMHD_BLK_COUNT(U16 *) = (size + BLK_SIZE - 1) >> BLK_SIZE_BITS;
                *BMHD_DAP_BLK(I64 *) = Clus2Blk(drive, de.clus);
//Bypass partition bounds-checking
                BlkDevLock(bd);
                AHCIAtaBlksRead(bd, &mbr, 0, 1);

                for (i = 0; i < BMHD_END - BMHD_CODE; i++)
                    mbr.code[i] = BMHD_CODE(U8 *)[i];

#assert sizeof(CMasterBoot.code) >= BMHD_END - BMHD_CODE

                for (; i < sizeof(CMasterBoot.code); i++)
                    mbr.code[i] = 0;

                if (!mbr.media_id)
                    mbr.media_id = RandU32;

                mbr.zero = 0;
                mbr.signature = 0xAA55;

                AHCIAtaBlksWrite(bd, &mbr, 0, 1);

                BlkDevUnlock(bd);
                res = TRUE;
            }
        }
        "\nZealOS Public Domain HDD MBR Boot Loader installed.\n\n";
    }
    catch
        PutExcept;
    return res;
}