//! Confidential Token Contract
//!
//! This contract implements privacy-preserving token transfers using:
//! - Pedersen commitments for hidden amounts
//! - Ring signatures for sender anonymity
//! - Stealth addresses for receiver privacy
//! - Range proofs to prevent inflation
//!
//! # Privacy Properties
//!
//! - **Amount Privacy**: All balances and transfer amounts are hidden in commitments
//! - **Sender Privacy**: Ring signatures hide which input is being spent
//! - **Receiver Privacy**: Stealth addresses make recipients unlinkable
//! - **Linkability Prevention**: Key images prevent double-spending without revealing identity
//!
//! # Methods
//!
//! - `init(name, symbol, decimals)` - Initialize the token
//! - `mint(commitment, range_proof, stealth_address)` - Mint new tokens (owner only)
//! - `transfer(inputs, outputs)` - Transfer tokens confidentially
//! - `burn(input, amount_reveal)` - Burn tokens (reveals amount)
//! - `is_key_image_used(key_image)` - Check if a key image was used

#![no_std]

extern crate alloc;

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

// =============================================================================
// EVENT TOPICS (pre-computed hashes for event names)
// =============================================================================

/// Hash an event name to a 32-byte topic
fn event_topic(name: &[u8]) -> [u8; 32] {
    use synor_sdk::crypto::blake3_hash;
    blake3_hash(name).0
}

// =============================================================================
// STORAGE KEYS
// =============================================================================

mod keys {
    /// Contract owner address
    pub const OWNER: &[u8] = b"ctoken:owner";
    /// Whether the contract has been initialized
    pub const INITIALIZED: &[u8] = b"ctoken:initialized";
    /// Token name
    pub const NAME: &[u8] = b"ctoken:name";
    /// Token symbol
    pub const SYMBOL: &[u8] = b"ctoken:symbol";
    /// Token decimals
    pub const DECIMALS: &[u8] = b"ctoken:decimals";
    /// Total supply commitment (sum of all outputs)
    pub const TOTAL_SUPPLY_COMMITMENT: &[u8] = b"ctoken:total_supply";
    /// Used key images (for double-spend prevention)
    pub const KEY_IMAGES: &[u8] = b"ctoken:key_images";
    /// Unspent transaction outputs (UTXOs)
    pub const UTXOS: &[u8] = b"ctoken:utxos";
    /// UTXO count for generating unique IDs
    pub const UTXO_COUNT: &[u8] = b"ctoken:utxo_count";
}

// =============================================================================
// DATA STRUCTURES
// =============================================================================

/// A confidential unspent transaction output (UTXO)
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct ConfidentialUtxo {
    /// Unique identifier
    pub id: u64,
    /// Amount commitment (32 bytes)
    pub commitment: [u8; 32],
    /// Stealth address (recipient's one-time address, 32 bytes)
    pub stealth_address: [u8; 32],
    /// Ephemeral public key (for recipient to derive key, 32 bytes)
    pub ephemeral_pubkey: [u8; 32],
    /// Range proof (proves amount is valid)
    pub range_proof: Vec<u8>,
    /// Encrypted amount (for recipient to decrypt)
    pub encrypted_amount: Option<[u8; 40]>,
    /// Whether this UTXO has been spent
    pub spent: bool,
    /// Block number when created
    pub created_at: u64,
}

/// A confidential transfer input
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct TransferInput {
    /// UTXO ID being spent
    pub utxo_id: u64,
    /// Key image (for double-spend prevention, 32 bytes)
    pub key_image: [u8; 32],
    /// Ring of public keys (each 32 bytes)
    pub ring: Vec<[u8; 32]>,
    /// Ring signature components
    pub ring_signature: RingSignatureData,
}

/// Ring signature data (LSAG format)
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct RingSignatureData {
    /// Initial challenge c_0 (32 bytes)
    pub c0: [u8; 32],
    /// Response values s_i (each 32 bytes)
    pub responses: Vec<[u8; 32]>,
}

