cw-cyber/contracts/litium-refer/src/contract.rs

use cosmwasm_std::{
    entry_point, to_json_binary, Binary, CosmosMsg, Deps, DepsMut, Empty, Env, MessageInfo,
    Response, StdResult, Uint128, WasmMsg,
};
use cw2::{get_contract_version, set_contract_version};

use crate::error::ContractError;
use crate::msg::{
    CommunityPoolBalanceResponse, ConfigResponse, ExecuteMsg, InstantiateMsg, QueryMsg,
    ReferralInfoResponse, ReferrerOfResponse, TestingOverrides, TotalPendingRewardsResponse,
};
use crate::state::{
    ReferConfig, ReferralStats, COMMUNITY_POOL_BALANCE, CONFIG, REFERRAL_STATS, REFERRER_BY_MINER,
    TOTAL_ACCRUED_REWARDS, TOTAL_CLAIMED_REWARDS,
};

const CONTRACT_NAME: &str = "crates.io:litium-refer";
const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");

#[derive(serde::Serialize)]
#[serde(rename_all = "snake_case")]
enum LitiumCoreExecuteMsg {
    Mint { to: String, amount: Uint128 },
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn instantiate(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    msg: InstantiateMsg,
) -> Result<Response, ContractError> {
    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;

    let admin = msg
        .admin
        .map(|a| deps.api.addr_validate(&a))
        .transpose()?
        .unwrap_or(info.sender.clone());

    let community_pool_addr = msg
        .community_pool_addr
        .map(|a| deps.api.addr_validate(&a))
        .transpose()?;

    let config = ReferConfig {
        core_contract: deps.api.addr_validate(&msg.core_contract)?,
        mine_contract: deps.api.addr_validate(&msg.mine_contract)?,
        community_pool_addr,
        admin,
        paused: false,
    };
    CONFIG.save(deps.storage, &config)?;
    COMMUNITY_POOL_BALANCE.save(deps.storage, &Uint128::zero())?;
    TOTAL_ACCRUED_REWARDS.save(deps.storage, &Uint128::zero())?;
    TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;

    Ok(Response::new()
        .add_attribute("action", "instantiate")
        .add_attribute("admin", config.admin.to_string()))
}

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn execute(
    deps: DepsMut,
    _env: Env,
    info: MessageInfo,
    msg: ExecuteMsg,
) -> Result<Response, ContractError> {
    match msg {
        ExecuteMsg::BindReferrer { miner, referrer } => {
            execute_bind_referrer(deps, info, miner, referrer)
        }
        ExecuteMsg::AccrueReward { miner, amount } => {
            execute_accrue_reward(deps, info, miner, amount)
        }
        ExecuteMsg::ClaimRewards {} => execute_claim_rewards(deps, info),
        ExecuteMsg::ClaimCommunityPool { to } => execute_claim_community_pool(deps, info, to),
        ExecuteMsg::UpdateConfig {
            core_contract,
            mine_contract,
            community_pool_addr,
            admin,
        } => execute_update_config(
            deps,
            info,
            core_contract,
            mine_contract,
            community_pool_addr,
            admin,
        ),
        ExecuteMsg::ApplyTestingOverrides { overrides } => {
            execute_apply_testing_overrides(deps, info, overrides)
        }
        ExecuteMsg::ResetState {} => execute_reset_state(deps, info),
        ExecuteMsg::Pause {} => execute_pause(deps, info),
        ExecuteMsg::Unpause {} => execute_unpause(deps, info),
    }
}

fn execute_bind_referrer(
    deps: DepsMut,
    info: MessageInfo,
    miner: String,
    referrer: String,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }
    if info.sender != config.mine_contract {
        return Err(ContractError::NotMineContract {});
    }

    let miner_addr = deps.api.addr_validate(&miner)?;
    let referrer_addr = deps.api.addr_validate(&referrer)?;

    if miner_addr == referrer_addr {
        return Err(ContractError::SelfReferral {});
    }

    if let Some(existing) = REFERRER_BY_MINER.may_load(deps.storage, &miner_addr)? {
        if existing != referrer_addr {
            return Err(ContractError::ReferrerMismatch {});
        }
        // Already bound to same referrer, no-op
        return Ok(Response::new()
            .add_attribute("action", "bind_referrer")
            .add_attribute("miner", miner)
            .add_attribute("referrer", referrer)
            .add_attribute("status", "already_bound"));
    }

