#include <nets/orphan_pool.h>
#include <constants.h>
#include <dlibc/vector.h>
#include <pthread.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

typedef struct {
    block_t* block;
    uint64_t height;
    uint64_t observedAtTipHeight; // local tip height when first seen; stamped once (reorg penalty)
    uint64_t sequence;            // insertion order, used to evict the oldest entry when full
    uint8_t hash[32];
} orphan_entry_t;

static vector_t* g_orphans = NULL;
static uint64_t g_nextSequence = 0;

// The pool is touched by the maintenance thread, by every per-peer TCP thread and by the REPL
// thread. It used to have no synchronisation at all, so a concurrent Insert could realloc the
// array out from under a scan that was holding a raw element pointer.
//
// Lock ordering: this mutex is never held while calling into chain.c (which takes g_chainLock).
// Candidate branches are collected under the lock, the lock is dropped, and only then is
// Chain_ReplaceBranch/Chain_AddBlock called.
static pthread_mutex_t g_orphanLock = PTHREAD_MUTEX_INITIALIZER;

// Element destructor for `g_orphans`: an entry owns its block, so dropping or evicting an entry
// frees the block along with it. OrphanPool_TakeByHashLocked hands a block out through
// vector_take_at(), which deliberately skips this. Block_Destroy accepts NULL.
static void OrphanPool_EntryDestructor(void* element) {
    orphan_entry_t* e = (orphan_entry_t*)element;
    Block_Destroy(e->block);
}

static void OrphanPool_InitLocked(void) {
    if (!g_orphans) {
        g_orphans = vector_create(sizeof(orphan_entry_t));
        vector_set_destructor(g_orphans, OrphanPool_EntryDestructor);
    }
}

void OrphanPool_Init(void) {
    pthread_mutex_lock(&g_orphanLock);
    OrphanPool_InitLocked();
    pthread_mutex_unlock(&g_orphanLock);
}

void OrphanPool_Destroy(void) {
    pthread_mutex_lock(&g_orphanLock);
    vector_destroy(&g_orphans); // The element destructor frees every pooled block
    pthread_mutex_unlock(&g_orphanLock);
}

static ssize_t OrphanPool_FindByHashLocked(const uint8_t blockHash[32]) {
    if (!g_orphans || !blockHash) {
        return -1;
    }

    size_t n = vector_size(g_orphans);
    for (size_t i = 0; i < n; ++i) {
        orphan_entry_t* e = (orphan_entry_t*)vector_get(g_orphans, i);
        if (e && memcmp(e->hash, blockHash, 32) == 0) {
            return (ssize_t)i;
        }
    }

    return -1;
}

// Drop the entry with the lowest sequence number, so a flood of unusable orphans cannot grow
// without bound. Returns true if something was evicted.
static bool OrphanPool_EvictOldestLocked(void) {
    if (!g_orphans) {
        return false;
    }

    size_t n = vector_size(g_orphans);
    if (n == 0) {
        return false;
    }

    size_t oldestIndex = 0;
    uint64_t oldestSequence = UINT64_MAX;
    for (size_t i = 0; i < n; ++i) {
        orphan_entry_t* e = (orphan_entry_t*)vector_get(g_orphans, i);
        if (e && e->sequence < oldestSequence) {
            oldestSequence = e->sequence;
            oldestIndex = i;
        }
    }

    vector_pop_at(g_orphans, oldestIndex); // The element destructor frees the victim's block
    return true;
}