/// A confidential transfer output
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct TransferOutput {
    /// Amount commitment (32 bytes)
    pub commitment: [u8; 32],
    /// Stealth address (32 bytes)
    pub stealth_address: [u8; 32],
    /// Ephemeral public key (32 bytes)
    pub ephemeral_pubkey: [u8; 32],
    /// Range proof
    pub range_proof: Vec<u8>,
    /// Encrypted amount (optional)
    pub encrypted_amount: Option<[u8; 40]>,
}

/// Balance proof (proves sum of inputs equals sum of outputs)
#[derive(Clone, Debug, BorshSerialize, BorshDeserialize)]
pub struct BalanceProof {
    /// Excess public key (32 bytes)
    pub excess_pubkey: [u8; 32],
    /// Schnorr signature over the excess (64 bytes)
    pub signature: [u8; 64],
}

// =============================================================================
// EVENTS
// =============================================================================

#[derive(BorshSerialize)]
pub struct TokenMinted {
    pub utxo_id: u64,
    pub commitment: [u8; 32],
    pub stealth_address: [u8; 32],
}

#[derive(BorshSerialize)]
pub struct TokenTransferred {
    pub input_count: u32,
    pub output_count: u32,
    pub key_images: Vec<[u8; 32]>,
    pub new_utxo_ids: Vec<u64>,
}

#[derive(BorshSerialize)]
pub struct TokenBurned {
    pub utxo_id: u64,
    pub amount: u64,
    pub burner: Address,
}

// =============================================================================
// STORAGE HELPERS
// =============================================================================

fn get_owner() -> Option<Address> {
    storage::get::<Address>(keys::OWNER)
}

fn is_owner(addr: &Address) -> bool {
    get_owner().map(|o| o == *addr).unwrap_or(false)
}

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

fn get_utxo_count() -> u64 {
    storage::get::<u64>(keys::UTXO_COUNT).unwrap_or(0)
}

fn set_utxo_count(count: u64) {
    storage::set(keys::UTXO_COUNT, &count);
}

fn get_utxo(id: u64) -> Option<ConfidentialUtxo> {
    storage::get_with_suffix::<ConfidentialUtxo>(keys::UTXOS, &id.to_le_bytes())
}

fn set_utxo(id: u64, utxo: &ConfidentialUtxo) {
    storage::set_with_suffix(keys::UTXOS, &id.to_le_bytes(), utxo);
}

fn is_key_image_used(key_image: &[u8; 32]) -> bool {
    storage::get_with_suffix::<bool>(keys::KEY_IMAGES, key_image).unwrap_or(false)
}

fn mark_key_image_used(key_image: &[u8; 32]) {
    storage::set_with_suffix(keys::KEY_IMAGES, key_image, &true);
}

// =============================================================================
// ENTRY POINTS
// =============================================================================

synor_sdk::entry_point!(init, call);

/// Initialize the confidential token contract.
fn init(params: &[u8]) -> Result<()> {
    require!(!is_initialized(), Error::invalid_args("Already initialized"));

    #[derive(BorshDeserialize)]
    struct InitParams {
        name: String,
        symbol: String,
        decimals: u8,
    }

    let params = InitParams::try_from_slice(params)
        .map_err(|_| Error::invalid_args("Invalid init params"))?;

    require!(!params.name.is_empty(), Error::invalid_args("Name required"));
    require!(!params.symbol.is_empty(), Error::invalid_args("Symbol required"));
    require!(params.decimals <= 18, Error::invalid_args("Decimals must be <= 18"));

    storage::set(keys::OWNER, &caller());
    storage::set(keys::INITIALIZED, &true);
    storage::set(keys::NAME, &params.name);
    storage::set(keys::SYMBOL, &params.symbol);
    storage::set(keys::DECIMALS, &params.decimals);
    set_utxo_count(0);

    Ok(())
}

