[package]
name = "my_project"
edition = "2024.beta"

[dependencies]
Sui = { git = "https://github.com/MystenLabs/sui.git", subdir = "crates/sui-framework/packages/sui-framework", rev = "framework/testnet" }
ika_dwallet_2pc_mpc = { git = "https://github.com/dwallet-labs/ika.git", subdir = "deployed_contracts/testnet/ika_dwallet_2pc_mpc", rev = "main" }
ika = { git = "https://github.com/dwallet-labs/ika.git", subdir = "deployed_contracts/testnet/ika", rev = "main" }

[addresses]
my_project = "0x0"

pnpm add @ika.xyz/sdk

import { getNetworkConfig, IkaClient } from '@ika.xyz/sdk';
import { getJsonRpcFullnodeUrl, SuiJsonRpcClient } from '@mysten/sui/jsonRpc';
const suiClient = new SuiJsonRpcClient({ url: getJsonRpcFullnodeUrl('testnet'), network: 'testnet' });
const ikaClient = new IkaClient({ suiClient, config: getNetworkConfig('testnet'), cache: true });
await ikaClient.initialize();

use ika::ika::IKA;
use ika_dwallet_2pc_mpc::{
    coordinator::DWalletCoordinator,
    coordinator_inner::{
        DWalletCap, ImportedKeyDWalletCap,
        UnverifiedPresignCap, VerifiedPresignCap,
        UnverifiedPartialUserSignatureCap, VerifiedPartialUserSignatureCap,
        MessageApproval, ImportedKeyMessageApproval
    },
    sessions_manager::SessionIdentifier
};
use sui::{balance::Balance, coin::Coin, sui::SUI};

public struct MyContract has key, store {
    id: UID,
    dwallet_cap: DWalletCap,
    presigns: vector<UnverifiedPresignCap>,
    ika_balance: Balance<IKA>,
    sui_balance: Balance<SUI>,
    dwallet_network_encryption_key_id: ID,
}

fun random_session(coordinator: &mut DWalletCoordinator, ctx: &mut TxContext): SessionIdentifier {
    coordinator.register_session_identifier(ctx.fresh_object_address().to_bytes(), ctx)
}

fun withdraw_payment_coins(self: &mut MyContract, ctx: &mut TxContext): (Coin<IKA>, Coin<SUI>) {
    let ika = self.ika_balance.withdraw_all().into_coin(ctx);
    let sui = self.sui_balance.withdraw_all().into_coin(ctx);
    (ika, sui)
}

fun return_payment_coins(self: &mut MyContract, ika: Coin<IKA>, sui: Coin<SUI>) {
    self.ika_balance.join(ika.into_balance());
    self.sui_balance.join(sui.into_balance());
}

public fun add_ika_balance(self: &mut MyContract, coin: Coin<IKA>) {
    self.ika_balance.join(coin.into_balance());
}

public fun add_sui_balance(self: &mut MyContract, coin: Coin<SUI>) {
    self.sui_balance.join(coin.into_balance());
}

let session = coordinator.register_session_identifier(ctx.fresh_object_address().to_bytes(), ctx);

let (dwallet_cap, _) = coordinator.request_dwallet_dkg_with_public_user_secret_key_share(
    dwallet_network_encryption_key_id, curve,
    centralized_public_key_share_and_proof, user_public_output, public_user_secret_key_share,
    option::none(), // sign_during_dkg_request
    session, &mut ika, &mut sui, ctx,
);

let (dwallet_cap, _) = coordinator.request_dwallet_dkg(
    dwallet_network_encryption_key_id, curve,
    centralized_public_key_share_and_proof, encrypted_centralized_secret_share_and_proof,
    encryption_key_address, user_public_output, signer_public_key,
    option::none(), session, &mut ika, &mut sui, ctx,
);

// Wait for DKG to complete and dWallet to reach AwaitingKeyHolderSignature state
const awaitingDWallet = await ikaClient.getDWalletInParticularState(dwalletId, 'AwaitingKeyHolderSignature');
// The encrypted user secret key share ID comes from the DKG transaction event,
// NOT from the dWallet ID.
const encryptedShare = await ikaClient.getEncryptedUserSecretKeyShare(encryptedUserSecretKeyShareId);

const tx = new Transaction();
const ikaTx = new IkaTransaction({ ikaClient, transaction: tx, userShareEncryptionKeys: keys });
await ikaTx.acceptEncryptedUserShare({
    dWallet: awaitingDWallet,
    encryptedUserSecretKeyShareId: encryptedShare.id,
    userPublicOutput: new Uint8Array(dkgData.userPublicOutput),
});
await suiClient.core.signAndExecuteTransaction({ transaction: tx, signer: keypair });

// dWallet is now Active and ready for signing
const activeDWallet = await ikaClient.getDWalletInParticularState(dwalletId, 'Active');