void OrphanPool_Insert(block_t* block, uint64_t height, uint64_t observedAtTipHeight) {
    if (!block) {
        return;
    }

    uint8_t blockHash[32];
    Block_CalculateHash(block, blockHash);

    pthread_mutex_lock(&g_orphanLock);
    OrphanPool_InitLocked();
    if (!g_orphans) {
        pthread_mutex_unlock(&g_orphanLock);
        Block_Destroy(block);
        return;
    }

    // Reject duplicates. The same block reaches us from every peer that relays it, and without
    // this each copy became its own permanently-resident entry.
    if (OrphanPool_FindByHashLocked(blockHash) >= 0) {
        pthread_mutex_unlock(&g_orphanLock);
        Block_Destroy(block);
        return;
    }

    while (vector_size(g_orphans) >= MAX_ORPHAN_BLOCKS) {
        if (!OrphanPool_EvictOldestLocked()) {
            break;
        }
    }

    orphan_entry_t e;
    memset(&e, 0, sizeof(e));
    e.block = block;
    e.height = height;
    e.observedAtTipHeight = observedAtTipHeight;
    e.sequence = g_nextSequence++;
    memcpy(e.hash, blockHash, 32);

    if (vector_push_back(g_orphans, &e) != 0) {
        pthread_mutex_unlock(&g_orphanLock);
        Block_Destroy(block); // never pooled, so the pool does not own it
        return;
    }

    pthread_mutex_unlock(&g_orphanLock);
}

bool OrphanPool_Contains(const uint8_t blockHash[32]) {
    pthread_mutex_lock(&g_orphanLock);
    bool found = OrphanPool_FindByHashLocked(blockHash) >= 0;
    pthread_mutex_unlock(&g_orphanLock);
    return found;
}

size_t OrphanPool_Size(void) {
    pthread_mutex_lock(&g_orphanLock);
    size_t n = vector_size(g_orphans);
    pthread_mutex_unlock(&g_orphanLock);
    return n;
}

// Remove the entry with this hash without freeing the block, and hand the block back. Used once a
// block has been given to the chain, which then owns its transaction array. vector_take_at() skips
// the element destructor, which is what moves ownership of the block to the caller.
static block_t* OrphanPool_TakeByHashLocked(const uint8_t blockHash[32]) {
    ssize_t index = OrphanPool_FindByHashLocked(blockHash);
    if (index < 0) {
        return NULL;
    }

    orphan_entry_t entry;
    if (vector_take_at(g_orphans, (size_t)index, &entry) != 0) {
        return NULL;
    }
    return entry.block;
}

static void OrphanPool_DropByHashLocked(const uint8_t blockHash[32]) {
    ssize_t index = OrphanPool_FindByHashLocked(blockHash);
    if (index < 0) {
        return;
    }

    vector_pop_at(g_orphans, (size_t)index); // The element destructor frees the block
}

/**
 * Copy out the orphan that extends `prevHash` at `height`, if any.
 * Returns false when there is no such orphan. Caller must hold the pool lock.
**/
static bool OrphanPool_FindChildLocked(uint64_t height,
                                       const uint8_t prevHash[32],
                                       block_t** outBlock,
                                       uint64_t* outObservedAtTipHeight,
                                       uint8_t outHash[32]) {
    if (!g_orphans) {
        return false;
    }

    size_t n = vector_size(g_orphans);
    for (size_t i = 0; i < n; ++i) {
        orphan_entry_t* e = (orphan_entry_t*)vector_get(g_orphans, i);
        if (!e || !e->block) {
            continue;
        }
        if (e->height != height) {
            continue;
        }
        if (memcmp(e->block->header.prevHash, prevHash, 32) != 0) {
            continue;
        }

        *outBlock = e->block;
        *outObservedAtTipHeight = e->observedAtTipHeight;
        memcpy(outHash, e->hash, 32);
        return true;
    }

    return false;
}