    REFERRER_BY_MINER.save(deps.storage, &miner_addr, &referrer_addr)?;

    let mut stats = REFERRAL_STATS
        .may_load(deps.storage, &referrer_addr)?
        .unwrap_or(ReferralStats {
            referral_rewards: Uint128::zero(),
            referrals_count: 0,
        });
    stats.referrals_count += 1;
    REFERRAL_STATS.save(deps.storage, &referrer_addr, &stats)?;

    Ok(Response::new()
        .add_attribute("action", "bind_referrer")
        .add_attribute("miner", miner)
        .add_attribute("referrer", referrer)
        .add_attribute("status", "bound"))
}

fn execute_accrue_reward(
    deps: DepsMut,
    info: MessageInfo,
    miner: String,
    amount: Uint128,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }
    if info.sender != config.mine_contract {
        return Err(ContractError::NotMineContract {});
    }

    if amount.is_zero() {
        return Ok(Response::new()
            .add_attribute("action", "accrue_reward")
            .add_attribute("miner", miner)
            .add_attribute("amount", "0")
            .add_attribute("destination", "none"));
    }

    let miner_addr = deps.api.addr_validate(&miner)?;
    TOTAL_ACCRUED_REWARDS.update(deps.storage, |v| -> StdResult<_> { Ok(v + amount) })?;

    // If no referrer is set, accrue to the community pool.
    if let Some(referrer_addr) = REFERRER_BY_MINER.may_load(deps.storage, &miner_addr)? {
        let mut stats = REFERRAL_STATS
            .may_load(deps.storage, &referrer_addr)?
            .unwrap_or(ReferralStats {
                referral_rewards: Uint128::zero(),
                referrals_count: 0,
            });
        stats.referral_rewards += amount;
        REFERRAL_STATS.save(deps.storage, &referrer_addr, &stats)?;

        Ok(Response::new()
            .add_attribute("action", "accrue_reward")
            .add_attribute("miner", miner)
            .add_attribute("amount", amount.to_string())
            .add_attribute("destination", referrer_addr.to_string()))
    } else {
        // No referrer → community pool
        COMMUNITY_POOL_BALANCE.update(deps.storage, |v| -> StdResult<_> { Ok(v + amount) })?;

        Ok(Response::new()
            .add_attribute("action", "accrue_reward")
            .add_attribute("miner", miner)
            .add_attribute("amount", amount.to_string())
            .add_attribute("destination", "community_pool"))
    }
}

fn execute_claim_rewards(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }

    let mut stats = REFERRAL_STATS
        .may_load(deps.storage, &info.sender)?
        .unwrap_or(ReferralStats {
            referral_rewards: Uint128::zero(),
            referrals_count: 0,
        });
    let amount = stats.referral_rewards;

    if amount.is_zero() {
        return Ok(Response::new()
            .add_attribute("action", "claim_rewards")
            .add_attribute("address", info.sender.to_string())
            .add_attribute("claimed", "0"));
    }

    stats.referral_rewards = Uint128::zero();
    REFERRAL_STATS.save(deps.storage, &info.sender, &stats)?;
    TOTAL_CLAIMED_REWARDS.update(deps.storage, |v| -> StdResult<_> { Ok(v + amount) })?;

    // Mint via litium-core
    let mint_msg = CosmosMsg::Wasm(WasmMsg::Execute {
        contract_addr: config.core_contract.to_string(),
        msg: cosmwasm_std::to_json_binary(&LitiumCoreExecuteMsg::Mint {
            to: info.sender.to_string(),
            amount,
        })?,
        funds: vec![],
    });

    Ok(Response::new()
        .add_message(mint_msg)
        .add_attribute("action", "claim_rewards")
        .add_attribute("address", info.sender.to_string())
        .add_attribute("claimed", amount.to_string()))
}

