/* SPDX-License-Identifier: GPL-2.0-only */
/*
 * GSP vGPU-IO Driver - Live Migration
 */

#include "gsp_core.h"
#include <linux/slab.h>
#include <linux/wait.h>
#include <linux/kthread.h>
#include <linux/delay.h>

/* Forward declarations */
static int gsp_migration_prepare(struct gsp_migration *mig);
static int gsp_migration_transfer_iter(struct gsp_migration *mig);
static int gsp_migration_commit(struct gsp_migration *mig);
static void gsp_migration_cleanup(struct gsp_migration *mig);

/* Migration Thread */
static int gsp_migration_thread(void *data)
{
    struct gsp_migration *mig = data;
    int ret = 0;
    
    while (!kthread_should_stop() && !mig->cancelled) {
        mutex_lock(&mig->lock);
        
        switch (mig->state) {
        case GSP_MIG_S_PREPARE:
            ret = gsp_migration_prepare(mig);
            if (!ret)
                mig->state = GSP_MIG_S_TRANSFER;
            else
                mig->state = GSP_MIG_S_FAILED;
            break;
            
        case GSP_MIG_S_TRANSFER:
            ret = gsp_migration_transfer_iter(mig);
            if (ret > 0) {  /* Transfer complete */
                mig->state = GSP_MIG_S_COMMIT;
            } else if (ret < 0) {
                mig->state = GSP_MIG_S_FAILED;
            }
            break;
            
        case GSP_MIG_S_COMMIT:
            ret = gsp_migration_commit(mig);
            if (!ret)
                mig->state = GSP_MIG_S_CLEANUP;
            else
                mig->state = GSP_MIG_S_FAILED;
            break;
            
        case GSP_MIG_S_CLEANUP:
            gsp_migration_cleanup(mig);
            mig->state = GSP_MIG_S_COMPLETED;
            complete(&mig->comp);
            mutex_unlock(&mig->lock);
            return 0;
            
        case GSP_MIG_S_FAILED:
            mig->error = ret;
            complete(&mig->comp);
            mutex_unlock(&mig->lock);
            return ret;
            
        default:
            break;
        }
        
        mutex_unlock(&mig->lock);
        
        if (mig->state == GSP_MIG_S_TRANSFER)
            msleep(10);  /* Throttle transfer iterations */
    }
    
    return 0;
}

/* Migration Query */
static int gsp_migration_query(struct gsp_vf *vf, bool *migratable)
{
    struct gsp_migrate_payload payload = {
        .phase = GSP_MIG_QUERY,
        .subcmd = 0,
        .src_vf = vf->id,
        .dst_vf = 0,
    };
    struct gsp_subcmd_desc desc;
    size_t desc_size = sizeof(desc);
    int ret;
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MIGRATE,
                            &payload, sizeof(payload), &desc, &desc_size);
    if (ret)
        return ret;
    
    *migratable = (desc.subcmd_id != 0);
    return 0;
}

/* Phase 1: Prepare */
static int gsp_migration_prepare(struct gsp_migration *mig)
{
    struct gsp_vf *vf = mig->vf;
    struct gsp_migrate_payload payload = {0};
    struct gsp_subcmd_desc desc;
    size_t desc_size = sizeof(desc);
    int ret;
    
    /* Quiesce VF (subcmd 3 -> returns ID 0x10) */
    payload.phase = GSP_MIG_PREPARE;
    payload.subcmd = 3;
    payload.src_vf = vf->id;
    payload.dst_vf = mig->dst_vf;
    payload.flags = mig->flags;
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MIGRATE,
                            &payload, sizeof(payload), &desc, &desc_size);
    if (ret)
        return ret;
    
    mig->token = desc.subcmd_id;
    
    /* Snapshot BAR state */
    ret = gsp_migration_bar_save(vf, mig->bar_state.buf, &mig->bar_state.size);
    if (ret)
        return ret;
    
    /* Snapshot WPR regions */
    /* TODO: WPR snapshot via migration RPC */
    
    dev_info(vf->gsp->dev, "VF %u: Migration prepare done, token=0x%llx\n",
             vf->id, mig->token);
    
    return 0;
}

/* Phase 2: Transfer (Iterative) */
static int gsp_migration_transfer_iter(struct gsp_migration *mig)
{
    struct gsp_vf *vf = mig->vf;
    struct gsp_migrate_payload payload = {0};
    struct gsp_subcmd_desc desc;
    size_t desc_size = sizeof(desc);
    int ret;
    
    if (!mig->progress.page_list || mig->progress.page_count == 0)
        return 1;  /* Nothing to transfer */
    
    payload.phase = GSP_MIG_TRANSFER;
    payload.subcmd = 0;
    payload.src_vf = vf->id;
    payload.dst_vf = mig->dst_vf;
    payload.data_ptr = (uint64_t)mig->progress.page_list;
    payload.data_size = mig->progress.page_count * sizeof(uint64_t);
    payload.token = mig->token;
    payload.flags = mig->flags;
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MIGRATE,
                            &payload, sizeof(payload), &desc, &desc_size);
    if (ret)
        return ret;
    
    /* Update progress from response */
    /* Note: Actual progress parsing depends on firmware response format */
    mig->progress.pages_transferred += 100;  /* Placeholder */
    mig->progress.bytes_transferred += 100 * 4096;
    
    /* Check for dirty pages */
    if (mig->progress.dirty_pages > 0 && !(mig->flags & GSP_MIG_POSTCOPY))
        return 0;  /* Continue iterating */
    
    return 1;  /* Transfer complete */
}