import { prepareDKGAsync, UserShareEncryptionKeys, Curve, createRandomSessionIdentifier } from '@ika.xyz/sdk';
const keys = await UserShareEncryptionKeys.fromRootSeedKey(seed, Curve.SECP256K1);
const bytesToHash = createRandomSessionIdentifier(); // bytes hashed to derive the session identifier
const dkgData = await prepareDKGAsync(ikaClient, Curve.SECP256K1, keys, bytesToHash, signerAddress);
const networkKey = await ikaClient.getLatestNetworkEncryptionKey();
// dkgData has: userDKGMessage, userPublicOutput, encryptedUserShareAndProof, userSecretKeyShare

let sign_req = coordinator.sign_during_dkg_request(verified_presign, hash_scheme, message, msg_sig);
// Pass option::some(sign_req) instead of option::none() in DKG call

let cap = coordinator.request_global_presign(
    dwallet_network_encryption_key_id, curve, signature_algorithm,
    session, &mut ika, &mut sui, ctx,
);
self.presigns.push_back(cap);

let cap = coordinator.request_presign(
    dwallet_id, signature_algorithm, session, &mut ika, &mut sui, ctx,
);

let is_ready = coordinator.is_presign_valid(&unverified_cap);
let verified = coordinator.verify_presign_cap(unverified_cap, ctx); // fails if not ready

// Pop from pool
let unverified = self.presigns.swap_remove(0);

// Auto-replenish after signing
if (self.presigns.length() < MIN_POOL) {
    let s = random_session(coordinator, ctx);
    self.presigns.push_back(coordinator.request_global_presign(
        self.dwallet_network_encryption_key_id, curve, sig_algo, s, &mut ika, &mut sui, ctx,
    ));
};

// Batch add
let mut i = 0;
while (i < count) {
    let s = random_session(coordinator, ctx);
    self.presigns.push_back(coordinator.request_global_presign(..., s, &mut ika, &mut sui, ctx));
    i = i + 1;
};

// 1. Verify presign
let verified = coordinator.verify_presign_cap(self.presigns.swap_remove(0), ctx);

// 2. Approve message
let approval = coordinator.approve_message(&self.dwallet_cap, sig_algo, hash_scheme, message);

// 3. Sign
let sign_id = coordinator.request_sign_and_return_id(
    verified, approval, message_centralized_signature, session, &mut ika, &mut sui, ctx,
);
// Also: coordinator.request_sign(...) without return ID

import { createUserSignMessageWithPublicOutput, Curve, SignatureAlgorithm, Hash } from '@ika.xyz/sdk';
const completedPresign = await ikaClient.getPresignInParticularState(presignId, 'Completed');
const protocolPublicParameters = await ikaClient.getProtocolPublicParameters(undefined, Curve.SECP256K1);
const msgSig = await createUserSignMessageWithPublicOutput(
    protocolPublicParameters,
    dWallet.state.Active!.public_output,
    dWallet.public_user_secret_key_share,
    completedPresign.presign,
    message,
    Hash.SHA256,
    SignatureAlgorithm.Taproot,
    Curve.SECP256K1,
);
// Pass msgSig as message_centralized_signature to Move

const signSession = await ikaClient.getSignInParticularState(
    signId, Curve.SECP256K1, SignatureAlgorithm.Taproot, 'Completed',
);
const sig = await parseSignatureFromSignOutput(Curve.SECP256K1, SignatureAlgorithm.Taproot, signSession.signature);

let partial_cap = coordinator.request_future_sign(
    self.dwallet_cap.dwallet_id(), verified_presign, message, hash_scheme,
    message_centralized_signature, session, &mut ika, &mut sui, ctx,
);
// Store partial_cap with request for later

// Verify partial sig
let verified_partial = coordinator.verify_partial_user_signature_cap(partial_cap, ctx);

// Create approval
let approval = coordinator.approve_message(&self.dwallet_cap, sig_algo, hash_scheme, message);

// Complete signature
let sign_id = coordinator.request_sign_with_partial_user_signature_and_return_id(
    verified_partial, approval, session, &mut ika, &mut sui, ctx,
);

let ready = coordinator.is_partial_user_signature_valid(&unverified_cap);
let matches = coordinator.match_partial_user_signature_with_message_approval(&verified, &approval);

let imported_cap = coordinator.request_imported_key_dwallet_verification(
    dwallet_network_encryption_key_id, curve, centralized_party_message,
    encrypted_centralized_secret_share_and_proof, encryption_key_address,
    user_public_output, signer_public_key, session, &mut ika, &mut sui, ctx,
);

coordinator.request_make_dwallet_user_secret_key_shares_public(
    dwallet_id, public_user_secret_key_shares, session, &mut ika, &mut sui, ctx,
);

module my_protocol::treasury;

use ika::ika::IKA;
use ika_dwallet_2pc_mpc::{
    coordinator::DWalletCoordinator,
    coordinator_inner::{DWalletCap, UnverifiedPresignCap, UnverifiedPartialUserSignatureCap}
};
use sui::{balance::Balance, coin::Coin, sui::SUI, table::{Self, Table}};

const SECP256K1: u32 = 0;
const TAPROOT: u32 = 1;
const SHA256: u32 = 0;
const MIN_PRESIGNS: u64 = 3;