fn execute_claim_community_pool(
    deps: DepsMut,
    info: MessageInfo,
    to: String,
) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if config.paused {
        return Err(ContractError::Paused {});
    }
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }

    let balance = COMMUNITY_POOL_BALANCE.load(deps.storage)?;
    if balance.is_zero() {
        return Ok(Response::new()
            .add_attribute("action", "claim_community_pool")
            .add_attribute("claimed", "0"));
    }

    COMMUNITY_POOL_BALANCE.save(deps.storage, &Uint128::zero())?;
    TOTAL_CLAIMED_REWARDS.update(deps.storage, |v| -> StdResult<_> { Ok(v + balance) })?;

    let to_addr = deps.api.addr_validate(&to)?;
    let mint_msg = CosmosMsg::Wasm(WasmMsg::Execute {
        contract_addr: config.core_contract.to_string(),
        msg: cosmwasm_std::to_json_binary(&LitiumCoreExecuteMsg::Mint {
            to: to_addr.to_string(),
            amount: balance,
        })?,
        funds: vec![],
    });

    Ok(Response::new()
        .add_message(mint_msg)
        .add_attribute("action", "claim_community_pool")
        .add_attribute("to", to_addr.to_string())
        .add_attribute("claimed", balance.to_string()))
}

fn execute_update_config(
    deps: DepsMut,
    info: MessageInfo,
    core_contract: Option<String>,
    mine_contract: Option<String>,
    community_pool_addr: Option<String>,
    admin: Option<String>,
) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    if let Some(c) = core_contract {
        config.core_contract = deps.api.addr_validate(&c)?;
    }
    if let Some(m) = mine_contract {
        config.mine_contract = deps.api.addr_validate(&m)?;
    }
    if let Some(a) = community_pool_addr {
        config.community_pool_addr = Some(deps.api.addr_validate(&a)?);
    }
    if let Some(a) = admin {
        config.admin = deps.api.addr_validate(&a)?;
    }
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "update_config"))
}

fn execute_apply_testing_overrides(
    deps: DepsMut,
    info: MessageInfo,
    overrides: TestingOverrides,
) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }

    if let Some(paused) = overrides.paused {
        config.paused = paused;
    }
    if let Some(addr) = overrides.community_pool_addr {
        config.community_pool_addr = Some(deps.api.addr_validate(&addr)?);
    }
    CONFIG.save(deps.storage, &config)?;

    if let Some(balance) = overrides.community_pool_balance {
        COMMUNITY_POOL_BALANCE.save(deps.storage, &balance)?;
    }
    if let Some(v) = overrides.total_accrued_rewards {
        TOTAL_ACCRUED_REWARDS.save(deps.storage, &v)?;
    }
    if let Some(v) = overrides.total_claimed_rewards {
        TOTAL_CLAIMED_REWARDS.save(deps.storage, &v)?;
    }

    Ok(Response::new().add_attribute("action", "apply_testing_overrides"))
}

/// Full state reset for daily testing. Zeroes referral stats and clears all bindings.
fn execute_reset_state(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }

    // Zero community pool and pending rewards
    COMMUNITY_POOL_BALANCE.save(deps.storage, &Uint128::zero())?;
    TOTAL_ACCRUED_REWARDS.save(deps.storage, &Uint128::zero())?;
    TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;

    // Clear all referrer bindings
    let miners: Vec<_> = REFERRER_BY_MINER
        .keys(deps.storage, None, None, cosmwasm_std::Order::Ascending)
        .collect::<Result<Vec<_>, _>>()?;
    for addr in &miners {
        REFERRER_BY_MINER.remove(deps.storage, addr);
    }

    // Clear all referral stats
    let referrers: Vec<_> = REFERRAL_STATS
        .keys(deps.storage, None, None, cosmwasm_std::Order::Ascending)
        .collect::<Result<Vec<_>, _>>()?;
    for addr in &referrers {
        REFERRAL_STATS.remove(deps.storage, addr);
    }

    Ok(Response::new()
        .add_attribute("action", "reset_state")
        .add_attribute("bindings_cleared", miners.len().to_string())
        .add_attribute("referrers_cleared", referrers.len().to_string()))
}

fn execute_pause(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    config.paused = true;
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "pause"))
}

fn execute_unpause(deps: DepsMut, info: MessageInfo) -> Result<Response, ContractError> {
    let mut config = CONFIG.load(deps.storage)?;
    if info.sender != config.admin {
        return Err(ContractError::Unauthorized {});
    }
    config.paused = false;
    CONFIG.save(deps.storage, &config)?;
    Ok(Response::new().add_attribute("action", "unpause"))
}

