// SPDX-License-Identifier: GPL-2.0 or MIT
/* Copyright 2025 ARM Limited. All rights reserved. */

#include <linux/platform_device.h>
#include <linux/interrupt.h>
#include <linux/cleanup.h>
#include <linux/iopoll.h>
#include <linux/wait.h>

#include <drm/drm_managed.h>
#include <drm/drm_print.h>

#include "panthor_device.h"
#include "panthor_hw.h"
#include "panthor_pwr.h"
#include "panthor_regs.h"

#define PWR_INTERRUPTS_MASK \
        (PWR_IRQ_POWER_CHANGED_SINGLE | \
         PWR_IRQ_POWER_CHANGED_ALL | \
         PWR_IRQ_DELEGATION_CHANGED | \
         PWR_IRQ_RESET_COMPLETED | \
         PWR_IRQ_RETRACT_COMPLETED | \
         PWR_IRQ_INSPECT_COMPLETED | \
         PWR_IRQ_COMMAND_NOT_ALLOWED | \
         PWR_IRQ_COMMAND_INVALID)

#define PWR_ALL_CORES_MASK              GENMASK_U64(63, 0)

#define PWR_DOMAIN_MAX_BITS             16

#define PWR_TRANSITION_TIMEOUT_US       (2ULL * USEC_PER_SEC)

#define PWR_RETRACT_TIMEOUT_US          (2ULL * USEC_PER_MSEC)

#define PWR_RESET_TIMEOUT_MS            500

/**
 * struct panthor_pwr - PWR_CONTROL block management data.
 */
struct panthor_pwr {
        /** @irq: PWR irq. */
        struct panthor_irq irq;

        /** @reqs_lock: Lock protecting access to pending_reqs. */
        spinlock_t reqs_lock;

        /** @pending_reqs: Pending PWR requests. */
        u32 pending_reqs;

        /** @reqs_acked: PWR request wait queue. */
        wait_queue_head_t reqs_acked;
};

static void panthor_pwr_irq_handler(struct panthor_device *ptdev, u32 status)
{
        spin_lock(&ptdev->pwr->reqs_lock);
        gpu_write(ptdev, PWR_INT_CLEAR, status);

        if (unlikely(status & PWR_IRQ_COMMAND_NOT_ALLOWED))
                drm_err(&ptdev->base, "PWR_IRQ: COMMAND_NOT_ALLOWED");

        if (unlikely(status & PWR_IRQ_COMMAND_INVALID))
                drm_err(&ptdev->base, "PWR_IRQ: COMMAND_INVALID");

        if (status & ptdev->pwr->pending_reqs) {
                ptdev->pwr->pending_reqs &= ~status;
                wake_up_all(&ptdev->pwr->reqs_acked);
        }
        spin_unlock(&ptdev->pwr->reqs_lock);
}
PANTHOR_IRQ_HANDLER(pwr, PWR, panthor_pwr_irq_handler);

static void panthor_pwr_write_command(struct panthor_device *ptdev, u32 command, u64 args)
{
        if (args)
                gpu_write64(ptdev, PWR_CMDARG, args);

        gpu_write(ptdev, PWR_COMMAND, command);
}

static bool reset_irq_raised(struct panthor_device *ptdev)
{
        return gpu_read(ptdev, PWR_INT_RAWSTAT) & PWR_IRQ_RESET_COMPLETED;
}

static bool reset_pending(struct panthor_device *ptdev)
{
        return (ptdev->pwr->pending_reqs & PWR_IRQ_RESET_COMPLETED);
}

static int panthor_pwr_reset(struct panthor_device *ptdev, u32 reset_cmd)
{
        scoped_guard(spinlock_irqsave, &ptdev->pwr->reqs_lock) {
                if (reset_pending(ptdev)) {
                        drm_WARN(&ptdev->base, 1, "Reset already pending");
                } else {
                        ptdev->pwr->pending_reqs |= PWR_IRQ_RESET_COMPLETED;
                        gpu_write(ptdev, PWR_INT_CLEAR, PWR_IRQ_RESET_COMPLETED);
                        panthor_pwr_write_command(ptdev, reset_cmd, 0);
                }
        }

        if (!wait_event_timeout(ptdev->pwr->reqs_acked, !reset_pending(ptdev),
                                msecs_to_jiffies(PWR_RESET_TIMEOUT_MS))) {
                guard(spinlock_irqsave)(&ptdev->pwr->reqs_lock);

                if (reset_pending(ptdev) && !reset_irq_raised(ptdev)) {
                        drm_err(&ptdev->base, "RESET timed out (0x%x)", reset_cmd);
                        return -ETIMEDOUT;
                }

                ptdev->pwr->pending_reqs &= ~PWR_IRQ_RESET_COMPLETED;
        }

        return 0;
}

