alloy-ens-extender 0.1.1

use alloy::{
    primitives::{Address, U256},
    providers::{Network, Provider},
    sol,
};
use alloy_ens::{EnsError, ProviderEnsExt};
use async_trait::async_trait;
use std::collections::HashMap;

mod constants;
mod utils;

use crate::constants::ID_URI_REGEX;
use crate::utils::{
    EnsAvatarNftInvalidMetadataError, NftUriNamespace, ParseAvatarError, ParseNftError,
    ParsedNftUri, UriItem, decode_base64_json, get_json_image, get_metadata_avatar_uri,
    parse_avatar_uri, parse_nft_uri,
};
use serde_json::Value;

sol! {
    /// ERC721 TokenURI contract for avatar resolution.
    #[sol(rpc)]
    contract ERC721TokenURI {
        function tokenURI(uint256 tokenId) external view returns (string memory);
    }
    /// ERC1155 TokenURI contract for avatar resolution.
    #[sol(rpc)]
    contract ERC1155TokenURI {
        function uri(uint256 id) external view returns (string memory);
    }
}

/// Simplified error type for ENS avatar resolution.
#[derive(Debug, thiserror::Error)]
pub enum AvatarError {
    /// ENS-related errors (name resolution, txt record lookup, etc.)
    #[error("ENS error: {0}")]
    Ens(#[from] EnsError),

    /// Contract call errors (ERC721/ERC1155 tokenURI calls)
    #[error("Failed to call contract method: {0}")]
    ContractCall(#[from] alloy::contract::Error),

    /// NFT URI parsing errors
    #[error("Failed to parse NFT URI: {0}")]
    NftUriParse(#[from] ParseNftError),

    /// Avatar URI parsing errors
    #[error("Failed to parse avatar URI: {0}")]
    AvatarUriParse(#[from] ParseAvatarError),

    /// Invalid metadata errors
    #[error("Invalid NFT metadata: {0}")]
    InvalidMetadata(#[from] EnsAvatarNftInvalidMetadataError),

    /// URI resolution failures
    #[error("Failed to resolve avatar URI: {uri}")]
    UriResolution { uri: String },

    /// JSON parsing errors
    #[error("Failed to parse JSON metadata: {0}")]
    JsonParse(String),
}

#[async_trait]
trait ProviderEnsExtend<N: alloy::providers::Network, P: alloy::providers::Provider<N>> {
    async fn get_erc721_token_uri(
        &self,
        address: Address,
        token_id: U256,
    ) -> Result<String, AvatarError>;

    async fn get_erc1155_token_uri(
        &self,
        address: Address,
        token_id: U256,
    ) -> Result<String, AvatarError>;

    async fn lookup_avatar(
        &self,
        name: &str,
        gateway_urls: Option<HashMap<String, String>>,
    ) -> Result<String, AvatarError>;

    async fn parse_nft_avatar_uri(
        &self,
        parsed_nft_uri: ParsedNftUri,
        gateway_urls: Option<HashMap<String, String>>,
    ) -> Result<UriItem, AvatarError>;
}

#[async_trait]
impl<N, P> ProviderEnsExtend<N, P> for P
where
    P: Provider<N>,
    N: Network,
{
    async fn get_erc721_token_uri(
        &self,
        address: Address,
        token_id: U256,
    ) -> Result<String, AvatarError> {
        let token = ERC721TokenURI::new(address, self);
        let uri = token.tokenURI(token_id).call().await?;
        Ok(uri)
    }

    async fn get_erc1155_token_uri(
        &self,
        address: Address,
        token_id: U256,
    ) -> Result<String, AvatarError> {
        let token = ERC1155TokenURI::new(address, self);
        let uri = token.uri(token_id).call().await?;
        Ok(uri)
    }

    async fn lookup_avatar(
        &self,
        name: &str,
        gateway_urls: Option<HashMap<String, String>>,
    ) -> Result<String, AvatarError> {
        // Get the avatar TXT record from ENS
        let record = self.lookup_txt(name, "avatar").await?;

        // Handle NFT avatars (eip155: format)
        if record.starts_with("eip155:") {
            let nft_decoded_uri = parse_nft_uri(&record)?;
            let token_id = nft_decoded_uri.token_id;

            let nft_uri_item = self
                .parse_nft_avatar_uri(nft_decoded_uri, gateway_urls.clone())
                .await?;

            // Handle on-chain metadata
            if nft_uri_item.is_on_chain
                && (nft_uri_item.uri.contains("data:application/json;base64,")
                    || nft_uri_item.uri.starts_with("{"))
            {
                let uri = if nft_uri_item.is_encoded {
                    decode_base64_json(&nft_uri_item.uri)
                        .map_err(|e| AvatarError::JsonParse(e.to_string()))?
                } else {
                    nft_uri_item.uri
                };

                let uri_image = get_json_image(&Value::String(uri))
                    .map_err(|e| AvatarError::JsonParse(e.to_string()))?;

                return Ok(parse_avatar_uri(&uri_image, gateway_urls)?.uri);
            }

            // Replace token ID placeholders and fetch metadata
            let replace_all = ID_URI_REGEX.replace_all(&nft_uri_item.uri, token_id.to_string());
            let final_uri = replace_all.into_owned();
            println!("FINAL_URI {}", final_uri);

            return Ok(get_metadata_avatar_uri(&final_uri, gateway_urls).await?);
        }

        // Handle regular avatar URIs (IPFS, HTTP, etc.)
        let uri_item = parse_avatar_uri(&record, gateway_urls)?;
        Ok(uri_item.uri)
    }

    async fn parse_nft_avatar_uri(
        &self,
        parsed_nft_uri: ParsedNftUri,
        gateway_urls: Option<HashMap<String, String>>,
    ) -> Result<UriItem, AvatarError> {
        // Get the token URI based on the NFT standard
        let nft_uri = match parsed_nft_uri.namespace {
            NftUriNamespace::ERC1155 => {
                self.get_erc1155_token_uri(parsed_nft_uri.contract_address, parsed_nft_uri.token_id)
                    .await?
            }
            NftUriNamespace::ERC721 => {
                self.get_erc721_token_uri(parsed_nft_uri.contract_address, parsed_nft_uri.token_id)
                    .await?
            }
        };

        let avatar_uri = parse_avatar_uri(&nft_uri, gateway_urls)?;
        Ok(avatar_uri)
    }
}

pub fn normalize(name: &str) -> String {
    let normalizer = ens_normalize_rs::EnsNameNormalizer::default();
    let processed = normalizer.process(name).unwrap();
    processed.normalize()
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy::providers::ProviderBuilder;

    #[test]
    fn test_normalize() {
        assert_eq!(normalize("‐Ξ1️⃣"), "-ξ1⃣");
    }
    #[tokio::test]
    async fn test_pub_resolver_fetching_avatar() {
        let provider = ProviderBuilder::new()
            .connect_http("http://reth-ethereum.ithaca.xyz/rpc".parse().unwrap());

        let name = "nick.eth";
        let res = provider.lookup_avatar(name, None).await.unwrap();
        assert_eq!(
            res,
            "https://i.seadn.io/gcs/files/3ae7be6c41ad4767bf3ecbc0493b4bfb.png?w=4000&auto=format"
        );
    }
}