/**
 * Follow prevHash links from `forkHeight` to build the longest branch the pool can offer.
 *
 * The old implementation took the first orphan found at each successive height with no linkage
 * check at all, which could splice blocks from two different forks into one incoherent branch.
 * Caller must hold the pool lock. The returned array borrows the pooled block pointers; the pool
 * still owns them (Chain_ReplaceBranch applies copies).
**/
static size_t OrphanPool_CollectBranchLocked(uint64_t forkHeight,
                                             const uint8_t forkParentHash[32],
                                             block_t*** outBlocks,
                                             uint8_t** outHashes,
                                             uint64_t* outObservedAtTipHeight) {
    *outBlocks = NULL;
    *outHashes = NULL;
    *outObservedAtTipHeight = 0;

    vector_t* collected = vector_create(sizeof(block_t*));
    vector_t* hashes = vector_create(32); // one whole 32-byte block hash per element, index-aligned with collected
    if (!collected || !hashes) {
        vector_destroy(&collected);
        vector_destroy(&hashes);
        return 0;
    }

    uint8_t expectedPrevHash[32];
    memcpy(expectedPrevHash, forkParentHash, 32);

    uint64_t earliestObserved = UINT64_MAX;
    uint64_t cursor = forkHeight;
    size_t count = 0;

    while (1) {
        block_t* child = NULL;
        uint64_t observed = 0;
        uint8_t childHash[32];
        if (!OrphanPool_FindChildLocked(cursor, expectedPrevHash, &child, &observed, childHash)) {
            break;
        }

        // A block only counts once both pushes have landed, so the two vectors never drift out of
        // step: the copy-out below reads index i from both.
        if (vector_push_back(collected, &child) != 0) {
            break;
        }
        if (vector_push_back(hashes, childHash) != 0) {
            break;
        }

        if (observed < earliestObserved) {
            earliestObserved = observed;
        }

        memcpy(expectedPrevHash, childHash, 32);
        cursor++;
        count++;
    }

    if (count == 0) {
        vector_destroy(&collected);
        vector_destroy(&hashes);
        return 0;
    }

    block_t** blocks = (block_t**)calloc(count, sizeof(block_t*));
    uint8_t* hashOut = (uint8_t*)calloc(count, 32);
    if (!blocks || !hashOut) {
        free(blocks);
        free(hashOut);
        vector_destroy(&collected);
        vector_destroy(&hashes);
        return 0;
    }

    for (size_t i = 0; i < count; ++i) {
        block_t** blockSlot = (block_t**)vector_get(collected, i);
        const uint8_t* hashSlot = (const uint8_t*)vector_get_const(hashes, i);
        if (!blockSlot || !hashSlot) {
            count = i; // never hand back a branch with a hole in it
            break;
        }
        blocks[i] = *blockSlot;
        memcpy(&hashOut[i * 32], hashSlot, 32);
    }

    vector_destroy(&collected);
    vector_destroy(&hashes);

    *outBlocks = blocks;
    *outHashes = hashOut;
    *outObservedAtTipHeight = earliestObserved == UINT64_MAX ? 0ULL : earliestObserved;
    return count;
}

