//! Simple AMM DEX Contract
//!
//! A Uniswap V2-style automated market maker for Synor.
//!
//! # Features
//! - Liquidity pools for token pairs
//! - Constant product formula (x * y = k)
//! - LP tokens for liquidity providers
//! - Swap with 0.3% fee
//!
//! # Methods
//! - `init(token_a, token_b)` - Create a new pair
//! - `add_liquidity(amount_a, amount_b, min_lp)` - Add liquidity
//! - `remove_liquidity(lp_amount, min_a, min_b)` - Remove liquidity
//! - `swap_a_for_b(amount_in, min_out)` - Swap token A for B
//! - `swap_b_for_a(amount_in, min_out)` - Swap token B for A
//! - `get_reserves() -> (u64, u64)` - Get current reserves
//! - `get_lp_balance(owner) -> u64` - Get LP token balance

#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use borsh::BorshDeserialize;
use synor_sdk::prelude::*;
use synor_sdk::require;

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

/// Swap fee: 0.3% (30 basis points)
const SWAP_FEE_BPS: u64 = 30;
const BPS_DENOMINATOR: u64 = 10000;

/// Minimum liquidity locked forever (prevents division by zero attacks)
const MINIMUM_LIQUIDITY: u64 = 1000;

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

mod keys {
    pub const TOKEN_A: &[u8] = b"dex:token_a";
    pub const TOKEN_B: &[u8] = b"dex:token_b";
    pub const RESERVE_A: &[u8] = b"dex:reserve_a";
    pub const RESERVE_B: &[u8] = b"dex:reserve_b";
    pub const LP_TOTAL_SUPPLY: &[u8] = b"dex:lp_total";
    pub const LP_BALANCES: &[u8] = b"dex:lp_balances";
    pub const OWNER: &[u8] = b"dex:owner";
    pub const INITIALIZED: &[u8] = b"dex:initialized";
}

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

fn lp_balances() -> storage::Map<[u8; 34], u64> {
    storage::Map::new(keys::LP_BALANCES)
}

fn get_reserves() -> (u64, u64) {
    let reserve_a: u64 = storage::get(keys::RESERVE_A).unwrap_or(0);
    let reserve_b: u64 = storage::get(keys::RESERVE_B).unwrap_or(0);
    (reserve_a, reserve_b)
}

fn set_reserves(reserve_a: u64, reserve_b: u64) {
    storage::set(keys::RESERVE_A, &reserve_a);
    storage::set(keys::RESERVE_B, &reserve_b);
}

fn get_lp_total_supply() -> u64 {
    storage::get(keys::LP_TOTAL_SUPPLY).unwrap_or(0)
}

fn set_lp_total_supply(supply: u64) {
    storage::set(keys::LP_TOTAL_SUPPLY, &supply);
}

fn mint_lp(to: &Address, amount: u64) {
    let balance = lp_balances().get_or_default(&to.0);
    lp_balances().set(&to.0, &(balance + amount));
    set_lp_total_supply(get_lp_total_supply() + amount);
}

fn burn_lp(from: &Address, amount: u64) -> Result<()> {
    let balance = lp_balances().get_or_default(&from.0);
    require!(
        balance >= amount,
        Error::InsufficientBalance {
            required: amount,
            available: balance
        }
    );
    lp_balances().set(&from.0, &(balance - amount));
    set_lp_total_supply(get_lp_total_supply().saturating_sub(amount));
    Ok(())
}

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

synor_sdk::entry_point!(init, call);

/// Initialize the DEX pair.
fn init(params: &[u8]) -> Result<()> {
    #[derive(BorshDeserialize)]
    struct InitParams {
        token_a: Address,
        token_b: Address,
    }

    let params = InitParams::try_from_slice(params)
        .map_err(|_| Error::invalid_args("Expected token_a, token_b"))?;

    require!(
        params.token_a != params.token_b,
        Error::invalid_args("Tokens must be different")
    );

    // Store token addresses
    storage::set(keys::TOKEN_A, &params.token_a);
    storage::set(keys::TOKEN_B, &params.token_b);
    storage::set(keys::OWNER, &caller());
    storage::set(keys::INITIALIZED, &true);

    // Initialize reserves to zero
    set_reserves(0, 0);
    set_lp_total_supply(0);

    emit(&PairCreated {
        token_a: params.token_a,
        token_b: params.token_b,
    });

    Ok(())
}