/* Phase 3: Commit */
static int gsp_migration_commit(struct gsp_migration *mig)
{
    struct gsp_vf *vf = mig->vf;
    struct gsp_migrate_payload payload = {0};
    struct gsp_subcmd_desc desc;
    size_t desc_size = sizeof(desc);
    int ret;
    
    /* Final dirty page transfer */
    if (mig->progress.dirty_pages > 0) {
        payload.phase = GSP_MIG_TRANSFER;
        payload.subcmd = 11;  /* Dirty tracking */
        payload.src_vf = vf->id;
        payload.dst_vf = mig->dst_vf;
        payload.token = mig->token;
        payload.flags = GSP_MIG_POSTCOPY;
        
        ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MIGRATE,
                                &payload, sizeof(payload), &desc, &desc_size);
        if (ret)
            return ret;
    }
    
    /* Commit (subcmd 12) */
    payload.phase = GSP_MIG_COMMIT;
    payload.subcmd = 12;
    payload.src_vf = vf->id;
    payload.dst_vf = mig->dst_vf;
    payload.token = mig->token;
    payload.flags = mig->flags | GSP_MIG_VERIFY;
    
    ret = gsp_rpc_send_sync(vf->gsp->rpc, GSP_RPC_MIGRATE,
                            &payload, sizeof(payload), &desc, &desc_size);
    if (ret)
        return ret;
    
    /* Restore BAR on destination */
    ret = gsp_migration_bar_restore(vf, mig->bar_state.buf, mig->bar_state.size);
    if (ret)
        return ret;
    
    /* Restore WPR on destination */
    /* TODO: WPR restore */
    
    dev_info(vf->gsp->dev, "VF %u: Migration commit done\n", vf->id);
    
    return 0;
}

/* Phase 4: Cleanup */
static void gsp_migration_cleanup(struct gsp_migration *mig)
{
    struct gsp_vf *vf = mig->vf;
    
    /* Free page list */
    kfree(mig->progress.page_list);
    mig->progress.page_list = NULL;
    mig->progress.page_count = 0;
    
    /* Free BAR state buffer */
    kfree(mig->bar_state.buf);
    mig->bar_state.buf = NULL;
    mig->bar_state.size = 0;
    
    dev_info(vf->gsp->dev, "VF %u: Migration cleanup done\n", vf->id);
}

/* Public Migration API */
int gsp_migration_start(struct gsp_vf *vf, uint64_t dst_vf)
{
    struct gsp_migration *mig;
    bool migratable;
    int ret;
    
    if (vf->migration)
        return -EBUSY;
    
    /* Query migratability */
    ret = gsp_migration_query(vf, &migratable);
    if (ret || !migratable)
        return ret ? ret : -EINVAL;
    
    mig = kzalloc(sizeof(*mig), GFP_KERNEL);
    if (!mig)
        return -ENOMEM;
    
    mig->vf = vf;
    mig->dst_vf = dst_vf;
    mig->flags = GSP_MIG_LIVE | GSP_MIG_VERIFY;
    mig->state = GSP_MIG_S_PREPARE;
    mutex_init(&mig->lock);
    init_completion(&mig->comp);
    
    /* Allocate page list (placeholder) */
    mig->progress.page_count = vf->pt_size / 4096;
    mig->progress.page_list = kcalloc(mig->progress.page_count, sizeof(uint64_t), GFP_KERNEL);
    if (!mig->progress.page_list) {
        kfree(mig);
        return -ENOMEM;
    }
    
    /* Allocate BAR state buffer */
    mig->bar_state.buf = kmalloc(65536, GFP_KERNEL);  /* 64KB max */
    if (!mig->bar_state.buf) {
        kfree(mig->progress.page_list);
        kfree(mig);
        return -ENOMEM;
    }
    
    vf->migration = mig;
    
    /* Start migration thread */
    mig->thread = kthread_run(gsp_migration_thread, mig, "gsp-mig-%u", vf->id);
    if (IS_ERR(mig->thread)) {
        ret = PTR_ERR(mig->thread);
        mig->thread = NULL;
        goto err_free;
    }
    
    return 0;

err_free:
    kfree(mig->progress.page_list);
    kfree(mig->bar_state.buf);
    kfree(mig);
    vf->migration = NULL;
    return ret;
}

void gsp_migration_cancel(struct gsp_migration *mig)
{
    if (!mig)
        return;
    
    mutex_lock(&mig->lock);
    mig->cancelled = true;
    mig->state = GSP_MIG_S_FAILED;
    mutex_unlock(&mig->lock);
    
    if (mig->thread)
        kthread_stop(mig->thread);
    
    wait_for_completion(&mig->comp);
    
    gsp_migration_cleanup(mig);
    kfree(mig);
}

int gsp_migration_status(struct gsp_migration *mig)
{
    if (!mig)
        return -EINVAL;
    
    mutex_lock(&mig->lock);
    int state = mig->state;
    int error = mig->error;
    mutex_unlock(&mig->lock);
    
    if (state == GSP_MIG_S_FAILED)
        return error;
    
    return state;
}