newton-chainio 0.5.2

//! PolicyClientRegistry controller for managing on-chain policy client registration

use crate::error::ChainIoError;
use alloy::{primitives::Address, rpc::types::TransactionReceipt};
use eigensdk::common::{get_provider, get_signer};
use newton_core::policy_client_registry::{IPolicyClientRegistry::ClientRecord, PolicyClientRegistry};
use tracing::info;

/// Controller for interacting with the PolicyClientRegistry contract
#[derive(Debug)]
pub struct PolicyClientRegistryController {
    rpc_url: String,
    signer: Option<String>,
    registry_address: Address,
}

impl PolicyClientRegistryController {
    /// Create a new controller with signer for write operations
    pub fn new(signer: String, rpc_url: String, registry_address: Address) -> Self {
        PolicyClientRegistryController {
            signer: Some(signer),
            rpc_url,
            registry_address,
        }
    }

    /// Create a read-only controller (no signer required)
    pub fn new_read_only(rpc_url: String, registry_address: Address) -> Self {
        PolicyClientRegistryController {
            signer: None,
            rpc_url,
            registry_address,
        }
    }

    /// Get the signer, returning an error if not configured
    #[allow(clippy::result_large_err)]
    fn require_signer(&self) -> Result<&str, ChainIoError> {
        self.signer
            .as_deref()
            .ok_or_else(|| ChainIoError::CreateNewTaskCallFail {
                reason: "signer required for write operations".to_string(),
            })
    }

    /// Register a policy client with the registry. The caller becomes the registered owner.
    pub async fn register_client(&self, client: Address) -> Result<TransactionReceipt, ChainIoError> {
        let signer = self.require_signer()?;
        let wallet = get_signer(signer, &self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, wallet);

        info!(
            "Registering policy client {} with registry {}",
            client, self.registry_address
        );

        let pending_tx = registry
            .registerClient(client)
            .send()
            .await
            .map_err(ChainIoError::ContractError)?;

        let receipt = pending_tx
            .get_receipt()
            .await
            .map_err(ChainIoError::AlloyProviderError)?;

        info!(
            "Policy client {} registered successfully, tx: {}",
            client, receipt.transaction_hash
        );
        Ok(receipt)
    }

    /// Deactivate a registered policy client. Only callable by the registered owner.
    pub async fn deactivate_client(&self, client: Address) -> Result<TransactionReceipt, ChainIoError> {
        let signer = self.require_signer()?;
        let wallet = get_signer(signer, &self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, wallet);

        info!(
            "Deactivating policy client {} in registry {}",
            client, self.registry_address
        );

        let pending_tx = registry
            .deactivateClient(client)
            .send()
            .await
            .map_err(ChainIoError::ContractError)?;

        let receipt = pending_tx
            .get_receipt()
            .await
            .map_err(ChainIoError::AlloyProviderError)?;

        info!("Policy client {} deactivated, tx: {}", client, receipt.transaction_hash);
        Ok(receipt)
    }

    /// Reactivate a previously deactivated policy client. Only callable by the registered owner.
    pub async fn activate_client(&self, client: Address) -> Result<TransactionReceipt, ChainIoError> {
        let signer = self.require_signer()?;
        let wallet = get_signer(signer, &self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, wallet);

        info!(
            "Activating policy client {} in registry {}",
            client, self.registry_address
        );

        let pending_tx = registry
            .activateClient(client)
            .send()
            .await
            .map_err(ChainIoError::ContractError)?;

        let receipt = pending_tx
            .get_receipt()
            .await
            .map_err(ChainIoError::AlloyProviderError)?;

        info!("Policy client {} activated, tx: {}", client, receipt.transaction_hash);
        Ok(receipt)
    }

    /// Transfer ownership of a registered client record. Only callable by the current owner.
    pub async fn set_client_owner(
        &self,
        client: Address,
        new_owner: Address,
    ) -> Result<TransactionReceipt, ChainIoError> {
        let signer = self.require_signer()?;
        let wallet = get_signer(signer, &self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, wallet);

        info!(
            "Transferring ownership of client {} to {} in registry {}",
            client, new_owner, self.registry_address
        );

        let pending_tx = registry
            .setClientOwner(client, new_owner)
            .send()
            .await
            .map_err(ChainIoError::ContractError)?;

        let receipt = pending_tx
            .get_receipt()
            .await
            .map_err(ChainIoError::AlloyProviderError)?;

        info!(
            "Ownership of client {} transferred to {}, tx: {}",
            client, new_owner, receipt.transaction_hash
        );
        Ok(receipt)
    }

    /// Get the registration record for a policy client
    pub async fn get_client_record(&self, client: Address) -> Result<ClientRecord, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, provider);

        info!(
            "Querying client record for {} from registry {}",
            client, self.registry_address
        );

        let record = registry
            .getClientRecord(client)
            .call()
            .await
            .map_err(ChainIoError::ContractError)?;

        info!(
            "Client {} record: owner={}, active={}, registeredAt={}",
            client, record.owner, record.active, record.registeredAt
        );
        Ok(record)
    }

    /// Get all policy client addresses owned by an address
    pub async fn get_clients_by_owner(&self, owner: Address) -> Result<Vec<Address>, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, provider);

        info!(
            "Querying clients owned by {} from registry {}",
            owner, self.registry_address
        );

        let clients = registry
            .getClientsByOwner(owner)
            .call()
            .await
            .map_err(ChainIoError::ContractError)?;

        info!("Owner {} has {} registered clients", owner, clients.len());
        Ok(clients)
    }

    /// Check if a client is registered and active
    pub async fn is_registered_client(&self, client: Address) -> Result<bool, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, provider);

        let is_registered = registry
            .isRegisteredClient(client)
            .call()
            .await
            .map_err(ChainIoError::ContractError)?;

        info!("Client {} registered and active: {}", client, is_registered);
        Ok(is_registered)
    }

    /// Get the number of clients owned by an address
    pub async fn get_client_count(&self, owner: Address) -> Result<u64, ChainIoError> {
        let provider = get_provider(&self.rpc_url);
        let registry = PolicyClientRegistry::new(self.registry_address, provider);

        let count = registry
            .getClientCount(owner)
            .call()
            .await
            .map_err(ChainIoError::ContractError)?;

        // U256 -> u64 is safe here since client count won't exceed u64::MAX
        let count_u64: u64 = count.to();
        info!("Owner {} has {} clients", owner, count_u64);
        Ok(count_u64)
    }
}