public struct Treasury has key, store {
    id: UID,
    dwallet_cap: DWalletCap,
    presigns: vector<UnverifiedPresignCap>,
    members: vector<address>,
    approval_threshold: u64,
    proposals: Table<u64, Proposal>,
    next_id: u64,
    ika_balance: Balance<IKA>,
    sui_balance: Balance<SUI>,
    dwallet_network_encryption_key_id: ID,
}

public struct Proposal has store {
    message: vector<u8>,
    partial_cap: Option<UnverifiedPartialUserSignatureCap>,
    approvals: u64,
    voters: Table<address, bool>,
    executed: bool,
}

// Create treasury with shared dWallet
public fun create(
    coordinator: &mut DWalletCoordinator,
    mut ika: Coin<IKA>, mut sui: Coin<SUI>,
    enc_key_id: ID,
    dkg_msg: vector<u8>, user_output: vector<u8>, user_share: vector<u8>,
    session_bytes: vector<u8>,
    members: vector<address>, threshold: u64,
    ctx: &mut TxContext,
) {
    let session = coordinator.register_session_identifier(session_bytes, ctx);
    let (dwallet_cap, _) = coordinator.request_dwallet_dkg_with_public_user_secret_key_share(
        enc_key_id, SECP256K1, dkg_msg, user_output, user_share,
        option::none(), session, &mut ika, &mut sui, ctx,
    );

    let treasury = Treasury {
        id: object::new(ctx), dwallet_cap, presigns: vector::empty(),
        members, approval_threshold: threshold,
        proposals: table::new(ctx), next_id: 0,
        ika_balance: ika.into_balance(), sui_balance: sui.into_balance(),
        dwallet_network_encryption_key_id: enc_key_id,
    };
    transfer::public_share_object(treasury);
}

// Create proposal with future sign (Phase 1)
public fun propose(
    self: &mut Treasury, coordinator: &mut DWalletCoordinator,
    message: vector<u8>, msg_sig: vector<u8>, ctx: &mut TxContext,
): u64 {
    assert!(self.members.contains(&ctx.sender()), 0);
    let (mut ika, mut sui) = self.withdraw_payment_coins(ctx);

    let verified = coordinator.verify_presign_cap(self.presigns.swap_remove(0), ctx);
    let session = random_session(coordinator, ctx);
    let partial = coordinator.request_future_sign(
        self.dwallet_cap.dwallet_id(), verified, message, SHA256,
        msg_sig, session, &mut ika, &mut sui, ctx,
    );

    let id = self.next_id;
    self.next_id = id + 1;
    self.proposals.add(id, Proposal {
        message, partial_cap: option::some(partial),
        approvals: 0, voters: table::new(ctx), executed: false,
    });

    // Auto-replenish
    if (self.presigns.length() < MIN_PRESIGNS) {
        let s = random_session(coordinator, ctx);
        self.presigns.push_back(coordinator.request_global_presign(
            self.dwallet_network_encryption_key_id, SECP256K1, TAPROOT, s,
            &mut ika, &mut sui, ctx,
        ));
    };
    self.return_payment_coins(ika, sui);
    id
}

// Vote
public fun vote(self: &mut Treasury, id: u64, approve: bool, ctx: &TxContext) {
    assert!(self.members.contains(&ctx.sender()), 0);
    let p = self.proposals.borrow_mut(id);
    assert!(!p.voters.contains(ctx.sender()) && !p.executed, 1);
    p.voters.add(ctx.sender(), approve);
    if (approve) { p.approvals = p.approvals + 1; };
}

// Execute after approval (Phase 2)
public fun execute(
    self: &mut Treasury, coordinator: &mut DWalletCoordinator,
    id: u64, ctx: &mut TxContext,
): ID {
    let p = self.proposals.borrow_mut(id);
    assert!(p.approvals >= self.approval_threshold && !p.executed, 2);
    let (mut ika, mut sui) = self.withdraw_payment_coins(ctx);

    let verified = coordinator.verify_partial_user_signature_cap(p.partial_cap.extract(), ctx);
    let approval = coordinator.approve_message(&self.dwallet_cap, TAPROOT, SHA256, p.message);
    let session = random_session(coordinator, ctx);
    let sign_id = coordinator.request_sign_with_partial_user_signature_and_return_id(
        verified, approval, session, &mut ika, &mut sui, ctx,
    );
    p.executed = true;
    self.return_payment_coins(ika, sui);
    sign_id
}

fun random_session(c: &mut DWalletCoordinator, ctx: &mut TxContext): SessionIdentifier {
    c.register_session_identifier(ctx.fresh_object_address().to_bytes(), ctx)
}
fun withdraw_payment_coins(self: &mut Treasury, ctx: &mut TxContext): (Coin<IKA>, Coin<SUI>) {
    (self.ika_balance.withdraw_all().into_coin(ctx), self.sui_balance.withdraw_all().into_coin(ctx))
}
fun return_payment_coins(self: &mut Treasury, ika: Coin<IKA>, sui: Coin<SUI>) {
    self.ika_balance.join(ika.into_balance());
    self.sui_balance.join(sui.into_balance());
}