import json
import os
import socket
from typing import Any, Dict, List, Optional, Union
from flowsint_core.core.enricher_base import Enricher
from flowsint_enrichers.registry import flowsint_enricher
from flowsint_types.domain import Domain
from flowsint_types.asn import ASN
from flowsint_core.utils import is_valid_domain
from flowsint_core.core.logger import Logger
from tools.network.asnmap import AsnmapTool


@flowsint_enricher
class DomainToAsnEnricher(Enricher):
    """[ASNMAP] Takes a domain and returns its corresponding ASN."""

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

    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,
        )
        self.domain_asn_mapping: List[tuple[Domain, ASN]] = []

    @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": "PDCP_API_KEY",
                "type": "vaultSecret",
                "description": "The ProjectDiscovery Cloud Platform API key for asnmap.",
                "required": True,
            },
        ]

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

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

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

    async def scan(self, data: List[InputType]) -> List[OutputType]:
        results: List[OutputType] = []
        self.domain_asn_mapping = []
        asnmap = AsnmapTool()

        # Retrieve API key from vault or environment
        api_key = self.get_secret("PDCP_API_KEY", os.getenv("PDCP_API_KEY"))

        for domain in data:
            try:
                # Use asnmap tool to get ASN info from domain, passing the API key
                asn_data = asnmap.launch(domain.domain, type="domain", api_key=api_key)

                if asn_data and "as_number" in asn_data:
                    # Parse ASN number from string like "AS16276" to integer 16276
                    asn_string = asn_data["as_number"]
                    asn_number = int(asn_string.replace("AS", "").replace("as", ""))

                    # Create ASN object with correct field mapping
                    asn = ASN(
                        number=asn_number,
                        name=asn_data.get("as_name", ""),
                        country=asn_data.get("as_country", ""),
                        description=asn_data.get("as_name", ""),
                    )
                    results.append(asn)
                    self.domain_asn_mapping.append((domain, asn))

                    Logger.info(
                        self.sketch_id,
                        {
                            "message": f"[ASNMAP] Found AS{asn.number} ({asn.name}) for domain {domain.domain}"
                        },
                    )
                else:
                    Logger.warn(
                        self.sketch_id,
                        {
                            "message": f"[ASNMAP] No ASN data or missing 'as_number' field for domain {domain.domain}"
                        },
                    )

            except Exception as e:
                Logger.error(
                    self.sketch_id,
                    {"message": f"Error getting ASN for domain {domain.domain}: {e}"},
                )
                continue

        return results

    def postprocess(
        self, results: List[OutputType], input_data: List[InputType] = None
    ) -> List[OutputType]:
        # Create Neo4j relationships between domains and their corresponding ASNs
        if self._graph_service:
            for domain, asn in self.domain_asn_mapping:
                # Create domain node
                self.create_node(domain)
                # Create ASN node
                self.create_node(asn)

                # Create relationship
                self.create_relationship(domain, asn, "HOSTED_IN")

                self.log_graph_message(
                    f"Domain {domain.domain} is hosted in AS{asn.number} ({asn.name})"
                )

        return results


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