use privora_sdk_program::prelude::*;
use borsh::{BorshSerialize, BorshDeserialize};

#[derive(BorshSerialize, BorshDeserialize)]
pub enum OrderSide { Buy, Sell }

#[derive(BorshSerialize, BorshDeserialize)]
pub enum OrderStatus { Open, Matched, Cancelled }

#[derive(BorshSerialize, BorshDeserialize)]
pub struct Order {
    pub owner: Pubkey,
    pub price_ref: EncryptedRef<u8>,   // 32 bytes
    pub qty_ref: EncryptedRef<u8>,     // 32 bytes
    pub side: OrderSide,
    pub status: OrderStatus,
    pub order_id: u64,
}

#[derive(BorshSerialize, BorshDeserialize)]
pub struct MatchResult {
    pub buy_order_id: u64,
    pub sell_order_id: u64,
    pub fill_price_ref: EncryptedRef<u8>,
    pub fill_qty_ref: EncryptedRef<u8>,
}

pub fn submit_order(
    accounts: &[AccountInfo],
    instruction_data: &[u8],
) -> ProgramResult {
    let account_info_iter = &mut accounts.iter();
    let order_account = next_account_info(account_info_iter)?;
    let owner_account = next_account_info(account_info_iter)?;

    // Verify signer
    if !owner_account.is_signer {
        return Err(ProgramError::MissingRequiredSignature);
    }

    // Parse instruction data
    #[derive(BorshDeserialize)]
    struct SubmitOrderData {
        price_ref: EncryptedRef<u8>,
        qty_ref: EncryptedRef<u8>,
        side: OrderSide,
    }
    let data = SubmitOrderData::try_from_slice(instruction_data)?;

    // Create order
    let order = Order {
        owner: *owner_account.key,
        price_ref: data.price_ref,
        qty_ref: data.qty_ref,
        side: data.side,
        status: OrderStatus::Open,
        order_id: generate_order_id(),
    };

    order.serialize(&mut *order_account.data.borrow_mut())?;
    Ok(())
}

pub fn match_orders(
    accounts: &[AccountInfo],
) -> ProgramResult {
    let account_info_iter = &mut accounts.iter();
    let buy_order_account = next_account_info(account_info_iter)?;
    let sell_order_account = next_account_info(account_info_iter)?;
    let match_result_account = next_account_info(account_info_iter)?;

    // Load orders
    let mut buy_order = Order::try_from_slice(&buy_order_account.data.borrow())?;
    let mut sell_order = Order::try_from_slice(&sell_order_account.data.borrow())?;

    // Verify order sides and status
    if buy_order.side != OrderSide::Buy || sell_order.side != OrderSide::Sell {
        return Err(ProgramError::InvalidArgument);
    }
    if buy_order.status != OrderStatus::Open || sell_order.status != OrderStatus::Open {
        return Err(ProgramError::InvalidArgument);
    }

    // Load encrypted values
    let buy_price = buy_order.price_ref.load()?;
    let sell_price = sell_order.price_ref.load()?;
    let buy_qty = buy_order.qty_ref.load()?;
    let sell_qty = sell_order.qty_ref.load()?;

    // FHE comparison: can match if buy_price >= sell_price
    let _can_match: EncryptedBool = buy_price.ge(&sell_price)?;

    // Calculate fill quantity: min(buy_qty, sell_qty)
    let fill_qty: Encrypted<u8> = buy_qty.min(&sell_qty)?;

    // Fill price is sell price (standard matching)
    let fill_price_ref = sell_order.price_ref;
    let fill_qty_ref = fill_qty.store()?;

    // Update order statuses
    buy_order.status = OrderStatus::Matched;
    sell_order.status = OrderStatus::Matched;

    buy_order.serialize(&mut *buy_order_account.data.borrow_mut())?;
    sell_order.serialize(&mut *sell_order_account.data.borrow_mut())?;

    // Store match result
    let result = MatchResult {
        buy_order_id: buy_order.order_id,
        sell_order_id: sell_order.order_id,
        fill_price_ref,
        fill_qty_ref,
    };
    result.serialize(&mut *match_result_account.data.borrow_mut())?;

    Ok(())
}

async function submitOrder(price: number, quantity: number, side: 'buy' | 'sell') {
  const privora = await Privora.connect('http://localhost:8899');
  const userCrypto = privora.userCrypto(keypair);

  // Encrypt with user recovery
  const encPrice = privora.encrypt(price, 'u8').withUserRecovery(userCrypto);
  const encQty = privora.encrypt(quantity, 'u8').withUserRecovery(userCrypto);

  // Submit to sequencer
  const priceHash = await privora.submit(encPrice);
  const qtyHash = await privora.submit(encQty);

  // Build instruction
  const data = Buffer.concat([
    Buffer.from([1]), // SubmitOrder
    Buffer.from(priceHash, 'hex'),
    Buffer.from(qtyHash, 'hex'),
    Buffer.from([side === 'buy' ? 0 : 1]),
  ]);

  // Send transaction
  return privora
    .transaction()
    .add({ programId: PROGRAM_ID, keys: [...], data })
    .withFheData([priceHash, qtyHash])
    .sign(keypair)
    .send();
}

#[test]
fn test_order_matching() {
    let mut env = FheTestEnv::new();
    env.deploy_program(PROGRAM_ID, include_bytes!("../program.so"));

    // Create orders: buy@100, sell@95
    let buy_price = env.encrypt_and_submit_u8(100);
    let buy_qty = env.encrypt_and_submit_u8(50);
    let sell_price = env.encrypt_and_submit_u8(95);
    let sell_qty = env.encrypt_and_submit_u8(30);

    // Submit orders and match...

    // Verify: fill_price = 95 (sell price)
    env.assert_encrypted_eq_u8(&fill_price_hash, 95);

    // Verify: fill_qty = min(50, 30) = 30
    env.assert_encrypted_eq_u8(&fill_qty_hash, 30);
}