static const char *get_domain_name(u8 domain)
{
        switch (domain) {
        case PWR_COMMAND_DOMAIN_L2:
                return "L2";
        case PWR_COMMAND_DOMAIN_TILER:
                return "Tiler";
        case PWR_COMMAND_DOMAIN_SHADER:
                return "Shader";
        case PWR_COMMAND_DOMAIN_BASE:
                return "Base";
        case PWR_COMMAND_DOMAIN_STACK:
                return "Stack";
        }
        return "Unknown";
}

static u32 get_domain_base(u8 domain)
{
        switch (domain) {
        case PWR_COMMAND_DOMAIN_L2:
                return PWR_L2_PRESENT;
        case PWR_COMMAND_DOMAIN_TILER:
                return PWR_TILER_PRESENT;
        case PWR_COMMAND_DOMAIN_SHADER:
                return PWR_SHADER_PRESENT;
        case PWR_COMMAND_DOMAIN_BASE:
                return PWR_BASE_PRESENT;
        case PWR_COMMAND_DOMAIN_STACK:
                return PWR_STACK_PRESENT;
        }
        return 0;
}

static u32 get_domain_ready_reg(u32 domain)
{
        return get_domain_base(domain) + (PWR_L2_READY - PWR_L2_PRESENT);
}

static u32 get_domain_pwrtrans_reg(u32 domain)
{
        return get_domain_base(domain) + (PWR_L2_PWRTRANS - PWR_L2_PRESENT);
}

static bool is_valid_domain(u32 domain)
{
        return get_domain_base(domain) != 0;
}

static bool has_rtu(struct panthor_device *ptdev)
{
        return ptdev->gpu_info.gpu_features & GPU_FEATURES_RAY_TRAVERSAL;
}

static u8 get_domain_subdomain(struct panthor_device *ptdev, u32 domain)
{
        if (domain == PWR_COMMAND_DOMAIN_SHADER && has_rtu(ptdev))
                return PWR_COMMAND_SUBDOMAIN_RTU;

        return 0;
}

static int panthor_pwr_domain_wait_transition(struct panthor_device *ptdev, u32 domain,
                                              u32 timeout_us)
{
        u32 pwrtrans_reg = get_domain_pwrtrans_reg(domain);
        u64 val;
        int ret = 0;

        ret = gpu_read64_poll_timeout(ptdev, pwrtrans_reg, val, !(PWR_ALL_CORES_MASK & val), 100,
                                      timeout_us);
        if (ret) {
                drm_err(&ptdev->base, "%s domain power in transition, pwrtrans(0x%llx)",
                        get_domain_name(domain), val);
                return ret;
        }

        return 0;
}

static void panthor_pwr_debug_info_show(struct panthor_device *ptdev)
{
        drm_info(&ptdev->base, "GPU_FEATURES:    0x%016llx", gpu_read64(ptdev, GPU_FEATURES));
        drm_info(&ptdev->base, "PWR_STATUS:      0x%016llx", gpu_read64(ptdev, PWR_STATUS));
        drm_info(&ptdev->base, "L2_PRESENT:      0x%016llx", gpu_read64(ptdev, PWR_L2_PRESENT));
        drm_info(&ptdev->base, "L2_PWRTRANS:     0x%016llx", gpu_read64(ptdev, PWR_L2_PWRTRANS));
        drm_info(&ptdev->base, "L2_READY:        0x%016llx", gpu_read64(ptdev, PWR_L2_READY));
        drm_info(&ptdev->base, "TILER_PRESENT:   0x%016llx", gpu_read64(ptdev, PWR_TILER_PRESENT));
        drm_info(&ptdev->base, "TILER_PWRTRANS:  0x%016llx", gpu_read64(ptdev, PWR_TILER_PWRTRANS));
        drm_info(&ptdev->base, "TILER_READY:     0x%016llx", gpu_read64(ptdev, PWR_TILER_READY));
        drm_info(&ptdev->base, "SHADER_PRESENT:  0x%016llx", gpu_read64(ptdev, PWR_SHADER_PRESENT));
        drm_info(&ptdev->base, "SHADER_PWRTRANS: 0x%016llx", gpu_read64(ptdev, PWR_SHADER_PWRTRANS));
        drm_info(&ptdev->base, "SHADER_READY:    0x%016llx", gpu_read64(ptdev, PWR_SHADER_READY));
}