/// Handle contract calls.
fn call(selector: &[u8], params: &[u8]) -> Result<Vec<u8>> {
    let add_liquidity_sel = synor_sdk::method_selector("add_liquidity");
    let remove_liquidity_sel = synor_sdk::method_selector("remove_liquidity");
    let swap_a_for_b_sel = synor_sdk::method_selector("swap_a_for_b");
    let swap_b_for_a_sel = synor_sdk::method_selector("swap_b_for_a");
    let get_reserves_sel = synor_sdk::method_selector("get_reserves");
    let get_lp_balance_sel = synor_sdk::method_selector("get_lp_balance");
    let get_tokens_sel = synor_sdk::method_selector("get_tokens");
    let quote_sel = synor_sdk::method_selector("quote");

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

            let lp_minted = do_add_liquidity(args.amount_a, args.amount_b, args.min_lp)?;
            Ok(borsh::to_vec(&lp_minted).unwrap())
        }

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

            let (amount_a, amount_b) = do_remove_liquidity(args.lp_amount, args.min_a, args.min_b)?;
            Ok(borsh::to_vec(&(amount_a, amount_b)).unwrap())
        }

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

            let amount_out = do_swap(args.amount_in, args.min_out, true)?;
            Ok(borsh::to_vec(&amount_out).unwrap())
        }

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

            let amount_out = do_swap(args.amount_in, args.min_out, false)?;
            Ok(borsh::to_vec(&amount_out).unwrap())
        }

        s if s == get_reserves_sel => {
            let reserves = get_reserves();
            Ok(borsh::to_vec(&reserves).unwrap())
        }

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

            let balance = lp_balances().get_or_default(&args.owner.0);
            Ok(borsh::to_vec(&balance).unwrap())
        }

        s if s == get_tokens_sel => {
            let token_a: Address = storage::get(keys::TOKEN_A).unwrap_or(Address::zero());
            let token_b: Address = storage::get(keys::TOKEN_B).unwrap_or(Address::zero());
            Ok(borsh::to_vec(&(token_a, token_b)).unwrap())
        }

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

            let (reserve_a, reserve_b) = get_reserves();
            let (reserve_in, reserve_out) = if args.a_to_b {
                (reserve_a, reserve_b)
            } else {
                (reserve_b, reserve_a)
            };

            let amount_out = calculate_output(args.amount_in, reserve_in, reserve_out);
            Ok(borsh::to_vec(&amount_out).unwrap())
        }

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

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

fn do_add_liquidity(amount_a: u64, amount_b: u64, min_lp: u64) -> Result<u64> {
    require!(amount_a > 0 && amount_b > 0, Error::invalid_args("Amounts must be positive"));

    let (reserve_a, reserve_b) = get_reserves();
    let lp_supply = get_lp_total_supply();

    let lp_to_mint = if lp_supply == 0 {
        // First liquidity provider - use geometric mean
        let liquidity = sqrt(amount_a as u128 * amount_b as u128) as u64;
        require!(
            liquidity > MINIMUM_LIQUIDITY,
            Error::invalid_args("Initial liquidity too low")
        );
        // Lock minimum liquidity forever
        mint_lp(&Address::zero(), MINIMUM_LIQUIDITY);
        liquidity - MINIMUM_LIQUIDITY
    } else {
        // Proportional to existing liquidity
        let lp_a = (amount_a as u128 * lp_supply as u128 / reserve_a as u128) as u64;
        let lp_b = (amount_b as u128 * lp_supply as u128 / reserve_b as u128) as u64;
        lp_a.min(lp_b)
    };

    require!(
        lp_to_mint >= min_lp,
        Error::invalid_args("Insufficient LP tokens minted")
    );

    // Update reserves
    set_reserves(reserve_a + amount_a, reserve_b + amount_b);

    // Mint LP tokens
    let provider = caller();
    mint_lp(&provider, lp_to_mint);

    emit(&LiquidityAdded {
        provider,
        amount_a,
        amount_b,
        lp_minted: lp_to_mint,
    });

    Ok(lp_to_mint)
}

