//! Token Staking Contract
//!
//! A staking contract with time-locked deposits and proportional rewards.
//!
//! # Features
//! - Stake tokens for rewards
//! - Lock periods with bonus multipliers
//! - Claim rewards anytime
//! - Emergency unstake (with penalty)
//!
//! # Methods
//! - `init(staking_token, reward_token, reward_rate)` - Initialize staking pool
//! - `stake(amount, lock_days)` - Stake tokens with optional lock period
//! - `unstake(stake_id)` - Unstake after lock period ends
//! - `emergency_unstake(stake_id)` - Force unstake with 10% penalty
//! - `claim(stake_id)` - Claim accumulated rewards
//! - `get_stake(owner, stake_id) -> StakeInfo` - Get stake details
//! - `get_stakes(owner) -> Vec<StakeInfo>` - Get all stakes for owner
//! - `get_pool_info() -> PoolInfo` - Get pool statistics

#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use borsh::{BorshDeserialize, BorshSerialize};
use synor_sdk::prelude::*;
use synor_sdk::{require, require_auth};

// ==================== Constants ====================

/// Emergency unstake penalty: 10%
const EMERGENCY_PENALTY_BPS: u64 = 1000;
const BPS_DENOMINATOR: u64 = 10000;

/// Seconds per day
const SECONDS_PER_DAY: u64 = 86400;

/// Lock period bonus multipliers (in basis points, 10000 = 1x)
const BONUS_NO_LOCK: u64 = 10000;      // 1x for no lock
const BONUS_30_DAYS: u64 = 12000;      // 1.2x for 30 days
const BONUS_90_DAYS: u64 = 15000;      // 1.5x for 90 days
const BONUS_180_DAYS: u64 = 20000;     // 2x for 180 days
const BONUS_365_DAYS: u64 = 30000;     // 3x for 365 days

// ==================== Storage Keys ====================

mod keys {
    pub const STAKING_TOKEN: &[u8] = b"stk:staking_token";
    pub const REWARD_TOKEN: &[u8] = b"stk:reward_token";
    pub const REWARD_RATE: &[u8] = b"stk:reward_rate";      // Rewards per second per staked token
    pub const TOTAL_STAKED: &[u8] = b"stk:total_staked";
    pub const TOTAL_REWARDS_PAID: &[u8] = b"stk:rewards_paid";
    pub const OWNER: &[u8] = b"stk:owner";
    pub const STAKES: &[u8] = b"stk:stakes";
    pub const STAKE_COUNT: &[u8] = b"stk:stake_count";
    pub const PAUSED: &[u8] = b"stk:paused";
}

// ==================== Data Structures ====================

#[derive(BorshSerialize, BorshDeserialize, Clone)]
pub struct StakeInfo {
    /// Unique stake ID
    pub id: u64,
    /// Amount staked
    pub amount: u64,
    /// Effective amount (with bonus multiplier)
    pub effective_amount: u64,
    /// Timestamp when staked
    pub staked_at: u64,
    /// Lock end timestamp (0 = no lock)
    pub lock_until: u64,
    /// Last time rewards were calculated
    pub last_reward_time: u64,
    /// Accumulated unclaimed rewards
    pub pending_rewards: u64,
    /// Whether this stake is active
    pub active: bool,
}

#[derive(BorshSerialize)]
pub struct PoolInfo {
    pub staking_token: Address,
    pub reward_token: Address,
    pub reward_rate: u64,
    pub total_staked: u64,
    pub total_rewards_paid: u64,
    pub paused: bool,
}

// ==================== Storage Helpers ====================

fn stake_key(owner: &Address, stake_id: u64) -> Vec<u8> {
    let mut key = Vec::with_capacity(42);
    key.extend_from_slice(owner.as_bytes());
    key.extend_from_slice(&stake_id.to_le_bytes());
    key
}

fn get_stake(owner: &Address, stake_id: u64) -> Option<StakeInfo> {
    let key = stake_key(owner, stake_id);
    storage::get_with_suffix::<StakeInfo>(keys::STAKES, &key)
}

fn set_stake(owner: &Address, stake_id: u64, stake: &StakeInfo) {
    let key = stake_key(owner, stake_id);
    storage::set_with_suffix(keys::STAKES, &key, stake);
}

