#!/usr/bin/env python3
"""
Resolve ENS Names and Contact Information

This script fetches ENS names and additional contact information for Ethereum addresses
in the database. It uses the Web3 library to query the Ethereum blockchain for ENS records
and text records containing social profiles and contact information.

Usage:
    python resolve_ens_names.py
"""

import os
import sys
import logging
import time
from typing import List, Dict, Any, Optional, Tuple
from web3 import Web3
from dotenv import load_dotenv

# Add parent directory to path to import utils
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.db_connector import DatabaseConnector
from utils.logger import setup_logger

# Load environment variables
load_dotenv()

# Setup logging
logger = setup_logger("ens_resolver")

class ENSResolver:
    """Resolver for ENS names and contact information from Ethereum addresses"""
    
    # ENS text record keys to check
    TEXT_RECORDS = [
        "name",        # Display name
        "email",       # Email address
        "url",         # Website URL
        "avatar",      # Avatar URL
        "description", # Bio/description
        "notice",      # Notice
        "keywords",    # Keywords/tags
        "com.twitter", # Twitter handle
        "com.github",  # GitHub username
        "org.telegram", # Telegram username
        "com.discord", # Discord username
        "com.reddit",  # Reddit username
        "xyz.farcaster", # Farcaster handle
        "social.picture", # Profile picture
        "vnd.twitter", # Alternative Twitter
        "vnd.github", # Alternative GitHub
    ]
    
    def __init__(self):
        """Initialize the resolver"""
        # Initialize database
        self.db = DatabaseConnector()
        
        # Initialize Web3 connection
        infura_key = os.getenv("INFURA_API_KEY")
        if not infura_key:
            raise ValueError("INFURA_API_KEY environment variable is required")
        
        self.w3 = Web3(Web3.HTTPProvider(f"https://mainnet.infura.io/v3/{infura_key}"))
        if not self.w3.is_connected():
            raise ConnectionError("Failed to connect to Ethereum node")
        
        logger.info(f"Connected to Ethereum node: {self.w3.client_version}")
    
    def get_contacts_without_ens(self) -> List[Dict[str, Any]]:
        """Get all contacts that have an Ethereum address but no ENS name"""
        query = """
        SELECT id, "ethereumAddress", name
        FROM "Contact"
        WHERE "ethereumAddress" IS NOT NULL
        AND "ensName" IS NULL
        """
        
        result = self.db.execute_query(query)
        logger.info(f"Found {len(result)} contacts without ENS names")
        return result
    
    def get_all_contacts_with_eth_address(self) -> List[Dict[str, Any]]:
        """Get all contacts that have an Ethereum address to check for additional info"""
        query = """
        SELECT id, "ethereumAddress", "ensName", name, twitter, discord, telegram, email, farcaster
        FROM "Contact"
        WHERE "ethereumAddress" IS NOT NULL
        """
        
        result = self.db.execute_query(query)
        logger.info(f"Found {len(result)} contacts with Ethereum addresses")
        return result
    
    def resolve_ens_name(self, address: str) -> Optional[str]:
        """Resolve ENS name for an Ethereum address"""
        try:
            # Ensure the address is properly formatted
            checksum_address = self.w3.to_checksum_address(address)
            
            # Try to get the ENS name
            ens_name = self.w3.ens.name(checksum_address)
            
            # If we got a name, verify it resolves back to the same address
            if ens_name:
                resolved_address = self.w3.ens.address(ens_name)
                if resolved_address and resolved_address.lower() == address.lower():
                    logger.info(f"Resolved ENS name for {address}: {ens_name}")
                    return ens_name
                else:
                    logger.warning(f"ENS name {ens_name} for {address} resolves to different address {resolved_address}")
            
            return None
        except Exception as e:
            logger.error(f"Error resolving ENS name for {address}: {e}")
            return None
    
    def get_ens_text_records(self, ens_name: str) -> Dict[str, str]:
        """Get text records for an ENS name"""
        text_records = {}
        
        try:
            for key in self.TEXT_RECORDS:
                try:
                    value = self.w3.ens.get_text(ens_name, key)
                    if value:
                        text_records[key] = value
                except Exception as e:
                    logger.debug(f"Error getting text record '{key}' for {ens_name}: {e}")
            
            if text_records:
                logger.info(f"Found {len(text_records)} text records for {ens_name}: {', '.join(text_records.keys())}")
            
            return text_records
        except Exception as e:
            logger.error(f"Error getting text records for {ens_name}: {e}")
            return {}
    
    def map_text_records_to_contact_fields(self, text_records: Dict[str, str]) -> Dict[str, str]:
        """Map ENS text records to Contact model fields"""
        contact_fields = {}
        
        # Map known fields
        if "name" in text_records:
            contact_fields["name"] = text_records["name"]
        
        if "email" in text_records:
            contact_fields["email"] = text_records["email"]
        
        # Twitter can be in different text records
        for twitter_key in ["com.twitter", "vnd.twitter"]:
            if twitter_key in text_records:
                twitter = text_records[twitter_key]
                # Remove @ if present
                if twitter.startswith("@"):
                    twitter = twitter[1:]
                contact_fields["twitter"] = twitter
                break
        
        # Discord
        if "com.discord" in text_records:
            contact_fields["discord"] = text_records["com.discord"]
        
        # Telegram
        if "org.telegram" in text_records:
            contact_fields["telegram"] = text_records["org.telegram"]
        
        # Farcaster
        if "xyz.farcaster" in text_records:
            contact_fields["farcaster"] = text_records["xyz.farcaster"]
        
        # Collect other social profiles
        other_social = []
        
        if "com.github" in text_records or "vnd.github" in text_records:
            github = text_records.get("com.github") or text_records.get("vnd.github")
            other_social.append(f"GitHub: {github}")
        
        if "com.reddit" in text_records:
            other_social.append(f"Reddit: {text_records['com.reddit']}")
        
        if "url" in text_records:
            other_social.append(f"Website: {text_records['url']}")
        
        if other_social:
            contact_fields["otherSocial"] = "; ".join(other_social)
        
        return contact_fields
    
    def update_contact_info(self, contact_id: str, ens_name: Optional[str] = None, contact_info: Optional[Dict[str, str]] = None) -> bool:
        """Update a contact with ENS name and additional contact information"""
        try:
            # Build the update query dynamically based on what fields we have
            update_fields = []
            params = {"contact_id": contact_id}
            
            if ens_name:
                update_fields.append('"ensName" = %(ens_name)s')
                params["ens_name"] = ens_name
            
            if contact_info:
                for field, value in contact_info.items():
                    update_fields.append(f'"{field}" = %({field})s')
                    params[field] = value
            
            if not update_fields:
                logger.warning(f"No fields to update for contact {contact_id}")
                return False
            
            query = f"""
            UPDATE "Contact"
            SET {", ".join(update_fields)},
                "updatedAt" = NOW()
            WHERE id = %(contact_id)s
            """
            
            self.db.execute_update(query, params)
            
            # Also update the name if it's currently a generic name and we have a better name
            if ens_name and "name" not in contact_info:
                name_query = """
                SELECT name FROM "Contact" WHERE id = %(contact_id)s
                """
                result = self.db.execute_query(name_query, {"contact_id": contact_id})
                current_name = result[0]["name"] if result else None
                
                # If the current name is generic (starts with MC_ or ETH_ or RG_), update it
                if current_name and (current_name.startswith("MC_") or current_name.startswith("ETH_") or current_name.startswith("RG_")):
                    # Use ENS name without .eth suffix as the name
                    name = ens_name[:-4] if ens_name.endswith('.eth') else ens_name
                    
                    update_name_query = """
                    UPDATE "Contact"
                    SET name = %(name)s,
                        "updatedAt" = NOW()
                    WHERE id = %(contact_id)s
                    """
                    
                    self.db.execute_update(update_name_query, {
                        "contact_id": contact_id,
                        "name": name
                    })
                    
                    logger.info(f"Updated contact {contact_id} name from '{current_name}' to '{name}'")
            
            fields_updated = []
            if ens_name:
                fields_updated.append("ENS name")
            if contact_info:
                fields_updated.extend(list(contact_info.keys()))
            
            logger.info(f"Updated contact {contact_id} with: {', '.join(fields_updated)}")
            return True
        except Exception as e:
            logger.error(f"Error updating contact {contact_id}: {e}")
            return False
    
    def process_contact(self, contact: Dict[str, Any]) -> Tuple[bool, bool]:
        """Process a single contact to resolve ENS name and contact info"""
        contact_id = contact["id"]
        address = contact["ethereumAddress"]
        current_ens = contact.get("ensName")
        
        ens_updated = False
        info_updated = False
        
        # Skip if no address
        if not address:
            return ens_updated, info_updated
        
        # Resolve ENS name if not already set
        ens_name = None
        if not current_ens:
            ens_name = self.resolve_ens_name(address)
            if ens_name:
                ens_updated = True
        else:
            ens_name = current_ens
        
        # Get contact info from ENS text records if we have an ENS name
        contact_info = {}
        if ens_name:
            text_records = self.get_ens_text_records(ens_name)
            if text_records:
                contact_info = self.map_text_records_to_contact_fields(text_records)
                
                # Only include fields that are different from what we already have
                for field in list(contact_info.keys()):
                    if field in contact and contact[field] == contact_info[field]:
                        del contact_info[field]
                
                if contact_info:
                    info_updated = True
        
        # Update the contact if we have new information
        if ens_updated or info_updated:
            self.update_contact_info(contact_id, ens_name if ens_updated else None, contact_info if info_updated else None)
        
        return ens_updated, info_updated
    
    def run(self, batch_size: int = 50, delay_seconds: float = 0.5, resolve_all: bool = False):
        """Run the resolver for contacts"""
        if resolve_all:
            contacts = self.get_all_contacts_with_eth_address()
        else:
            contacts = self.get_contacts_without_ens()
        
        if not contacts:
            logger.info("No contacts found to process")
            return 0, 0
        
        ens_updated_count = 0
        info_updated_count = 0
        
        # Process in batches to avoid rate limiting
        for i in range(0, len(contacts), batch_size):
            batch = contacts[i:i+batch_size]
            logger.info(f"Processing batch {i//batch_size + 1}/{(len(contacts) + batch_size - 1)//batch_size}")
            
            for contact in batch:
                ens_updated, info_updated = self.process_contact(contact)
                
                if ens_updated:
                    ens_updated_count += 1
                if info_updated:
                    info_updated_count += 1
                
                # Add a small delay to avoid rate limiting
                time.sleep(delay_seconds)
        
        logger.info(f"Updated ENS names for {ens_updated_count} contacts and contact info for {info_updated_count} contacts out of {len(contacts)} processed")
        return ens_updated_count, info_updated_count

def main():
    """Main function"""
    try:
        import argparse
        
        parser = argparse.ArgumentParser(description="Resolve ENS names and contact information")
        parser.add_argument("--all", action="store_true", help="Process all contacts with Ethereum addresses, not just those without ENS names")
        parser.add_argument("--batch-size", type=int, default=50, help="Number of contacts to process in each batch")
        parser.add_argument("--delay", type=float, default=0.5, help="Delay in seconds between processing contacts")
        
        args = parser.parse_args()
        
        resolver = ENSResolver()
        ens_count, info_count = resolver.run(
            batch_size=args.batch_size,
            delay_seconds=args.delay,
            resolve_all=args.all
        )
        
        logger.info(f"ENS resolution completed successfully. Updated {ens_count} ENS names and {info_count} contact info records.")
        return 0
    except Exception as e:
        logger.exception(f"Error running ENS resolver: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())