import os
import re
from typing import Any, Dict, List, Optional, Set

from dotenv import load_dotenv
from flowsint_core.core.enricher_base import Enricher
from flowsint_core.core.logger import Logger
from flowsint_types.address import Location
from flowsint_types.domain import Domain
from flowsint_types.email import Email
from flowsint_types.individual import Individual
from flowsint_types.phone import Phone

from flowsint_enrichers.registry import flowsint_enricher
from tools.network.whoxy import WhoxyTool

load_dotenv()

WHOXY_API_KEY = os.getenv("WHOXY_API_KEY")


@flowsint_enricher
class EmailToDomainsEnricher(Enricher):
    """[WHOXY] Takes an email and returns the domains it registered."""

    # Define types as class attributes - base class handles schema generation automatically
    InputType = Email
    OutputType = Domain

    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 get_params_schema(cls) -> List[Dict[str, Any]]:
        """Declare required parameters for this enricher"""
        return [
            {
                "name": "WHOXY_API_KEY",
                "type": "vaultSecret",
                "description": "The Whoxy API key to use for domain lookups.",
                "required": True,
            },
        ]

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

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

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

    async def scan(self, data: List[InputType]) -> List[OutputType]:
        """Find domains related to emails using whoxy api."""
        domains: List[OutputType] = []
        self._extracted_data = []  # Store all extracted data for postprocess
        api_key = self.get_secret("WHOXY_API_KEY", os.getenv("WHOXY_API_KEY"))

        for email in data:
            infos_data = self.__get_infos_from_whoxy(email.email, api_key)
            if infos_data and "search_result" in infos_data:
                # Process each domain result
                for result in infos_data["search_result"]:
                    if self.__is_valid_domain_result(result):
                        domain_name = result.get("domain_name")
                        if domain_name:
                            domain = Domain(domain=domain_name, root=True)
                            domains.append(domain)

                            # Store extracted data for postprocess
                            extracted_info = {
                                "email": email,
                                "domain": domain,
                                "domain_data": result,
                                "contacts": {
                                    "registrant": result.get("registrant_contact", {}),
                                    "administrative": result.get(
                                        "administrative_contact", {}
                                    ),
                                    "technical": result.get("technical_contact", {}),
                                    "billing": result.get("billing_contact", {}),
                                },
                            }
                            self._extracted_data.append(extracted_info)
            else:
                Logger.info(
                    self.sketch_id,
                    {"message": f"[WHOXY] No domain found for email {email.email}."},
                )
        return domains

    def __get_infos_from_whoxy(self, email: str, api_key: str) -> Dict[str, Any]:
        infos: Dict[str, Any] = {}
        whoxy = WhoxyTool()
        try:
            params = {
                "key": api_key,
                "reverse": "whois",
                "email": email,
            }
            infos = whoxy.launch(params=params)
        except Exception as e:
            Logger.error(
                self.sketch_id, {"message": f"[WHOXY] Whoxy exception for {email}: {e}"}
            )
        return infos

    def __is_valid_domain_result(self, result: Dict[str, Any]) -> bool:
        """Check if a domain result is valid."""
        domain_name = result.get("domain_name")
        if not domain_name:
            return False
        return True

    def __extract_individual_from_contact(
        self, contact: Dict[str, Any], contact_type: str
    ) -> Individual:
        """Extract individual information from contact data."""
        full_name = contact.get("full_name", "")

        # Skip if name is redacted
        if "REDACTED" in full_name or "REDACTED FOR PRIVACY" in full_name:
            return None

        # Parse full name into first and last name
        name_parts = full_name.strip().split()
        first_name = name_parts[0] if name_parts else "N/A"
        last_name = " ".join(name_parts[1:]) if len(name_parts) > 1 else "N/A"

        # Extract email and phone
        email = contact.get("email_address", "")
        phone = contact.get("phone_number", "")

        # Skip if email is redacted or invalid
        if (
            "REDACTED" in email
            or "REDACTED FOR PRIVACY" in email
            or not self.__is_valid_email(email)
        ):
            email = ""

        # Skip if phone is redacted
        if "REDACTED" in phone or "REDACTED FOR PRIVACY" in phone:
            phone = ""

        # Extract address information
        address = contact.get("mailing_address", "")
        city = contact.get("city_name", "")
        zip_code = contact.get("zip_code", "")
        country = contact.get("country_name", "")

        # Skip if address is redacted
        if "REDACTED" in address or "REDACTED FOR PRIVACY" in address:
            address = ""
        if "REDACTED" in city or "REDACTED FOR PRIVACY" in city:
            city = ""
        if "REDACTED" in zip_code or "REDACTED FOR PRIVACY" in zip_code:
            zip_code = ""
        if "REDACTED" in country or "REDACTED FOR PRIVACY" in country:
            country = ""

        # Create individual object
        individual = Individual(
            first_name=first_name,
            last_name=last_name,
            full_name=full_name,
            email_addresses=[email] if email else None,
            phone_numbers=[phone] if phone else None,
        )

        return individual

    def __is_valid_email(self, email: str) -> bool:
        """Check if email is valid."""
        if not email:
            return False
        # Basic email validation
        pattern = r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$"
        return bool(re.match(pattern, email))

    def __extract_physical_address(self, contact: Dict[str, Any]) -> Location:
        """Extract physical address from contact data."""
        address = contact.get("mailing_address", "")
        city = contact.get("city_name", "")
        zip_code = contact.get("zip_code", "")
        country = contact.get("country_name", "")

        # Skip if any part is redacted
        if any("REDACTED" in field for field in [address, city, zip_code, country]):
            return None

        if not all([address, city, zip_code, country]):
            return None

        return Location(address=address, city=city, zip=zip_code, country=country)

    def postprocess(
        self, results: List[OutputType], original_input: List[InputType]
    ) -> List[OutputType]:
        """Create Neo4j nodes and relationships from extracted data."""
        if not self._graph_service:
            return results

        # Track processed entities to avoid duplicates
        processed_domains: Set[str] = set()
        processed_individuals: Set[str] = set()
        processed_emails: Set[str] = set()
        processed_phones: Set[str] = set()
        processed_addresses: Set[str] = set()

        for extracted_info in self._extracted_data:
            email = extracted_info["email"]
            domain = extracted_info["domain"]
            domain_data = extracted_info["domain_data"]
            contacts = extracted_info["contacts"]

            domain_name = domain.domain
            if domain_name in processed_domains:
                continue
            processed_domains.add(domain_name)

            # Create email node
            self.create_node(email)

            # Create domain node
            domain_obj = Domain(domain=domain_name)
            self.create_node(domain_obj)

            # Create relationship between email and domain
            domain_obj_email = Domain(domain=domain_name)
            self.create_relationship(email, domain_obj_email, "HAS_REGISTERED_DOMAIN")

            # Process all contact types
            for contact_type, contact in contacts.items():
                if contact:
                    self.__process_contact(
                        contact,
                        contact_type.upper(),
                        domain_name,
                        email.email,
                        processed_individuals,
                        processed_emails,
                        processed_phones,
                        processed_addresses,
                    )

            self.log_graph_message(
                f"Processed domain {domain_name} for email {email.email}"
            )

        return results

    def __process_contact(
        self,
        contact: Dict[str, Any],
        contact_type: str,
        domain_name: str,
        email_address: str,
        processed_individuals: Set[str],
        processed_emails: Set[str],
        processed_phones: Set[str],
        processed_addresses: Set[str],
    ):
        """Process a contact and create all related entities and relationships."""

        # Extract individual
        individual = self.__extract_individual_from_contact(contact, contact_type)
        if not individual:
            return

        individual_id = (
            f"{individual.first_name}_{individual.last_name}_{individual.full_name}"
        )
        if individual_id in processed_individuals:
            return

        processed_individuals.add(individual_id)

        # Create individual node
        self.create_node(individual)

        # Create relationship between individual and domain
        domain_obj_ind = Domain(domain=domain_name)
        self.create_relationship(
            individual, domain_obj_ind, f"IS_{contact_type}_CONTACT"
        )

        # Create relationship between individual and email
        email_obj_ind = Email(email=email_address)
        self.create_relationship(individual, email_obj_ind, "WORKS_FOR")

        # Process email addresses
        if individual.email_addresses:
            for email_obj in individual.email_addresses:
                email_str = email_obj.email
                if email_str and email_str not in processed_emails:
                    processed_emails.add(email_str)

                    # Create email node
                    email_node = Email(email=email_str)
                    self.create_node(email_node)

                    # Create relationship between individual and email
                    self.create_relationship(individual, email_node, "HAS_EMAIL")

        # Process phone numbers
        if individual.phone_numbers:
            for phone_obj in individual.phone_numbers:
                phone_str = phone_obj.number
                if phone_str and phone_str not in processed_phones:
                    processed_phones.add(phone_str)

                    # Create phone node
                    phone_node = Phone(number=phone_str)
                    self.create_node(phone_node)

                    # Create relationship between individual and phone
                    self.create_relationship(individual, phone_node, "HAS_PHONE")

        # Process physical address
        address = self.__extract_physical_address(contact)
        if address:
            address_id = (
                f"{address.address}_{address.city}_{address.zip}_{address.country}"
            )
            if address_id not in processed_addresses:
                processed_addresses.add(address_id)

                # Create address node
                self.create_node(address)

                # Create relationship between individual and address
                self.create_relationship(individual, address, "LIVES_AT")


InputType = EmailToDomainsEnricher.InputType
OutputType = EmailToDomainsEnricher.OutputType
