/****************************************************************************
 * Copyright (C) 2021 Maschell
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 ****************************************************************************/
#include <malloc.h>
#include <utils/WiiUScreen.h>
#include <common/common.h>
#include <iosuhax.h>
#include <fs/FSUtils.h>
#include <utils/StringTools.h>
#include <utils/utils.h>
#include <WUD/content/WiiUDiscContentsHeader.h>
#include "WUDDumperState.h"

WUDDumperState::WUDDumperState(WUDDumperState::eDumpTargetFormat pTargetFormat, eDumpTarget pTargetDevice)
        : targetFormat(pTargetFormat), targetDevice(pTargetDevice) {
    this->sectorBufSize = READ_SECTOR_SIZE * READ_NUM_SECTORS;
    this->state = STATE_OPEN_ODD1;
}

WUDDumperState::~WUDDumperState() {
    if (this->oddFd >= 0) {
        IOSUHAX_FSA_RawClose(gFSAfd, oddFd);
    }
    free(sectorBuf);
    free(emptySector);
}

ApplicationState::eSubState WUDDumperState::update(Input *input) {
    if (this->state == STATE_ERROR) {
        if (entrySelected(input)) {
            return ApplicationState::SUBSTATE_RETURN;
        }
        return ApplicationState::SUBSTATE_RUNNING;
    }
    if (this->state == STATE_OPEN_ODD1) {
        if (this->currentSector > 0) {
            auto ret = IOSUHAX_FSA_RawOpen(gFSAfd, "/dev/odd01", &(this->oddFd));
            if (ret >= 0) {
                // continue!
                this->state = STATE_DUMP_DISC;
            } else {
                this->oddFd = -1;
                this->state = STATE_WAIT_USER_ERROR_CONFIRM;
            }
            return ApplicationState::SUBSTATE_RUNNING;
        }
        if (this->retryCount-- <= 0) {
            this->state = STATE_PLEASE_INSERT_DISC;
            return ApplicationState::SUBSTATE_RUNNING;
        }
        auto ret = IOSUHAX_FSA_RawOpen(gFSAfd, "/dev/odd01", &(this->oddFd));
        if (ret >= 0) {
            if (this->sectorBuf == nullptr) {
                this->sectorBuf = (void *) memalign(0x100, this->sectorBufSize);
                if (this->sectorBuf == nullptr) {
                    DEBUG_FUNCTION_LINE("ERROR_MALLOC_FAILED");
                    this->setError(ERROR_MALLOC_FAILED);
                    return ApplicationState::SUBSTATE_RUNNING;
                }
            }
            DEBUG_FUNCTION_LINE("Opened /dev/odd01 %d", this->oddFd);
            this->state = STATE_READ_DISC_INFO;
        }
    } else if (this->state == STATE_PLEASE_INSERT_DISC) {
        if (entrySelected(input)) {
            return SUBSTATE_RETURN;
        }
    } else if (this->state == STATE_READ_DISC_INFO) {
        if (IOSUHAX_FSA_RawRead(gFSAfd, this->sectorBuf, READ_SECTOR_SIZE, 1, 0, this->oddFd) >= 0) {
            this->discId[10] = '\0';
            memcpy(this->discId.data(), sectorBuf, 10);
            if (this->discId[0] == 0) {
                setError(ERROR_NO_DISC_ID);
                return ApplicationState::SUBSTATE_RUNNING;
            }
            this->state = STATE_READ_DISC_INFO_DONE;
            return ApplicationState::SUBSTATE_RUNNING;
        }

        this->setError(ERROR_READ_FIRST_SECTOR);
        return ApplicationState::SUBSTATE_RUNNING;
    } else if (this->state == STATE_READ_DISC_INFO_DONE) {
        this->state = STATE_DUMP_DISC_KEY;
    } else if (this->state == STATE_DUMP_DISC_KEY) {
        // Read the WiiUDiscContentsHeader to determine if we need disckey and if it's the correct one.
        auto res = IOSUHAX_FSA_RawRead(gFSAfd, this->sectorBuf, READ_SECTOR_SIZE, 1, 3, this->oddFd);
        uint8_t discKey[16];
        bool hasDiscKey = false;
        if (res >= 0) {
            if (((uint32_t *) this->sectorBuf)[0] != WiiUDiscContentsHeader::MAGIC) {
                auto discKeyRes = IOSUHAX_ODM_GetDiscKey(discKey);
                if (discKeyRes >= 0) {
                    hasDiscKey = true;
                }
            }
        }

        if (hasDiscKey) {
            if (!FSUtils::CreateSubfolder(StringTools::fmt("%swudump/%s", getPathForDevice(targetDevice).c_str(), discId))) {
                setError(ERROR_WRITE_FAILED);
                return SUBSTATE_RUNNING;
            }
            if (!FSUtils::saveBufferToFile(StringTools::fmt("%swudump/%s/game.key", getPathForDevice(targetDevice).c_str(), discId), discKey, 16)) {
                setError(ERROR_WRITE_FAILED);
                return SUBSTATE_RUNNING;
            }
        }
        this->state = STATE_DUMP_DISC_START;
    } else if (this->state == STATE_DUMP_DISC_START) {
        if (!FSUtils::CreateSubfolder(StringTools::fmt("%swudump/%s", getPathForDevice(targetDevice).c_str(), discId))) {
            setError(ERROR_WRITE_FAILED);
            return ApplicationState::SUBSTATE_RUNNING;
        }
        if (targetFormat == DUMP_AS_WUX) {
            this->fileHandle = std::make_unique<WUXFileWriter>(StringTools::fmt("%swudump/%s/game.wux", getPathForDevice(targetDevice).c_str(), discId), READ_SECTOR_SIZE * WRITE_BUFFER_NUM_SECTORS,
                                                               SECTOR_SIZE, targetDevice == TARGET_SD);
        } else {
            this->fileHandle = std::make_unique<WUDFileWriter>(StringTools::fmt("%swudump/%s/game.wud", getPathForDevice(targetDevice).c_str(), discId), READ_SECTOR_SIZE * WRITE_BUFFER_NUM_SECTORS,
                                                               SECTOR_SIZE, targetDevice == TARGET_SD);
        }
        if (!this->fileHandle->isOpen()) {
            DEBUG_FUNCTION_LINE("Failed to open file");
            this->setError(ERROR_FILE_OPEN_FAILED);
            return ApplicationState::SUBSTATE_RUNNING;
        }

        this->startTime = OSGetTime();

        this->state = STATE_DUMP_DISC;
        this->totalSectorCount = (WUD_FILE_SIZE / SECTOR_SIZE);
        this->currentSector = 0;
        this->writtenSectors = 0;
        this->retryCount = 10;
    } else if (this->state == STATE_DUMP_DISC) {
        size_t numSectors = this->currentSector + READ_NUM_SECTORS > this->totalSectorCount ? this->totalSectorCount - this->currentSector : READ_NUM_SECTORS;
        if ((this->readResult = IOSUHAX_FSA_RawRead(gFSAfd, sectorBuf, READ_SECTOR_SIZE, numSectors, this->currentSector, this->oddFd)) >= 0) {
            auto curWrittenSectors = fileHandle->writeSector((const uint8_t *) this->sectorBuf, numSectors);
            if (curWrittenSectors < 0) {
                this->setError(ERROR_WRITE_FAILED);
                return ApplicationState::SUBSTATE_RUNNING;
            }
            currentSector += numSectors;
            this->writtenSectors += curWrittenSectors;

            this->retryCount = 10;
            if (this->currentSector >= this->totalSectorCount) {
                this->state = STATE_DUMP_DISC_DONE;
                if (this->fileHandle->isOpen()) {
                    if (!this->fileHandle->flush()) {
                        DEBUG_FUNCTION_LINE("Flush failed");
                        this->setError(ERROR_WRITE_FAILED);
                        return ApplicationState::SUBSTATE_RUNNING;
                    }
                    this->fileHandle->finalize();
                    this->fileHandle->close();
                }
            }
        } else {
            this->state = STATE_WAIT_USER_ERROR_CONFIRM;
            if (this->oddFd >= 0) {
                IOSUHAX_FSA_RawClose(gFSAfd, this->oddFd);
                this->oddFd = -1;
            }
            return ApplicationState::SUBSTATE_RUNNING;
        }
    } else if (this->state == STATE_WAIT_USER_ERROR_CONFIRM) {
        if (this->autoSkipOnError) {
            if (this->oddFd >= 0) {
                IOSUHAX_FSA_RawClose(gFSAfd, this->oddFd);
                this->oddFd = -1;
            }
        }
        if (this->autoSkipOnError || (input->data.buttons_d & Input::BUTTON_A)) {
            // this->log.fwrite("Skipped sector %d : 0x%ll016X-0x%ll016X, filled with 0's\n", this->currentSector, this->currentSector * READ_SECTOR_SIZE, (this->currentSector + 1) * READ_SECTOR_SIZE);
            this->state = STATE_OPEN_ODD1;
            this->skippedSectors.push_back(this->currentSector);
            // We can't use seek because we may have cached values.

            if (this->emptySector == nullptr) {
                this->emptySector = memalign(0x100, READ_SECTOR_SIZE);
                if (this->emptySector == nullptr) {
                    this->setError(ERROR_MALLOC_FAILED);
                    return ApplicationState::SUBSTATE_RUNNING;
                }
            }

            auto curWrittenSectors = fileHandle->writeSector((uint8_t *) emptySector, 1);
            if (curWrittenSectors < 0) {
                this->setError(ERROR_WRITE_FAILED);
                return ApplicationState::SUBSTATE_RUNNING;
            }

            this->currentSector += 1;
            this->writtenSectors += curWrittenSectors;
            this->readResult = 0;
        } else if (input->data.buttons_d & Input::BUTTON_B) {
            this->state = STATE_OPEN_ODD1;
            this->readResult = 0;
        } else if (input->data.buttons_d & Input::BUTTON_Y) {
            this->autoSkipOnError = true;
        }
    } else if (this->state == STATE_DUMP_DISC_DONE) {
        WiiUScreen::drawLinef("Dumping done! Press A to continue");
        if (entrySelected(input)) {
            return SUBSTATE_RETURN;
        }

        return ApplicationState::SUBSTATE_RUNNING;
    }
    return ApplicationState::SUBSTATE_RUNNING;
}