fn do_remove_liquidity(lp_amount: u64, min_a: u64, min_b: u64) -> Result<(u64, u64)> {
    require!(lp_amount > 0, Error::invalid_args("Amount must be positive"));

    let (reserve_a, reserve_b) = get_reserves();
    let lp_supply = get_lp_total_supply();

    require!(lp_supply > 0, Error::invalid_args("No liquidity"));

    // Calculate proportional amounts
    let amount_a = (lp_amount as u128 * reserve_a as u128 / lp_supply as u128) as u64;
    let amount_b = (lp_amount as u128 * reserve_b as u128 / lp_supply as u128) as u64;

    require!(
        amount_a >= min_a && amount_b >= min_b,
        Error::invalid_args("Slippage too high")
    );

    // Burn LP tokens
    let provider = caller();
    burn_lp(&provider, lp_amount)?;

    // Update reserves
    set_reserves(reserve_a - amount_a, reserve_b - amount_b);

    emit(&LiquidityRemoved {
        provider,
        amount_a,
        amount_b,
        lp_burned: lp_amount,
    });

    Ok((amount_a, amount_b))
}

fn do_swap(amount_in: u64, min_out: u64, a_to_b: bool) -> Result<u64> {
    require!(amount_in > 0, Error::invalid_args("Amount must be positive"));

    let (reserve_a, reserve_b) = get_reserves();

    let (reserve_in, reserve_out) = if a_to_b {
        (reserve_a, reserve_b)
    } else {
        (reserve_b, reserve_a)
    };

    require!(reserve_in > 0 && reserve_out > 0, Error::invalid_args("No liquidity"));

    let amount_out = calculate_output(amount_in, reserve_in, reserve_out);

    require!(
        amount_out >= min_out,
        Error::invalid_args("Slippage too high")
    );

    // Update reserves
    let (new_a, new_b) = if a_to_b {
        (reserve_a + amount_in, reserve_b - amount_out)
    } else {
        (reserve_a - amount_out, reserve_b + amount_in)
    };
    set_reserves(new_a, new_b);

    emit(&Swap {
        trader: caller(),
        amount_in,
        amount_out,
        a_to_b,
    });

    Ok(amount_out)
}

/// Calculate output amount using constant product formula with fee.
/// amount_out = (amount_in * fee_multiplier * reserve_out) / (reserve_in + amount_in * fee_multiplier)
fn calculate_output(amount_in: u64, reserve_in: u64, reserve_out: u64) -> u64 {
    let amount_in_with_fee = amount_in as u128 * (BPS_DENOMINATOR - SWAP_FEE_BPS) as u128;
    let numerator = amount_in_with_fee * reserve_out as u128;
    let denominator = reserve_in as u128 * BPS_DENOMINATOR as u128 + amount_in_with_fee;

    if denominator == 0 {
        0
    } else {
        (numerator / denominator) as u64
    }
}

/// Integer square root using Newton's method.
fn sqrt(n: u128) -> u128 {
    if n == 0 {
        return 0;
    }
    let mut x = n;
    let mut y = (x + 1) / 2;
    while y < x {
        x = y;
        y = (x + n / x) / 2;
    }
    x
}

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

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

#[derive(borsh::BorshSerialize)]
struct PairCreated {
    token_a: Address,
    token_b: Address,
}

impl Event for PairCreated {
    fn topics(&self) -> alloc::vec::Vec<Topic> {
        alloc::vec![
            Topic::from_name("PairCreated"),
            Topic::from(&self.token_a),
            Topic::from(&self.token_b),
        ]
    }
    fn data(&self) -> alloc::vec::Vec<u8> {
        alloc::vec::Vec::new()
    }
}

#[derive(borsh::BorshSerialize)]
struct LiquidityAdded {
    provider: Address,
    amount_a: u64,
    amount_b: u64,
    lp_minted: u64,
}

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

#[derive(borsh::BorshSerialize)]
struct LiquidityRemoved {
    provider: Address,
    amount_a: u64,
    amount_b: u64,
    lp_burned: u64,
}

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

#[derive(borsh::BorshSerialize)]
struct Swap {
    trader: Address,
    amount_in: u64,
    amount_out: u64,
    a_to_b: bool,
}

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