fn get_stake_count(owner: &Address) -> u64 {
    storage::get_with_suffix::<u64>(keys::STAKE_COUNT, owner.as_bytes()).unwrap_or(0)
}

fn increment_stake_count(owner: &Address) -> u64 {
    let count = get_stake_count(owner) + 1;
    storage::set_with_suffix(keys::STAKE_COUNT, owner.as_bytes(), &count);
    count
}

fn get_bonus_multiplier(lock_days: u64) -> u64 {
    if lock_days >= 365 {
        BONUS_365_DAYS
    } else if lock_days >= 180 {
        BONUS_180_DAYS
    } else if lock_days >= 90 {
        BONUS_90_DAYS
    } else if lock_days >= 30 {
        BONUS_30_DAYS
    } else {
        BONUS_NO_LOCK
    }
}

fn is_paused() -> bool {
    storage::get::<bool>(keys::PAUSED).unwrap_or(false)
}

// ==================== Entry Points ====================

synor_sdk::entry_point!(init, call);

/// Initialize the staking pool.
fn init(params: &[u8]) -> Result<()> {
    #[derive(BorshDeserialize)]
    struct InitParams {
        staking_token: Address,
        reward_token: Address,
        reward_rate: u64,  // Rewards per second per 1e8 staked tokens
    }

    let params = InitParams::try_from_slice(params)
        .map_err(|_| Error::invalid_args("Expected staking_token, reward_token, reward_rate"))?;

    storage::set(keys::STAKING_TOKEN, &params.staking_token);
    storage::set(keys::REWARD_TOKEN, &params.reward_token);
    storage::set(keys::REWARD_RATE, &params.reward_rate);
    storage::set(keys::OWNER, &caller());
    storage::set(keys::TOTAL_STAKED, &0u64);
    storage::set(keys::TOTAL_REWARDS_PAID, &0u64);
    storage::set(keys::PAUSED, &false);

    emit(&PoolCreated {
        staking_token: params.staking_token,
        reward_token: params.reward_token,
        reward_rate: params.reward_rate,
    });

    Ok(())
}

/// Handle contract calls.
fn call(selector: &[u8], params: &[u8]) -> Result<Vec<u8>> {
    let stake_sel = synor_sdk::method_selector("stake");
    let unstake_sel = synor_sdk::method_selector("unstake");
    let emergency_unstake_sel = synor_sdk::method_selector("emergency_unstake");
    let claim_sel = synor_sdk::method_selector("claim");
    let get_stake_sel = synor_sdk::method_selector("get_stake");
    let get_stakes_sel = synor_sdk::method_selector("get_stakes");
    let get_pool_info_sel = synor_sdk::method_selector("get_pool_info");
    let pause_sel = synor_sdk::method_selector("pause");
    let unpause_sel = synor_sdk::method_selector("unpause");
    let set_reward_rate_sel = synor_sdk::method_selector("set_reward_rate");

    match selector {
        s if s == stake_sel => {
            require!(!is_paused(), Error::invalid_args("Pool is paused"));

            #[derive(BorshDeserialize)]
            struct Args {
                amount: u64,
                lock_days: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected amount, lock_days"))?;

            let stake_id = do_stake(args.amount, args.lock_days)?;
            Ok(borsh::to_vec(&stake_id).unwrap())
        }

        s if s == unstake_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                stake_id: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected stake_id"))?;

            let amount = do_unstake(args.stake_id, false)?;
            Ok(borsh::to_vec(&amount).unwrap())
        }

        s if s == emergency_unstake_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                stake_id: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected stake_id"))?;

            let amount = do_unstake(args.stake_id, true)?;
            Ok(borsh::to_vec(&amount).unwrap())
        }

        s if s == claim_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                stake_id: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected stake_id"))?;

            let rewards = do_claim(args.stake_id)?;
            Ok(borsh::to_vec(&rewards).unwrap())
        }

        s if s == get_stake_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                owner: Address,
                stake_id: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected owner, stake_id"))?;

            let stake = get_stake(&args.owner, args.stake_id);
            Ok(borsh::to_vec(&stake).unwrap())
        }

        s if s == get_stakes_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                owner: Address,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected owner"))?;

            let count = get_stake_count(&args.owner);
            let mut stakes = Vec::new();
            for i in 1..=count {
                if let Some(stake) = get_stake(&args.owner, i) {
                    if stake.active {
                        stakes.push(stake);
                    }
                }
            }
            Ok(borsh::to_vec(&stakes).unwrap())
        }

        s if s == get_pool_info_sel => {
            let info = PoolInfo {
                staking_token: storage::get(keys::STAKING_TOKEN).unwrap_or(Address::zero()),
                reward_token: storage::get(keys::REWARD_TOKEN).unwrap_or(Address::zero()),
                reward_rate: storage::get(keys::REWARD_RATE).unwrap_or(0),
                total_staked: storage::get(keys::TOTAL_STAKED).unwrap_or(0),
                total_rewards_paid: storage::get(keys::TOTAL_REWARDS_PAID).unwrap_or(0),
                paused: is_paused(),
            };
            Ok(borsh::to_vec(&info).unwrap())
        }

        s if s == pause_sel => {
            let owner: Address = storage::get(keys::OWNER).ok_or(Error::Unauthorized)?;
            require_auth!(owner);
            storage::set(keys::PAUSED, &true);
            Ok(Vec::new())
        }

        s if s == unpause_sel => {
            let owner: Address = storage::get(keys::OWNER).ok_or(Error::Unauthorized)?;
            require_auth!(owner);
            storage::set(keys::PAUSED, &false);
            Ok(Vec::new())
        }

        s if s == set_reward_rate_sel => {
            let owner: Address = storage::get(keys::OWNER).ok_or(Error::Unauthorized)?;
            require_auth!(owner);

            #[derive(BorshDeserialize)]
            struct Args {
                new_rate: u64,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Expected new_rate"))?;

            storage::set(keys::REWARD_RATE, &args.new_rate);
            Ok(Vec::new())
        }

        _ => Err(Error::InvalidMethod),
    }
}