void WUDDumperState::render() {
    WiiUScreen::clearScreen();
    ApplicationState::printHeader();
    if (this->state == STATE_ERROR) {
        WiiUScreen::drawLine();
        WiiUScreen::drawLinef("Error:       %s", ErrorMessage().c_str());
        WiiUScreen::drawLinef("Description: %s", ErrorDescription().c_str());
        WiiUScreen::drawLine();
        WiiUScreen::drawLine("Press A to return.");
    } else if (this->state == STATE_OPEN_ODD1) {
        WiiUScreen::drawLine("Open /dev/odd01");
    } else if (this->state == STATE_PLEASE_INSERT_DISC) {
        WiiUScreen::drawLine("Please insert a Wii U disc and try again.\n\nPress A to return");
    } else if (this->state == STATE_DUMP_DISC_KEY) {
        WiiUScreen::drawLine("Read disc key");
    } else if (this->state == STATE_READ_DISC_INFO) {
        WiiUScreen::drawLine("Read disc information");
    } else if (this->state == STATE_READ_DISC_INFO_DONE) {
        WiiUScreen::drawLinef("Dumping: %s", this->discId);
    } else if (this->state == STATE_DUMP_DISC_START || this->state == STATE_DUMP_DISC || this->state == STATE_WAIT_USER_ERROR_CONFIRM) {
        WiiUScreen::drawLinef("Dumping: %s", this->discId);

        float percent = this->currentSector / (WUD_FILE_SIZE / READ_SECTOR_SIZE * 1.0f) * 100.0f;
        WiiUScreen::drawLinef("Progress: %0.2f MiB / %5.2f MiB (%2.1f %%)", this->currentSector * (READ_SECTOR_SIZE / 1024.0f / 1024.0f), WUD_FILE_SIZE / 1024.0f / 1024.0f, percent);
        if (targetFormat == DUMP_AS_WUX) {
            WiiUScreen::drawLinef("Written %0.2f MiB. Compression ratio 1:%0.2f", writtenSectors * (READ_SECTOR_SIZE / 1024.0f / 1024.0f),
                                  1.0f / (writtenSectors / (float) this->currentSector));
        }

        if (this->readResult < 0 || this->oddFd < 0) {
            WiiUScreen::drawLine();

            if (this->oddFd < 0) {
                WiiUScreen::drawLine("Failed to open disc, try again.");
            } else {
                WiiUScreen::drawLinef("Error: Failed to read sector - Error %d", this->readResult);
            }
            WiiUScreen::drawLine();
            WiiUScreen::drawLine("Press A to skip this sector (will be replaced by 0's)");
            WiiUScreen::drawLine("Press B to try again");
        } else {
            OSTime curTime = OSGetTime();
            float remaining = (WUD_FILE_SIZE - (READ_SECTOR_SIZE * this->currentSector)) / 1024.0f / 1024.0f;
            float curSpeed = READ_SECTOR_SIZE * ((this->currentSector / 1000.0f) / OSTicksToMilliseconds(curTime - startTime));
            int32_t remainingSec = remaining / curSpeed;
            int32_t minutes = (remainingSec / 60) % 60;
            int32_t seconds = remainingSec % 60;
            int32_t hours = remainingSec / 3600;

            WiiUScreen::drawLinef("Speed: %.2f MiB/s ETA: %02dh %02dm %02ds", curSpeed, remaining, hours, minutes, seconds);
        }

        WiiUScreen::drawLine();
        if (!this->skippedSectors.empty()) {
            WiiUScreen::drawLinef("Skipped dumping %d sectors", this->skippedSectors.size());
        }
    } else if (this->state == STATE_DUMP_DISC_DONE) {
        WiiUScreen::drawLinef("Dumping done! Press A to continue");
    }

    ApplicationState::printFooter();
    WiiUScreen::flipBuffers();
}