static int panthor_pwr_domain_transition(struct panthor_device *ptdev, u32 cmd, u32 domain,
                                         u64 mask, u32 timeout_us)
{
        u32 ready_reg = get_domain_ready_reg(domain);
        u32 pwr_cmd = PWR_COMMAND_DEF(cmd, domain, get_domain_subdomain(ptdev, domain));
        u64 expected_val = 0;
        u64 val;
        int ret = 0;

        if (drm_WARN_ON(&ptdev->base, !is_valid_domain(domain)))
                return -EINVAL;

        switch (cmd) {
        case PWR_COMMAND_POWER_DOWN:
                expected_val = 0;
                break;
        case PWR_COMMAND_POWER_UP:
                expected_val = mask;
                break;
        default:
                drm_err(&ptdev->base, "Invalid power domain transition command (0x%x)", cmd);
                return -EINVAL;
        }

        ret = panthor_pwr_domain_wait_transition(ptdev, domain, timeout_us);
        if (ret)
                return ret;

        /* domain already in target state, return early */
        if ((gpu_read64(ptdev, ready_reg) & mask) == expected_val)
                return 0;

        panthor_pwr_write_command(ptdev, pwr_cmd, mask);

        ret = gpu_read64_poll_timeout(ptdev, ready_reg, val, (mask & val) == expected_val, 100,
                                      timeout_us);
        if (ret) {
                drm_err(&ptdev->base,
                        "timeout waiting on %s power domain transition, cmd(0x%x), arg(0x%llx)",
                        get_domain_name(domain), pwr_cmd, mask);
                panthor_pwr_debug_info_show(ptdev);
                return ret;
        }

        return 0;
}

#define panthor_pwr_domain_power_off(__ptdev, __domain, __mask, __timeout_us)            \
        panthor_pwr_domain_transition(__ptdev, PWR_COMMAND_POWER_DOWN, __domain, __mask, \
                                      __timeout_us)

#define panthor_pwr_domain_power_on(__ptdev, __domain, __mask, __timeout_us) \
        panthor_pwr_domain_transition(__ptdev, PWR_COMMAND_POWER_UP, __domain, __mask, __timeout_us)

/**
 * retract_domain() - Retract control of a domain from MCU
 * @ptdev: Device.
 * @domain: Domain to retract the control
 *
 * Retracting L2 domain is not expected since it won't be delegated.
 *
 * Return: 0 on success or retracted already.
 *         -EPERM if domain is L2.
 *         A negative error code otherwise.
 */
static int retract_domain(struct panthor_device *ptdev, u32 domain)
{
        const u32 pwr_cmd = PWR_COMMAND_DEF(PWR_COMMAND_RETRACT, domain, 0);
        const u64 pwr_status = gpu_read64(ptdev, PWR_STATUS);
        const u64 delegated_mask = PWR_STATUS_DOMAIN_DELEGATED(domain);
        const u64 allow_mask = PWR_STATUS_DOMAIN_ALLOWED(domain);
        u64 val;
        int ret;

        if (drm_WARN_ON(&ptdev->base, domain == PWR_COMMAND_DOMAIN_L2))
                return -EPERM;

        ret = gpu_read64_poll_timeout(ptdev, PWR_STATUS, val, !(PWR_STATUS_RETRACT_PENDING & val),
                                      0, PWR_RETRACT_TIMEOUT_US);
        if (ret) {
                drm_err(&ptdev->base, "%s domain retract pending", get_domain_name(domain));
                return ret;
        }

        if (!(pwr_status & delegated_mask)) {
                drm_dbg(&ptdev->base, "%s domain already retracted", get_domain_name(domain));
                return 0;
        }

        panthor_pwr_write_command(ptdev, pwr_cmd, 0);

        /*
         * On successful retraction
         * allow-flag will be set with delegated-flag being cleared.
         */
        ret = gpu_read64_poll_timeout(ptdev, PWR_STATUS, val,
                                      ((delegated_mask | allow_mask) & val) == allow_mask, 10,
                                      PWR_TRANSITION_TIMEOUT_US);
        if (ret) {
                drm_err(&ptdev->base, "Retracting %s domain timeout, cmd(0x%x)",
                        get_domain_name(domain), pwr_cmd);
                return ret;
        }

        return 0;
}

