#include <ogcsys.h>
#include <malloc.h>
#include <string.h>

#include "usbloader/sdhc.h"
#include "usbloader/usbstorage2.h"
#include "usbloader/wdvd.h"
#include "wbfs_rw.h"

/* Constants */
#define MAX_NB_SECTORS  32

u32 sector_size = 512;
rw_sector_callback_t readCallback  = NULL;
rw_sector_callback_t writeCallback = NULL;

void SetSectorSize( u32 size )
{
    sector_size = size;
}

s32 __ReadDVD( void *fp, u32 lba, u32 len, void *iobuf )
{
    void *buffer = NULL;

    u64 offset;
    u32 mod, size;
    s32 ret;

    /* Calculate offset */
    offset = ( ( u64 )lba ) << 2;

    /* Calcualte sizes */
    mod  = len % 32;
    size = len - mod;

    /* Read aligned data */
    if ( size )
    {
        ret = WDVD_UnencryptedRead( iobuf, size, offset );
        if ( ret < 0 )
            goto out;
    }

    /* Read non-aligned data */
    if ( mod )
    {
        /* Allocate memory */
        buffer = memalign( 32, 0x20 );
        if ( !buffer )
            return -1;

        /* Read data */
        ret = WDVD_UnencryptedRead( buffer, 0x20, offset + size );
        if ( ret < 0 )
            goto out;

        /* Copy data */
        void *ptr = ( ( u8 * ) iobuf ) + size;
        memcpy( ptr, buffer, mod );
    }

    /* Success */
    ret = 0;

out:
    /* Free memory */
    if ( buffer )
        free( buffer );

    return ret;
}

s32 __ReadUSB( void *fp, u32 lba, u32 count, void *iobuf )
{
    u32 cnt = 0;
    s32 ret;

    /* Do reads */
    while ( cnt < count )
    {
        void *ptr     = ( ( u8 * )iobuf ) + ( cnt * sector_size );
        u32   sectors = ( count - cnt );

        /* Read sectors is too big */
        if ( sectors > MAX_NB_SECTORS )
            sectors = MAX_NB_SECTORS;

        /* USB read */
        ret = USBStorage2_ReadSectors( lba + cnt, sectors, ptr );
        if ( ret < 0 )
            return ret;

        /* Increment counter */
        cnt += sectors;
    }

    return 0;
}

s32 __WriteUSB( void *fp, u32 lba, u32 count, void *iobuf )
{
    u32 cnt = 0;
    s32 ret;

    /* Do writes */
    while ( cnt < count )
    {
        void *ptr     = ( ( u8 * )iobuf ) + ( cnt * sector_size );
        u32   sectors = ( count - cnt );

        /* Write sectors is too big */
        if ( sectors > MAX_NB_SECTORS )
            sectors = MAX_NB_SECTORS;

        /* USB write */
        ret = USBStorage2_WriteSectors( lba + cnt, sectors, ptr );
        if ( ret < 0 )
            return ret;

        /* Increment counter */
        cnt += sectors;
    }

    return 0;
}

s32 __ReadSDHC( void *fp, u32 lba, u32 count, void *iobuf )
{
    u32 cnt = 0;
    s32 ret;

    /* Do reads */
    while ( cnt < count )
    {
        void *ptr     = ( ( u8 * )iobuf ) + ( cnt * sector_size );
        u32   sectors = ( count - cnt );

        /* Read sectors is too big */
        if ( sectors > MAX_NB_SECTORS )
            sectors = MAX_NB_SECTORS;

        /* SDHC read */
        ret = SDHC_ReadSectors( lba + cnt, sectors, ptr );
        if ( !ret )
            return -1;

        /* Increment counter */
        cnt += sectors;
    }

    return 0;
}

s32 __WriteSDHC( void *fp, u32 lba, u32 count, void *iobuf )
{
    u32 cnt = 0;
    s32 ret;

    /* Do writes */
    while ( cnt < count )
    {
        void *ptr     = ( ( u8 * )iobuf ) + ( cnt * sector_size );
        u32   sectors = ( count - cnt );

        /* Write sectors is too big */
        if ( sectors > MAX_NB_SECTORS )
            sectors = MAX_NB_SECTORS;

        /* SDHC write */
        ret = SDHC_WriteSectors( lba + cnt, sectors, ptr );
        if ( !ret )
            return -1;

        /* Increment counter */
        cnt += sectors;
    }

    return 0;
}