/// Handle contract method calls.
fn call(selector: &[u8], params: &[u8]) -> Result<Vec<u8>> {
    let mint_sel = synor_sdk::method_selector("mint");
    let transfer_sel = synor_sdk::method_selector("transfer");
    let burn_sel = synor_sdk::method_selector("burn");
    let name_sel = synor_sdk::method_selector("name");
    let symbol_sel = synor_sdk::method_selector("symbol");
    let decimals_sel = synor_sdk::method_selector("decimals");
    let owner_sel = synor_sdk::method_selector("owner");
    let get_utxo_sel = synor_sdk::method_selector("get_utxo");
    let is_key_image_used_sel = synor_sdk::method_selector("is_key_image_used");

    match selector {
        // ===== Read Methods =====

        s if s == name_sel => {
            let name = storage::get::<String>(keys::NAME).unwrap_or_default();
            Ok(borsh::to_vec(&name).unwrap())
        }

        s if s == symbol_sel => {
            let symbol = storage::get::<String>(keys::SYMBOL).unwrap_or_default();
            Ok(borsh::to_vec(&symbol).unwrap())
        }

        s if s == decimals_sel => {
            let decimals = storage::get::<u8>(keys::DECIMALS).unwrap_or(18);
            Ok(borsh::to_vec(&decimals).unwrap())
        }

        s if s == owner_sel => {
            let owner = get_owner().unwrap_or(Address::zero());
            Ok(borsh::to_vec(&owner).unwrap())
        }

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

            let utxo = get_utxo(args.utxo_id);
            Ok(borsh::to_vec(&utxo).unwrap())
        }

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

            let used = is_key_image_used(&args.key_image);
            Ok(borsh::to_vec(&used).unwrap())
        }

        // ===== Write Methods =====

        s if s == mint_sel => {
            // Only owner can mint
            let owner = get_owner().ok_or(Error::Unauthorized)?;
            require_auth!(owner);

            #[derive(BorshDeserialize)]
            struct Args {
                commitment: [u8; 32],
                stealth_address: [u8; 32],
                ephemeral_pubkey: [u8; 32],
                range_proof: Vec<u8>,
                encrypted_amount: Option<[u8; 40]>,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Invalid mint params"))?;

            // Create new UTXO
            let utxo_id = get_utxo_count();
            let utxo = ConfidentialUtxo {
                id: utxo_id,
                commitment: args.commitment,
                stealth_address: args.stealth_address,
                ephemeral_pubkey: args.ephemeral_pubkey,
                range_proof: args.range_proof,
                encrypted_amount: args.encrypted_amount,
                spent: false,
                created_at: timestamp(),
            };

            set_utxo(utxo_id, &utxo);
            set_utxo_count(utxo_id + 1);

            // Emit event
            emit_raw(
                &[event_topic(b"TokenMinted")],
                &borsh::to_vec(&TokenMinted {
                    utxo_id,
                    commitment: args.commitment,
                    stealth_address: args.stealth_address,
                })
                .unwrap(),
            );

            Ok(borsh::to_vec(&utxo_id).unwrap())
        }

        s if s == transfer_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                inputs: Vec<TransferInput>,
                outputs: Vec<TransferOutput>,
                balance_proof: BalanceProof,
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Invalid transfer params"))?;

            require!(!args.inputs.is_empty(), Error::invalid_args("Need at least 1 input"));
            require!(!args.outputs.is_empty(), Error::invalid_args("Need at least 1 output"));

            let mut input_commitments = Vec::new();
            let mut key_images = Vec::new();

            // Validate inputs
            for input in &args.inputs {
                // Get the UTXO
                let utxo = get_utxo(input.utxo_id)
                    .ok_or_else(|| Error::invalid_args("UTXO not found"))?;

                // Check not already spent
                require!(!utxo.spent, Error::invalid_args("UTXO already spent"));

                // Check key image not used (double-spend prevention)
                require!(
                    !is_key_image_used(&input.key_image),
                    Error::invalid_args("Key image already used")
                );

                // Verify ring signature (simplified - full verification would be more complex)
                require!(
                    verify_ring_signature_basic(&input),
                    Error::invalid_args("Invalid ring signature")
                );

                input_commitments.push(utxo.commitment);
                key_images.push(input.key_image);
            }

            // Verify balance proof (simplified)
            require!(
                verify_balance_proof(&args.balance_proof),
                Error::invalid_args("Invalid balance proof")
            );

            // Mark inputs as spent and key images as used
            for input in &args.inputs {
                mark_key_image_used(&input.key_image);

                if let Some(mut utxo) = get_utxo(input.utxo_id) {
                    utxo.spent = true;
                    set_utxo(input.utxo_id, &utxo);
                }
            }

            // Create output UTXOs
            let mut new_utxo_ids = Vec::new();
            for output in &args.outputs {
                let utxo_id = get_utxo_count();
                let utxo = ConfidentialUtxo {
                    id: utxo_id,
                    commitment: output.commitment,
                    stealth_address: output.stealth_address,
                    ephemeral_pubkey: output.ephemeral_pubkey,
                    range_proof: output.range_proof.clone(),
                    encrypted_amount: output.encrypted_amount,
                    spent: false,
                    created_at: timestamp(),
                };
                set_utxo(utxo_id, &utxo);
                set_utxo_count(utxo_id + 1);
                new_utxo_ids.push(utxo_id);
            }

            // Emit event
            emit_raw(
                &[event_topic(b"TokenTransferred")],
                &borsh::to_vec(&TokenTransferred {
                    input_count: args.inputs.len() as u32,
                    output_count: args.outputs.len() as u32,
                    key_images,
                    new_utxo_ids: new_utxo_ids.clone(),
                })
                .unwrap(),
            );

            Ok(borsh::to_vec(&new_utxo_ids).unwrap())
        }

        s if s == burn_sel => {
            #[derive(BorshDeserialize)]
            struct Args {
                input: TransferInput,
                amount: u64,  // Revealed amount for burning
            }
            let args = Args::try_from_slice(params)
                .map_err(|_| Error::invalid_args("Invalid burn params"))?;

            // Get the UTXO
            let utxo = get_utxo(args.input.utxo_id)
                .ok_or_else(|| Error::invalid_args("UTXO not found"))?;

            // Check not already spent
            require!(!utxo.spent, Error::invalid_args("UTXO already spent"));

            // Check key image
            require!(
                !is_key_image_used(&args.input.key_image),
                Error::invalid_args("Key image already used")
            );

            // Verify ring signature
            require!(
                verify_ring_signature_basic(&args.input),
                Error::invalid_args("Invalid ring signature")
            );

            // Mark as spent
            mark_key_image_used(&args.input.key_image);
            let mut utxo = utxo;
            utxo.spent = true;
            set_utxo(args.input.utxo_id, &utxo);

            // Emit event with revealed amount
            emit_raw(
                &[event_topic(b"TokenBurned")],
                &borsh::to_vec(&TokenBurned {
                    utxo_id: args.input.utxo_id,
                    amount: args.amount,
                    burner: caller(),
                })
                .unwrap(),
            );

            Ok(borsh::to_vec(&true).unwrap())
        }

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

// =============================================================================
// VERIFICATION HELPERS
// =============================================================================

/// Basic ring signature verification (structure validation)
/// Full cryptographic verification would require the privacy crate
fn verify_ring_signature_basic(input: &TransferInput) -> bool {
    // Check ring has valid size
    if input.ring.len() < 2 || input.ring.len() > 32 {
        return false;
    }

    // Check responses count matches ring size
    if input.ring_signature.responses.len() != input.ring.len() {
        return false;
    }

    // Check key image is not zero
    if input.key_image == [0u8; 32] {
        return false;
    }

    // In a full implementation, we would verify the LSAG signature here
    // by recomputing the challenge chain and checking if it closes
    true
}

/// Basic balance proof verification (structure validation)
fn verify_balance_proof(proof: &BalanceProof) -> bool {
    // Check excess pubkey is not zero
    if proof.excess_pubkey == [0u8; 32] {
        return false;
    }

    // Check signature is not zero
    if proof.signature == [0u8; 64] {
        return false;
    }

    // In a full implementation, we would verify the Schnorr signature
    // and check that sum(inputs) - sum(outputs) = excess * H
    true
}
