/*
 * Copyright (c) Souldbminer, Lightos and Horizon OC Contributors
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms and conditions of the GNU General Public License,
 * version 2, as published by the Free Software Foundation.
 *
 * This program is distributed in the hope 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 <atomic>
#include <cstddef>
#include <cstring>
#include <switch.h>

#include "../board/board.hpp"
#include "../file/config.hpp"
#include "../file/file_utils.hpp"
#include "../mapping/mem_map.hpp"
#include "../tsensor/soctherm.hpp"
#include "bpmp.hpp"

#include "bpmpfw_bin.h"

namespace bpmp {

    namespace {

        constexpr u64 FlowCtlrPhysBase = 0x60007000;
        constexpr u64 EvpPhysBase = 0x6000F000;

        constexpr u32 ClkRstRstDevLSet = 0x300;
        constexpr u32 ClkRstRstDevLClr = 0x304;
        constexpr u32 RstDevLCopRstBit = (1u << 1); /* SWR_COP_RST */

        constexpr u32 ApbMiscSlaveSecurityReg0 = 0xC00;
        constexpr u32 ApbMiscSlaveSecurityReg1 = 0xC04;
        constexpr u32 ApbMiscSlaveSecurityReg2 = 0xC08;

        /* Relative to actmonVirtAddr */
        constexpr u32 ActmonGlbPeriodCtrl = 0x004;
        constexpr u32 ActmonCopCtrl       = 0xC0;
        constexpr u32 ActmonCopUpperWmark = 0xC4;
        constexpr u32 ActmonCopIntrStatus = 0xE4;

        constexpr u32 ActmonGlbPeriodCtrlUsec  = (1u << 8);
        constexpr u32 ActmonCopCtrlStockArmed  = (1u << 31) | (1u << 30) | (1u << 18);

        /* MC SMMU handling for the BPMP. */
        constexpr u64 McPhysBase = 0x70019000;

        constexpr u32 McSmmuTlbConfig  = 0x014;
        constexpr u32 McSmmuPtbAsid    = 0x01C;
        constexpr u32 McSmmuPtbData    = 0x020;
        constexpr u32 McSmmuTlbFlush   = 0x030;
        constexpr u32 McSmmuPtcFlush   = 0x034;
        constexpr u32 McSmmuAvpcAsid   = 0x23C;
        constexpr u32 McSmmuPpcs1Asid  = 0x298;

        /* Page table is backed by the security engine's MMIO aperture. */
        constexpr u64 BpmpSmmuPdeBase = 0x70012000;
        constexpr u32 BpmpSmmuAsid    = 1;

        /* MC_SMMU_PTB_ASID */
        constexpr u32 SmmuPtbAsidCurrentAsid = BpmpSmmuAsid; /* [0:7] */

        /* MC_SMMU_PTB_DATA */
        constexpr u32 SmmuPtbDataPdeMask = 0x7FFFFF; /* [0:22] */

        /* MC_SMMU_AVPC_ASID / MC_SMMU_PPCS1_ASID */
        constexpr u32 SmmuClientSmmuEnableBit = (1u << 31);
        constexpr u32 SmmuAvpcAsidValue   = (SmmuClientSmmuEnableBit | BpmpSmmuAsid);
        constexpr u32 SmmuPpcs1AsidValue  = (SmmuClientSmmuEnableBit | BpmpSmmuAsid);

        constexpr u32 EvpCopResetVector         = 0x200;
        constexpr u32 EvpCopUndefVector         = 0x204;
        constexpr u32 EvpCopSwiVector           = 0x208;
        constexpr u32 EvpCopPrefetchAbortVector = 0x20C;
        constexpr u32 EvpCopDataAbortVector     = 0x210;
        constexpr u32 EvpCopRsvdVector          = 0x214;
        constexpr u32 EvpCopIrqVector           = 0x218;
        constexpr u32 EvpCopFiqVector           = 0x21C;

        constexpr u32 StockBpmpExceptionVector = 0x7D000000;

        constexpr u32 FlowCtlrHaltCopEvents = 0x004;
        constexpr u32 FlowModeNone          = 0x00000000;
        constexpr u32 FlowModeStop          = (2u << 29);

        constexpr u64 FwIramPhysBase = 0x40004000;
        constexpr u32 FwStagingSize  = 0x8000;
        constexpr u64 WorkRamPhysBase = FwIramPhysBase + FwStagingSize; // matches bpmpfw's WorkRamStart

        /* secmon will return error unless we apply thiss workaround */
        alignas(4096) u8 s_fwStageBuf[FwStagingSize];

        HocClkBpmpSharedInfo s_sharedInfo = {};

        inline volatile u32 &Mmio(u64 base_va, u32 offset) {
            return *reinterpret_cast<volatile u32 *>(base_va + offset);
        }

        // Late in the boot chain, otherwise a fatal will occur from invalid boot state
        constexpr u32 pscDependencies[] = { PscPmModuleId_Olsc };
        constexpr PscPmModuleId PscModuleId = (PscPmModuleId)705;

        PscPmModule s_pscModule;
        bool s_pscPrepared = false;
        Thread s_pscThread;
        bool s_pscExit = false;

        std::atomic_bool s_isAwake = true;

 
        /* Must be called to prevent secmon/erpt panic on wake */
        void RequestBpmpExitAndWait() {
            if (!IsPatchedExosphere()) {
                return;
            }

            const u32 one = 1;
            Result rc = SmcCopyToIram(WorkRamPhysBase + offsetof(HocClkBpmpSharedInfo, exitRequested), &one, sizeof(one));
            if (R_FAILED(rc)) {
                fileUtils::LogLine("[bpmp] RequestBpmpExitAndWait: SmcCopyToIram failed: 0x%x", rc);
                return;
            }

            constexpr int MaxAttempts = 20;
            constexpr u64 PollIntervalNs = 20'000'000; // 20ms

            HocClkBpmpSharedInfo info = {};
            for (int i = 0; i < MaxAttempts; i++) {
                svcSleepThread(PollIntervalNs);

                rc = SmcCopyFromIram(&info, WorkRamPhysBase, sizeof(info));
                if (R_SUCCEEDED(rc) && info.magic != HOCCLK_BPMP_MAGIC) {
                    return;
                }
            }

            fileUtils::LogLine("[bpmp] RequestBpmpExitAndWait: BPMP timeout (cannot sleep)");
        }

        void EnableBpmpSmmu() {
            /* Point the ASID page table at the security engine aperture. */
            const u32 pde_base = static_cast<u32>(BpmpSmmuPdeBase / 0x1000);

            /* Select the ASID and configure its page table entry. */
            SmcReadWriteRegister(McPhysBase + McSmmuPtbAsid, 0xFFFFFFFF, SmmuPtbAsidCurrentAsid);
            SmcReadWriteRegister(McPhysBase + McSmmuPtbData, 0xFFFFFFFF, pde_base & SmmuPtbDataPdeMask);

            /* Route BPMP (AVPC) and PPCS1 through the SMMU. */
            SmcReadWriteRegister(McPhysBase + McSmmuAvpcAsid,  0xFFFFFFFF, SmmuAvpcAsidValue);
            SmcReadWriteRegister(McPhysBase + McSmmuPpcs1Asid, 0xFFFFFFFF, SmmuPpcs1AsidValue);

            /* Flush the page table cache. */
            SmcReadWriteRegister(McPhysBase + McSmmuPtcFlush, 0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuTlbConfig, 0, 0);

            /* Flush the translation lookaside buffer. */
            SmcReadWriteRegister(McPhysBase + McSmmuTlbFlush, 0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuTlbConfig, 0, 0);
        }

        void DisableBpmpSmmu() {
            /* Clear the ASID page table. */
            SmcReadWriteRegister(McPhysBase + McSmmuPtbAsid,  0xFFFFFFFF, SmmuPtbAsidCurrentAsid);
            SmcReadWriteRegister(McPhysBase + McSmmuPtbData,  0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuAvpcAsid,  0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuPpcs1Asid, 0xFFFFFFFF, 0);

            /* Flush the page table cache. */
            SmcReadWriteRegister(McPhysBase + McSmmuPtcFlush, 0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuTlbConfig, 0, 0);

            /* Flush the translation lookaside buffer. */
            SmcReadWriteRegister(McPhysBase + McSmmuTlbFlush, 0xFFFFFFFF, 0);
            SmcReadWriteRegister(McPhysBase + McSmmuTlbConfig, 0, 0);
        }

        void RestoreStockState() {
            RequestBpmpExitAndWait();

            SmcReadWriteRegister(EvpPhysBase + EvpCopResetVector,         0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopUndefVector,         0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopSwiVector,           0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopPrefetchAbortVector, 0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopDataAbortVector,     0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopRsvdVector,          0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopIrqVector,           0xFFFFFFFF, StockBpmpExceptionVector);
            SmcReadWriteRegister(EvpPhysBase + EvpCopFiqVector,           0xFFFFFFFF, StockBpmpExceptionVector);

            SmcReadWriteRegister(FlowCtlrPhysBase + FlowCtlrHaltCopEvents, 0xFFFFFFFF, FlowModeStop);

            Mmio(board::actmonVirtAddr, ActmonGlbPeriodCtrl) = ActmonGlbPeriodCtrlUsec;
            Mmio(board::actmonVirtAddr, ActmonCopUpperWmark) = 0;
            Mmio(board::actmonVirtAddr, ActmonCopIntrStatus) = Mmio(board::actmonVirtAddr, ActmonCopIntrStatus);
            Mmio(board::actmonVirtAddr, ActmonCopCtrl) = ActmonCopCtrlStockArmed;

            EnableBpmpSmmu();
        }

        void PscThreadFunc(void *) {
            while (!s_pscExit) {
                Result rc = eventWait(&s_pscModule.event, 1'000'000'000ULL);
                if (R_FAILED(rc)) {
                    continue; // Timeout
                }

                PscPmState state;
                u32 flags;
                rc = pscPmModuleGetRequest(&s_pscModule, &state, &flags);
                if (R_SUCCEEDED(rc)) {
                    if (state == PscPmState_Awake) {
                        svcSleepThread(500'000'000); // 500ms

                        /* I have no clue why this is needed as it wasn't needed before ¯\(ツ)/¯ */
                        tsensor::InitializeSoctherm();

                        Result wrc = StartBpmfwExecution();
                        if (R_FAILED(wrc)) {
                            fileUtils::LogLine("[bpmp] restart after wake failed: 0x%x", wrc);
                        }

                        s_isAwake = true;
                    } else if (state == PscPmState_ReadySleep || state == PscPmState_ReadyShutdown) {
                        RestoreStockState();
                        s_isAwake = false;
                    }

                    pscPmModuleAcknowledge(&s_pscModule, state);
                }
            }
        }

    } // namespace

    Result StartBpmfwExecution() {
        if (!IsPatchedExosphere()) {
            fileUtils::LogLine("[bpmp] Cannot start without patched exosphere");
            return (Result)1;
        }

        const size_t fw_size = bpmpfw_bin_size;
        if (fw_size == 0 || fw_size > FwStagingSize) {
            fileUtils::LogLine("[bpmp] StartBpmfwExecution: bpmpfw of %zu bytes is invalid (should not happen)!", fw_size);
            return (Result)1;
        }

        std::memcpy(s_fwStageBuf, bpmpfw_bin, fw_size);

        fileUtils::LogLine("[bpmp] StartBpmfwExecution: loading %zu bytes at 0x%llx",
                            fw_size, static_cast<unsigned long long>(FwIramPhysBase));

        /* Disable the SMMU so DRAM can be accessed */
        DisableBpmpSmmu();

        /* Set BPMP reset */
        Mmio(board::clkVirtAddr, ClkRstRstDevLSet) = RstDevLCopRstBit;

        /* Grant BPMP access to all a bunch of registers that are useful */
        Mmio(board::apbVirtAddr, ApbMiscSlaveSecurityReg0) = 0;
        Mmio(board::apbVirtAddr, ApbMiscSlaveSecurityReg1) = 0;
        Mmio(board::apbVirtAddr, ApbMiscSlaveSecurityReg2) = 0;

        /* Set up exeption vectors for bpmp */
        const u32 fw_reset_vector = static_cast<u32>(FwIramPhysBase + 0x0);
        const u32 fw_panic_vector = static_cast<u32>(FwIramPhysBase + 0x4);

        SmcReadWriteRegister(EvpPhysBase + EvpCopResetVector,         0xFFFFFFFF, fw_reset_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopUndefVector,         0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopSwiVector,           0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopPrefetchAbortVector, 0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopDataAbortVector,     0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopRsvdVector,          0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopIrqVector,           0xFFFFFFFF, fw_panic_vector);
        SmcReadWriteRegister(EvpPhysBase + EvpCopFiqVector,           0xFFFFFFFF, fw_panic_vector);

        /* Disable actmon monitoring to avoid hang on FW >4.0.0 */
        Mmio(board::actmonVirtAddr, ActmonCopCtrl) = 0;
        Mmio(board::actmonVirtAddr, ActmonCopIntrStatus) = Mmio(board::actmonVirtAddr, ActmonCopIntrStatus);

        /* Load the BPMPFW into iram per page. TODO: determine if this is nessesary */
        constexpr size_t PageSize = 4096;
        for (size_t ofs = 0; ofs < fw_size; ofs += PageSize) {
            const size_t chunk = (fw_size - ofs) < PageSize ? (fw_size - ofs) : PageSize;

            Result copy_rc = SmcCopyToIram(FwIramPhysBase + ofs, s_fwStageBuf + ofs, static_cast<u32>(chunk));
            if (R_FAILED(copy_rc)) {
                fileUtils::LogLine("[bpmp] SmcCopyToIram failed at +0x%zx: 0x%x", ofs, copy_rc);
                return copy_rc;
            }
        }

        /* Clear BPMP reset. */
        Mmio(board::clkVirtAddr, ClkRstRstDevLClr) = RstDevLCopRstBit;

        /* Clear the halt written by bootloader. */
        SmcReadWriteRegister(FlowCtlrPhysBase + FlowCtlrHaltCopEvents, 0xFFFFFFFF, FlowModeNone);

        return 0;
    }

    void StartSleepMonitorThread() {
        if (!IsPatchedExosphere()) {
            fileUtils::LogLine("[bpmp] Cant start BPMP without exosphere patch");
            return;
        }

        Result rc = pscmInitialize();
        if (R_FAILED(rc)) {
            fileUtils::LogLine("[bpmp] pscmInitialize failed: 0x%x", rc);
            return;
        }

        rc = pscmGetPmModule(&s_pscModule, PscModuleId, pscDependencies, sizeof(pscDependencies) / sizeof(u32), true);
        if (R_FAILED(rc)) {
            fileUtils::LogLine("[bpmp] pscmGetPmModule failed: 0x%x", rc);
            pscmExit();
            return;
        }
        s_pscPrepared = true;

        rc = threadCreate(&s_pscThread, PscThreadFunc, nullptr, NULL, 0x1000, 0x10, 3);
        if (R_FAILED(rc)) {
            fileUtils::LogLine("[bpmp] failed to create psc thread: 0x%x", rc);
            return;
        }
        threadStart(&s_pscThread);
    }

    void StopSleepMonitorThread() {
        if (!s_pscPrepared) {
            return;
        }

        s_pscExit = true;
        eventFire(&s_pscModule.event);
        threadWaitForExit(&s_pscThread);
        threadClose(&s_pscThread);

        pscPmModuleFinalize(&s_pscModule);
        pscPmModuleClose(&s_pscModule);
        pscmExit();
        s_pscPrepared = false;
    }

    const HocClkBpmpSharedInfo *GetSharedInfo() {
        if (!IsPatchedExosphere()) {
            s_sharedInfo = {};
            return &s_sharedInfo;
        }

        SmcCopyFromIram(&s_sharedInfo, WorkRamPhysBase, sizeof(s_sharedInfo));
        return &s_sharedInfo;
    }

    bool IsAwake() {
        return s_isAwake;
    }

} // namespace bpmp