void WUDDumperState::setError(WUDDumperState::eErrorState err) {
    this->state = STATE_ERROR;
    this->errorState = err;
    //OSEnableHomeButtonMenu(true);
}

std::string WUDDumperState::ErrorMessage() const {
    if (this->errorState == ERROR_READ_FIRST_SECTOR) {
        return "ERROR_READ_FIRST_SECTOR";
    } else if (this->errorState == ERROR_FILE_OPEN_FAILED) {
        return "ERROR_FILE_OPEN_FAILED";
    } else if (this->errorState == ERROR_MALLOC_FAILED) {
        return "ERROR_MALLOC_FAILED";
    } else if (this->errorState == ERROR_NO_DISC_ID) {
        return "ERROR_NO_DISC_ID";
    } else if (this->errorState == ERROR_WRITE_FAILED) {
        return "ERROR_WRITE_FAILED";
    }
    return "UNKNOWN_ERROR";
}

std::string WUDDumperState::ErrorDescription() const {
    if (this->errorState == ERROR_READ_FIRST_SECTOR) {
        return "Failed to read first sector.";
    } else if (this->errorState == ERROR_MALLOC_FAILED) {
        return "Failed to allocate data.";
    } else if (this->errorState == ERROR_FILE_OPEN_FAILED) {
        return "Failed to create file";
    } else if (this->errorState == ERROR_NO_DISC_ID) {
        return "Failed to get the disc id";
    } else if (this->errorState == ERROR_WRITE_FAILED) {
        return "Failed to write the file. \nMake sure to have enough space on the storage";
    }
    return "UNKNOWN_ERROR";
}

std::string WUDDumperState::getPathForDevice(eDumpTarget target) const {
    if (target == TARGET_NTFS) {
        return "ntfs0:/";
    }
    return "fs:/vol/external01/";
}