// ============================================================
// Migrate
// ============================================================

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn migrate(deps: DepsMut, _env: Env, _msg: Empty) -> Result<Response, ContractError> {
    let stored = get_contract_version(deps.storage)?;
    let stored_ver = stored
        .version
        .parse::<semver::Version>()
        .map_err(|_| ContractError::MigrationError {})?;
    let current_ver = CONTRACT_VERSION
        .parse::<semver::Version>()
        .map_err(|_| ContractError::MigrationError {})?;
    if stored_ver >= current_ver {
        return Err(ContractError::MigrationError {});
    }
    // v0.1.1 → v0.1.2: migrate TOTAL_PENDING_REWARDS → monotonic counters
    if TOTAL_ACCRUED_REWARDS.may_load(deps.storage)?.is_none() {
        let old_pending: Uint128 = deps
            .storage
            .get(b"total_pending_rewards")
            .and_then(|v| cosmwasm_std::from_json(&v).ok())
            .unwrap_or_default();
        TOTAL_ACCRUED_REWARDS.save(deps.storage, &old_pending)?;
        TOTAL_CLAIMED_REWARDS.save(deps.storage, &Uint128::zero())?;
        deps.storage.remove(b"total_pending_rewards");
    }

    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
    Ok(Response::default())
}

// ============================================================
// Queries
// ============================================================

#[cfg_attr(not(feature = "library"), entry_point)]
pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
    match msg {
        QueryMsg::Config {} => to_json_binary(&query_config(deps)?),
        QueryMsg::ReferrerOf { miner } => to_json_binary(&query_referrer_of(deps, miner)?),
        QueryMsg::ReferralInfo { address } => to_json_binary(&query_referral_info(deps, address)?),
        QueryMsg::CommunityPoolBalance {} => to_json_binary(&query_community_pool_balance(deps)?),
        QueryMsg::TotalPendingRewards {} => to_json_binary(&query_total_pending_rewards(deps)?),
    }
}

fn query_config(deps: Deps) -> StdResult<ConfigResponse> {
    let config = CONFIG.load(deps.storage)?;
    Ok(ConfigResponse {
        core_contract: config.core_contract.to_string(),
        mine_contract: config.mine_contract.to_string(),
        community_pool_addr: config.community_pool_addr.map(|a| a.to_string()),
        admin: config.admin.to_string(),
        paused: config.paused,
    })
}

fn query_referrer_of(deps: Deps, miner: String) -> StdResult<ReferrerOfResponse> {
    let miner_addr = deps.api.addr_validate(&miner)?;
    let referrer = REFERRER_BY_MINER
        .may_load(deps.storage, &miner_addr)?
        .map(|a| a.to_string());
    Ok(ReferrerOfResponse { miner, referrer })
}

fn query_referral_info(deps: Deps, address: String) -> StdResult<ReferralInfoResponse> {
    let addr = deps.api.addr_validate(&address)?;
    let stats = REFERRAL_STATS
        .may_load(deps.storage, &addr)?
        .unwrap_or(ReferralStats {
            referral_rewards: Uint128::zero(),
            referrals_count: 0,
        });
    Ok(ReferralInfoResponse {
        address,
        referral_rewards: stats.referral_rewards,
        referrals_count: stats.referrals_count,
    })
}

fn query_community_pool_balance(deps: Deps) -> StdResult<CommunityPoolBalanceResponse> {
    let balance = COMMUNITY_POOL_BALANCE
        .may_load(deps.storage)?
        .unwrap_or_else(Uint128::zero);
    Ok(CommunityPoolBalanceResponse { balance })
}

fn query_total_pending_rewards(deps: Deps) -> StdResult<TotalPendingRewardsResponse> {
    let accrued = TOTAL_ACCRUED_REWARDS
        .may_load(deps.storage)?
        .unwrap_or_else(Uint128::zero);
    let claimed = TOTAL_CLAIMED_REWARDS
        .may_load(deps.storage)?
        .unwrap_or_else(Uint128::zero);
    Ok(TotalPendingRewardsResponse {
        total_pending_rewards: accrued.saturating_sub(claimed),
    })
}