#![cfg_attr(not(feature = "export-abi"), no_main)]
extern crate alloc;

use stylus_sdk::prelude::*;
use stylus_sdk::{msg, call};
use stylus_sdk::storage::{StorageMap, StorageAddress, StorageU256};
use alloy_primitives::{Address, U256};

#[storage]
#[entrypoint]
pub struct VendingMachine {
    owner: StorageAddress,
    cupcake_balances: StorageMap<Address, U256>,
    cupcake_price: StorageU256,
}

#[public]
impl VendingMachine {
    pub fn init(&mut self) {
        self.owner.set(msg::sender());
        self.cupcake_price.set(U256::from(1_000_000_000_000_000_u64)); // 0.001 ETH
    }

    pub fn get_balance(&self, addr: Address) -> U256 {
        self.cupcake_balances.get(addr)
    }

    // Restock cupcakes (owner only)
    pub fn restock(&mut self, amount: U256) {
        assert_eq!(msg::sender(), self.owner.get(), "Only owner");
        let owner = self.owner.get();
        let current = self.cupcake_balances.get(owner);
        self.cupcake_balances.insert(owner, current + amount);
    }

    // Purchase cupcakes with ETH
    #[payable]
    pub fn purchase(&mut self, amount: U256) {
        let price = self.cupcake_price.get();
        let total_cost = price * amount;

        assert!(msg::value() >= total_cost, "Not enough ETH");

        let owner = self.owner.get();
        let owner_balance = self.cupcake_balances.get(owner);
        assert!(owner_balance >= amount, "Not enough cupcakes");

        self.cupcake_balances.insert(owner, owner_balance - amount);

        let buyer = msg::sender();
        let buyer_balance = self.cupcake_balances.get(buyer);
        self.cupcake_balances.insert(buyer, buyer_balance + amount);
    }

    // Withdraw ETH (owner only)
    pub fn withdraw(&mut self) {
        assert_eq!(msg::sender(), self.owner.get(), "Only owner");
        let balance = contract::balance();
        call::transfer_eth(self.owner.get(), balance).unwrap();
    }
}