/**
 * delegate_domain() - Delegate control of a domain to MCU
 * @ptdev: Device.
 * @domain: Domain to delegate the control
 *
 * Delegating L2 domain is prohibited.
 *
 * Return:
 * *       0 on success or delegated already.
 * *       -EPERM if domain is L2.
 * *       A negative error code otherwise.
 */
static int delegate_domain(struct panthor_device *ptdev, u32 domain)
{
        const u32 pwr_cmd = PWR_COMMAND_DEF(PWR_COMMAND_DELEGATE, domain, 0);
        const u64 pwr_status = gpu_read64(ptdev, PWR_STATUS);
        const u64 allow_mask = PWR_STATUS_DOMAIN_ALLOWED(domain);
        const u64 delegated_mask = PWR_STATUS_DOMAIN_DELEGATED(domain);
        u64 val;
        int ret;

        if (drm_WARN_ON(&ptdev->base, domain == PWR_COMMAND_DOMAIN_L2))
                return -EPERM;

        /* Already delegated, exit early */
        if (pwr_status & delegated_mask)
                return 0;

        /* Check if the command is allowed before delegating. */
        if (!(pwr_status & allow_mask)) {
                drm_warn(&ptdev->base, "Delegating %s domain not allowed", get_domain_name(domain));
                return -EPERM;
        }

        ret = panthor_pwr_domain_wait_transition(ptdev, domain, PWR_TRANSITION_TIMEOUT_US);
        if (ret)
                return ret;

        panthor_pwr_write_command(ptdev, pwr_cmd, 0);

        /*
         * On successful delegation
         * allow-flag will be cleared with delegated-flag being set.
         */
        ret = gpu_read64_poll_timeout(ptdev, PWR_STATUS, val,
                                      ((delegated_mask | allow_mask) & val) == delegated_mask,
                                      10, PWR_TRANSITION_TIMEOUT_US);
        if (ret) {
                drm_err(&ptdev->base, "Delegating %s domain timeout, cmd(0x%x)",
                        get_domain_name(domain), pwr_cmd);
                return ret;
        }

        return 0;
}

static int panthor_pwr_delegate_domains(struct panthor_device *ptdev)
{
        int ret;

        if (!ptdev->pwr)
                return 0;

        ret = delegate_domain(ptdev, PWR_COMMAND_DOMAIN_SHADER);
        if (ret)
                return ret;

        ret = delegate_domain(ptdev, PWR_COMMAND_DOMAIN_TILER);
        if (ret)
                goto err_retract_shader;

        return 0;

err_retract_shader:
        retract_domain(ptdev, PWR_COMMAND_DOMAIN_SHADER);

        return ret;
}

/**
 * panthor_pwr_domain_force_off - Forcefully power down a domain.
 * @ptdev: Device.
 * @domain: Domain to forcefully power down.
 *
 * This function will attempt to retract and power off the requested power
 * domain. However, if retraction fails, the operation is aborted. If power off
 * fails, the domain will remain retracted and under the host control.
 *
 * Return: 0 on success or a negative error code on failure.
 */
static int panthor_pwr_domain_force_off(struct panthor_device *ptdev, u32 domain)
{
        const u64 domain_ready = gpu_read64(ptdev, get_domain_ready_reg(domain));
        int ret;

        /* Domain already powered down, early exit. */
        if (!domain_ready)
                return 0;

        /* Domain has to be in host control to issue power off command. */
        ret = retract_domain(ptdev, domain);
        if (ret)
                return ret;

        return panthor_pwr_domain_power_off(ptdev, domain, domain_ready, PWR_TRANSITION_TIMEOUT_US);
}