// Discard orphans that can no longer ever be applied: anything at or below the current tip whose
// hash does not match the block we actually have there. Without this the pool only ever grew, and
// permanently-invalid entries were retried on every 1 Hz maintenance tick.
static void OrphanPool_PruneStale(blockchain_t* chain) {
    if (!chain) {
        return;
    }

    const size_t chainSize = Chain_Size(chain);

    // Collect the hashes to drop first, so we never call into chain.c while holding the pool lock.
    vector_t* doomed = vector_create(32); // one whole 32-byte block hash per element
    if (!doomed) {
        return;
    }

    pthread_mutex_lock(&g_orphanLock);
    size_t n = vector_size(g_orphans);
    vector_t* candidates = vector_create(32); // index-aligned with candidateHeights
    vector_t* candidateHeights = vector_create(sizeof(uint64_t));
    if (candidates && candidateHeights) {
        for (size_t i = 0; i < n; ++i) {
            orphan_entry_t* e = (orphan_entry_t*)vector_get(g_orphans, i);
            if (!e || !e->block) {
                continue;
            }
            if (e->height >= (uint64_t)chainSize) {
                continue; // still ahead of us; may attach later
            }
            if (vector_push_back(candidates, e->hash) != 0) {
                break;
            }
            if (vector_push_back(candidateHeights, &e->height) != 0) {
                vector_pop_back(candidates); // keep the two index-aligned
                break;
            }
        }
    }
    pthread_mutex_unlock(&g_orphanLock);

    size_t candidateCount = vector_size(candidateHeights);
    for (size_t i = 0; i < candidateCount; ++i) {
        const uint64_t* heightSlot = (const uint64_t*)vector_get_const(candidateHeights, i);
        const uint8_t* orphanHash = (const uint8_t*)vector_get_const(candidates, i);
        if (!heightSlot || !orphanHash) {
            continue;
        }
        uint64_t height = *heightSlot;

        block_t* local = NULL;
        if (!Chain_GetBlockCopy(chain, (size_t)height, &local) || !local) {
            continue;
        }

        uint8_t localHash[32];
        Block_CalculateHash(local, localHash);
        Block_Destroy(local);

        // Same block we already have: pure duplicate, drop it. A different block at a height we
        // have already passed is kept, because it may yet be the base of a heavier branch.
        if (memcmp(localHash, orphanHash, 32) == 0) {
            (void)vector_push_back(doomed, orphanHash);
        }
    }

    size_t doomedCount = vector_size(doomed);
    if (doomedCount > 0) {
        pthread_mutex_lock(&g_orphanLock);
        for (size_t i = 0; i < doomedCount; ++i) {
            const uint8_t* h = (const uint8_t*)vector_get_const(doomed, i);
            if (h) {
                OrphanPool_DropByHashLocked(h);
            }
        }
        pthread_mutex_unlock(&g_orphanLock);
    }

    vector_destroy(&candidates);
    vector_destroy(&candidateHeights);
    vector_destroy(&doomed);
}

/**
 * Try to extend the current tip directly with pooled orphans.
 * Returns the number of blocks attached.
**/
static size_t OrphanPool_ExtendTip(blockchain_t* chain) {
    size_t attached = 0;

    while (1) {
        const size_t chainSize = Chain_Size(chain);

        uint8_t tipHash[32];
        memset(tipHash, 0, sizeof(tipHash));
        if (chainSize > 0) {
            block_t* tip = NULL;
            if (!Chain_GetBlockCopy(chain, chainSize - 1, &tip) || !tip) {
                break;
            }
            Block_CalculateHash(tip, tipHash);
            Block_Destroy(tip);
        }

        // Take a copy of the candidate under the lock, then release it before touching the chain.
        pthread_mutex_lock(&g_orphanLock);
        block_t* pooled = NULL;
        uint64_t observed = 0;
        uint8_t candidateHash[32];
        bool found = OrphanPool_FindChildLocked((uint64_t)chainSize, tipHash, &pooled, &observed, candidateHash);
        block_t* candidate = found ? Block_Copy(pooled) : NULL;
        pthread_mutex_unlock(&g_orphanLock);

        if (!found) {
            break;
        }
        if (!candidate) {
            break;
        }

        if (!Chain_AddBlock(chain, candidate)) {
            // Permanent rejection for this block at this height (bad coinbase, wrong difficulty,
            // ...). Drop it rather than retrying it on every maintenance tick forever.
            Block_Destroy(candidate);
            pthread_mutex_lock(&g_orphanLock);
            OrphanPool_DropByHashLocked(candidateHash);
            pthread_mutex_unlock(&g_orphanLock);
            continue;
        }

        // Chain_AddBlock took ownership of the transaction array and cleared our pointer to it.
        Block_Destroy(candidate);

        pthread_mutex_lock(&g_orphanLock);
        block_t* taken = OrphanPool_TakeByHashLocked(candidateHash);
        pthread_mutex_unlock(&g_orphanLock);
        if (taken) {
            Block_Destroy(taken); // the pool's own copy is independent of the one we applied
        }

        attached++;
    }

    return attached;
}