// ==================== Internal Functions ====================

fn do_stake(amount: u64, lock_days: u64) -> Result<u64> {
    require!(amount > 0, Error::invalid_args("Amount must be positive"));

    let staker = caller();
    let now = timestamp();

    // Calculate effective amount with bonus
    let multiplier = get_bonus_multiplier(lock_days);
    let effective_amount = (amount as u128 * multiplier as u128 / BPS_DENOMINATOR as u128) as u64;

    // Calculate lock end time
    let lock_until = if lock_days > 0 {
        now + lock_days * SECONDS_PER_DAY
    } else {
        0
    };

    // Create new stake
    let stake_id = increment_stake_count(&staker);
    let stake = StakeInfo {
        id: stake_id,
        amount,
        effective_amount,
        staked_at: now,
        lock_until,
        last_reward_time: now,
        pending_rewards: 0,
        active: true,
    };

    set_stake(&staker, stake_id, &stake);

    // Update total staked
    let total: u64 = storage::get(keys::TOTAL_STAKED).unwrap_or(0);
    storage::set(keys::TOTAL_STAKED, &(total + effective_amount));

    emit(&Staked {
        staker,
        stake_id,
        amount,
        effective_amount,
        lock_until,
    });

    Ok(stake_id)
}

fn do_unstake(stake_id: u64, emergency: bool) -> Result<u64> {
    let staker = caller();
    let now = timestamp();

    let mut stake = get_stake(&staker, stake_id).ok_or(Error::invalid_args("Stake not found"))?;
    require!(stake.active, Error::invalid_args("Stake already withdrawn"));

    // Check lock period (unless emergency)
    if !emergency && stake.lock_until > 0 {
        require!(
            now >= stake.lock_until,
            Error::invalid_args("Stake is still locked")
        );
    }

    // Calculate and claim pending rewards first
    let rewards = calculate_rewards(&stake, now);
    stake.pending_rewards += rewards;

    // Calculate return amount
    let return_amount = if emergency && stake.lock_until > now {
        // Apply penalty for early withdrawal
        let penalty = stake.amount * EMERGENCY_PENALTY_BPS / BPS_DENOMINATOR;
        stake.amount - penalty
    } else {
        stake.amount
    };

    // Mark stake as inactive
    stake.active = false;
    set_stake(&staker, stake_id, &stake);

    // Update total staked
    let total: u64 = storage::get(keys::TOTAL_STAKED).unwrap_or(0);
    storage::set(keys::TOTAL_STAKED, &total.saturating_sub(stake.effective_amount));

    emit(&Unstaked {
        staker,
        stake_id,
        amount: return_amount,
        rewards: stake.pending_rewards,
        emergency,
    });

    Ok(return_amount)
}

