import os
from typing import List, Dict, Any, Optional, Union
import requests
from flowsint_core.core.enricher_base import Enricher
from flowsint_enrichers.registry import flowsint_enricher
from flowsint_types.wallet import CryptoWallet, CryptoNFT
from flowsint_core.core.logger import Logger
from dotenv import load_dotenv

load_dotenv()


@flowsint_enricher
class CryptoWalletAddressToNFTs(Enricher):
    """[ETHERSCAN] Resolve NFTs for a wallet address (ETH)."""

    # Define types as class attributes - base class handles schema generation automatically
    InputType = CryptoWallet
    OutputType = CryptoNFT

    def __init__(
        self,
        sketch_id: Optional[str] = None,
        scan_id: Optional[str] = None,
        vault=None,
        params: Optional[Dict[str, Any]] = None,
    ):
        super().__init__(
            sketch_id=sketch_id,
            scan_id=scan_id,
            params_schema=self.get_params_schema(),
            vault=vault,
            params=params,
        )

    @classmethod
    def required_params(cls) -> bool:
        return True

    @classmethod
    def icon(cls) -> str | None:
        return "cryptowallet"

    @classmethod
    def get_params_schema(cls) -> List[Dict[str, Any]]:
        """Declare required parameters for this enricher"""
        return [
            {
                "name": "ETHERSCAN_API_KEY",
                "type": "vaultSecret",
                "description": "The Etherscan API key to use for the NFT lookup.",
                "required": True,
            },
            {
                "name": "ETHERSCAN_API_URL",
                "type": "url",
                "description": "The Etherscan API URL to use for the NFT lookup.",
                "required": False,
                "default": "https://api.etherscan.io/v2/api",
            },
        ]

    @classmethod
    def name(cls) -> str:
        return "cryptowallet_to_nfts"

    @classmethod
    def category(cls) -> str:
        return "CryptoWallet"

    @classmethod
    def key(cls) -> str:
        return "address"

    async def scan(self, data: List[InputType]) -> List[OutputType]:
        results: List[OutputType] = []
        api_key = self.get_secret("ETHERSCAN_API_KEY", os.getenv("ETHERSCAN_API_KEY"))
        api_url = self.get_params().get(
            "ETHERSCAN_API_URL", "https://api.etherscan.io/v2/api"
        )
        for d in data:
            try:
                nfts = self._get_nfts(d.address, api_key, api_url)
                results.append(nfts)
            except Exception as e:
                print(f"Error resolving nfts for {d.address}: {e}")
        return results

    def _get_nfts(self, address: str, api_key: str, api_url: str) -> List[CryptoNFT]:
        nfts = []
        """Get nfts for a wallet address."""
        params = {
            "module": "account",
            "action": "tokennfttx",
            "address": address,
            "startblock": 0,
            "endblock": 99999999,
            "page": 1,
            "offset": 10000,
            "sort": "asc",
            "apikey": api_key,
        }
        response = requests.get(api_url, params=params)
        data = response.json()
        results = data["result"]
        for tx in results:
            nfts.append(
                CryptoNFT(
                    wallet=CryptoWallet(address=address),
                    contract_address=tx["contractAddress"],
                    token_id=tx["tokenID"],
                    collection_name=tx["collectionName"],
                    metadata_url=tx["metadataURL"],
                    image_url=tx["imageURL"],
                    name=tx["name"],
                )
            )
        return nfts

    def postprocess(self, results: List[OutputType], original_input: List[InputType]) -> List[OutputType]:
        if not self._graph_service:
            return results

        for nfts in results:
            for nft in nfts:
                # Create or update wallet node
                self.create_node(nft.wallet)
                # Create or update NFT node
                nft_key = f"{nft.contract_address}_{nft.token_id}"
                self.create_node(nft)
                # Create relationship from wallet to NFT
                wallet_obj = CryptoWallet(address=nft.wallet.address)
                nft_obj = CryptoNFT(
                    wallet=wallet_obj,
                    contract_address=nft.contract_address,
                    token_id=nft.token_id,
                    collection_name=nft.collection_name,
                    metadata_url=nft.metadata_url,
                    image_url=nft.image_url,
                    name=nft.name,
                )
                self.create_relationship(wallet_obj, nft_obj, "OWNS")
                self.log_graph_message(
                    f"Found NFT for {nft.wallet.address}: {nft.contract_address} - {nft.token_id}"
                )

        return results


# Make types available at module level for easy access
InputType = CryptoWalletAddressToNFTs.InputType
OutputType = CryptoWalletAddressToNFTs.OutputType