/**
 * Look for a competing branch that forks below our tip and is worth adopting.
 * The work comparison, the reorg penalty and the atomicity all live in Chain_ReplaceBranch.
**/
static size_t OrphanPool_TryAdoptBranch(blockchain_t* chain, bool bypassPenalty) {
    const size_t chainSize = Chain_Size(chain);
    if (chainSize == 0) {
        return 0;
    }

    // Walk fork points from just below the tip downwards; the shallowest fork wins, which is also
    // the one with the smallest reorg penalty.
    for (size_t forkHeight = chainSize; forkHeight >= 1; --forkHeight) {
        block_t* parent = NULL;
        if (!Chain_GetBlockCopy(chain, forkHeight - 1, &parent) || !parent) {
            continue;
        }
        uint8_t parentHash[32];
        Block_CalculateHash(parent, parentHash);
        Block_Destroy(parent);

        pthread_mutex_lock(&g_orphanLock);
        block_t** branch = NULL;
        uint8_t* branchHashes = NULL;
        uint64_t observedAtTipHeight = 0;
        size_t branchCount = OrphanPool_CollectBranchLocked((uint64_t)forkHeight, parentHash,
                                                            &branch, &branchHashes, &observedAtTipHeight);
        // Copy the branch so the pool lock can be released before we call into the chain.
        block_t** branchCopies = NULL;
        if (branchCount > 0) {
            branchCopies = (block_t**)calloc(branchCount, sizeof(block_t*));
            if (branchCopies) {
                for (size_t i = 0; i < branchCount; ++i) {
                    branchCopies[i] = Block_Copy(branch[i]);
                }
            }
        }
        pthread_mutex_unlock(&g_orphanLock);

        free(branch);

        if (branchCount == 0 || !branchCopies) {
            free(branchHashes);
            if (branchCopies) {
                free(branchCopies);
            }
            if (forkHeight == 1) break;
            continue;
        }

        bool copiedAll = true;
        for (size_t i = 0; i < branchCount; ++i) {
            if (!branchCopies[i]) {
                copiedAll = false;
            }
        }

        bool adopted = false;
        if (copiedAll) {
            adopted = Chain_ReplaceBranch(chain, forkHeight, branchCopies, branchCount, observedAtTipHeight,
                                          bypassPenalty);
        }

        for (size_t i = 0; i < branchCount; ++i) {
            if (branchCopies[i]) {
                Block_Destroy(branchCopies[i]); // Chain_ReplaceBranch applied its own copies
            }
        }
        free(branchCopies);

        if (adopted) {
            printf("Adopted competing branch of %zu block(s) at fork height %zu\n", branchCount, forkHeight);
            pthread_mutex_lock(&g_orphanLock);
            for (size_t i = 0; i < branchCount; ++i) {
                OrphanPool_DropByHashLocked(&branchHashes[i * 32]);
            }
            pthread_mutex_unlock(&g_orphanLock);
            free(branchHashes);
            return branchCount;
        }

        free(branchHashes);

        if (forkHeight == 1) {
            break;
        }
    }

    return 0;
}

size_t OrphanPool_AttemptAttach(blockchain_t* chain) {
    return OrphanPool_AttemptAttachForced(chain, false);
}

size_t OrphanPool_AttemptAttachForced(blockchain_t* chain, bool bypassPenalty) {
    if (!chain) {
        return 0;
    }

    pthread_mutex_lock(&g_orphanLock);
    bool empty = vector_is_empty(g_orphans); // also true while the pool is not initialised
    pthread_mutex_unlock(&g_orphanLock);
    if (empty) {
        return 0;
    }

    size_t attached = 0;

    // Extending the tip is always preferable to a reorg, so try that to exhaustion first, and only
    // then consider replacing part of our chain with a competing branch.
    while (1) {
        size_t extended = OrphanPool_ExtendTip(chain);
        attached += extended;

        size_t adopted = OrphanPool_TryAdoptBranch(chain, bypassPenalty);
        attached += adopted;

        if (extended == 0 && adopted == 0) {
            break;
        }
    }

    OrphanPool_PruneStale(chain);

    return attached;
}