void panthor_pwr_unplug(struct panthor_device *ptdev)
{
        unsigned long flags;

        if (!ptdev->pwr)
                return;

        /* Make sure the IRQ handler is not running after that point. */
        panthor_pwr_irq_suspend(&ptdev->pwr->irq);

        /* Wake-up all waiters. */
        spin_lock_irqsave(&ptdev->pwr->reqs_lock, flags);
        ptdev->pwr->pending_reqs = 0;
        wake_up_all(&ptdev->pwr->reqs_acked);
        spin_unlock_irqrestore(&ptdev->pwr->reqs_lock, flags);
}

int panthor_pwr_init(struct panthor_device *ptdev)
{
        struct panthor_pwr *pwr;
        int err, irq;

        if (!panthor_hw_has_pwr_ctrl(ptdev))
                return 0;

        pwr = drmm_kzalloc(&ptdev->base, sizeof(*pwr), GFP_KERNEL);
        if (!pwr)
                return -ENOMEM;

        spin_lock_init(&pwr->reqs_lock);
        init_waitqueue_head(&pwr->reqs_acked);
        ptdev->pwr = pwr;

        irq = platform_get_irq_byname(to_platform_device(ptdev->base.dev), "gpu");
        if (irq < 0)
                return irq;

        err = panthor_request_pwr_irq(ptdev, &pwr->irq, irq, PWR_INTERRUPTS_MASK);
        if (err)
                return err;

        return 0;
}

int panthor_pwr_reset_soft(struct panthor_device *ptdev)
{
        if (!(gpu_read64(ptdev, PWR_STATUS) & PWR_STATUS_ALLOW_SOFT_RESET)) {
                drm_err(&ptdev->base, "RESET_SOFT not allowed");
                return -EOPNOTSUPP;
        }

        return panthor_pwr_reset(ptdev, PWR_COMMAND_RESET_SOFT);
}

void panthor_pwr_l2_power_off(struct panthor_device *ptdev)
{
        const u64 l2_allow_mask = PWR_STATUS_DOMAIN_ALLOWED(PWR_COMMAND_DOMAIN_L2);
        const u64 pwr_status = gpu_read64(ptdev, PWR_STATUS);

        /* Abort if L2 power off constraints are not satisfied */
        if (!(pwr_status & l2_allow_mask)) {
                drm_warn(&ptdev->base, "Power off L2 domain not allowed");
                return;
        }

        /* It is expected that when halting the MCU, it would power down its
         * delegated domains. However, an unresponsive or hung MCU may not do
         * so, which is why we need to check and retract the domains back into
         * host control to be powered down in the right order before powering
         * down the L2.
         */
        if (panthor_pwr_domain_force_off(ptdev, PWR_COMMAND_DOMAIN_TILER))
                return;

        if (panthor_pwr_domain_force_off(ptdev, PWR_COMMAND_DOMAIN_SHADER))
                return;

        panthor_pwr_domain_power_off(ptdev, PWR_COMMAND_DOMAIN_L2, ptdev->gpu_info.l2_present,
                                     PWR_TRANSITION_TIMEOUT_US);
}

int panthor_pwr_l2_power_on(struct panthor_device *ptdev)
{
        const u32 pwr_status = gpu_read64(ptdev, PWR_STATUS);
        const u32 l2_allow_mask = PWR_STATUS_DOMAIN_ALLOWED(PWR_COMMAND_DOMAIN_L2);
        int ret;

        if ((pwr_status & l2_allow_mask) == 0) {
                drm_warn(&ptdev->base, "Power on L2 domain not allowed");
                return -EPERM;
        }

        ret = panthor_pwr_domain_power_on(ptdev, PWR_COMMAND_DOMAIN_L2, ptdev->gpu_info.l2_present,
                                          PWR_TRANSITION_TIMEOUT_US);
        if (ret)
                return ret;

        /* Delegate control of the shader and tiler power domains to the MCU as
         * it can better manage which shader/tiler cores need to be powered up
         * or can be powered down based on currently running jobs.
         *
         * If the shader and tiler domains are already delegated to the MCU,
         * this call would just return early.
         */
        return panthor_pwr_delegate_domains(ptdev);
}

void panthor_pwr_suspend(struct panthor_device *ptdev)
{
        if (!ptdev->pwr)
                return;

        panthor_pwr_irq_suspend(&ptdev->pwr->irq);
}

void panthor_pwr_resume(struct panthor_device *ptdev)
{
        if (!ptdev->pwr)
                return;

        panthor_pwr_irq_resume(&ptdev->pwr->irq, PWR_INTERRUPTS_MASK);
}