fn do_claim(stake_id: u64) -> Result<u64> {
    let staker = caller();
    let now = timestamp();

    let mut stake = get_stake(&staker, stake_id).ok_or(Error::invalid_args("Stake not found"))?;
    require!(stake.active, Error::invalid_args("Stake not active"));

    // Calculate rewards
    let new_rewards = calculate_rewards(&stake, now);
    let total_rewards = stake.pending_rewards + new_rewards;

    require!(total_rewards > 0, Error::invalid_args("No rewards to claim"));

    // Update stake
    stake.pending_rewards = 0;
    stake.last_reward_time = now;
    set_stake(&staker, stake_id, &stake);

    // Update total rewards paid
    let total_paid: u64 = storage::get(keys::TOTAL_REWARDS_PAID).unwrap_or(0);
    storage::set(keys::TOTAL_REWARDS_PAID, &(total_paid + total_rewards));

    emit(&RewardsClaimed {
        staker,
        stake_id,
        amount: total_rewards,
    });

    Ok(total_rewards)
}

fn calculate_rewards(stake: &StakeInfo, now: u64) -> u64 {
    if !stake.active || now <= stake.last_reward_time {
        return 0;
    }

    let reward_rate: u64 = storage::get(keys::REWARD_RATE).unwrap_or(0);
    let duration = now - stake.last_reward_time;

    // rewards = effective_amount * reward_rate * duration / 1e8
    (stake.effective_amount as u128 * reward_rate as u128 * duration as u128 / 100_000_000) as u64
}

// ==================== Events ====================

use synor_sdk::events::{Event, Topic};

#[derive(BorshSerialize)]
struct PoolCreated {
    staking_token: Address,
    reward_token: Address,
    reward_rate: u64,
}

impl Event for PoolCreated {
    fn topics(&self) -> alloc::vec::Vec<Topic> {
        alloc::vec![
            Topic::from_name("PoolCreated"),
            Topic::from(&self.staking_token),
            Topic::from(&self.reward_token),
        ]
    }
    fn data(&self) -> alloc::vec::Vec<u8> {
        self.reward_rate.to_le_bytes().to_vec()
    }
}

#[derive(BorshSerialize)]
struct Staked {
    staker: Address,
    stake_id: u64,
    amount: u64,
    effective_amount: u64,
    lock_until: u64,
}

impl Event for Staked {
    fn topics(&self) -> alloc::vec::Vec<Topic> {
        alloc::vec![
            Topic::from_name("Staked"),
            Topic::from(&self.staker),
            Topic::from(self.stake_id),
        ]
    }
    fn data(&self) -> alloc::vec::Vec<u8> {
        borsh::to_vec(self).unwrap_or_default()
    }
}

#[derive(BorshSerialize)]
struct Unstaked {
    staker: Address,
    stake_id: u64,
    amount: u64,
    rewards: u64,
    emergency: bool,
}

impl Event for Unstaked {
    fn topics(&self) -> alloc::vec::Vec<Topic> {
        alloc::vec![
            Topic::from_name("Unstaked"),
            Topic::from(&self.staker),
            Topic::from(self.stake_id),
        ]
    }
    fn data(&self) -> alloc::vec::Vec<u8> {
        borsh::to_vec(self).unwrap_or_default()
    }
}

#[derive(BorshSerialize)]
struct RewardsClaimed {
    staker: Address,
    stake_id: u64,
    amount: u64,
}

impl Event for RewardsClaimed {
    fn topics(&self) -> alloc::vec::Vec<Topic> {
        alloc::vec![
            Topic::from_name("RewardsClaimed"),
            Topic::from(&self.staker),
            Topic::from(self.stake_id),
        ]
    }
    fn data(&self) -> alloc::vec::Vec<u8> {
        self.amount.to_le_bytes().to_vec()
    }
}