# dnsrecon/providers/correlation_provider.py

import re
from typing import Dict, Any, List

from .base_provider import BaseProvider
from core.provider_result import ProviderResult
from core.graph_manager import NodeType, GraphManager

class CorrelationProvider(BaseProvider):
    """
    A provider that finds correlations between nodes in the graph.
    """

    def __init__(self, name: str = "correlation", session_config=None):
        """
        Initialize the correlation provider.
        """
        super().__init__(name, session_config=session_config)
        self.graph: GraphManager | None = None
        self.correlation_index = {}
        self.date_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}')
        self.EXCLUDED_KEYS = [
            'cert_source',
            'cert_issuer_ca_id',
            'cert_common_name',
            'cert_validity_period_days',
            'cert_issuer_name',
            'cert_entry_timestamp',
            'cert_serial_number', # useless
            'cert_not_before',
            'cert_not_after',
            'dns_ttl',
            'timestamp',
            'last_update',
            'updated_timestamp',
            'discovery_timestamp',
            'query_timestamp',
        ]

    def get_name(self) -> str:
        """Return the provider name."""
        return "correlation"

    def get_display_name(self) -> str:
        """Return the provider display name for the UI."""
        return "Correlation Engine"

    def requires_api_key(self) -> bool:
        """Return True if the provider requires an API key."""
        return False

    def get_eligibility(self) -> Dict[str, bool]:
        """Return a dictionary indicating if the provider can query domains and/or IPs."""
        return {'domains': True, 'ips': True}

    def is_available(self) -> bool:
        """Check if the provider is available and properly configured."""
        return True

    def query_domain(self, domain: str) -> ProviderResult:
        """
        Query the provider for information about a domain.
        """
        return self._find_correlations(domain)

    def query_ip(self, ip: str) -> ProviderResult:
        """
        Query the provider for information about an IP address.
        """
        return self._find_correlations(ip)

    def set_graph_manager(self, graph_manager: GraphManager):
        """
        Set the graph manager for the provider to use.
        """
        self.graph = graph_manager

    def _find_correlations(self, node_id: str) -> ProviderResult:
        """
        Find correlations for a given node.
        """
        result = ProviderResult()
        # FIXED: Ensure self.graph is not None before proceeding.
        if not self.graph or not self.graph.graph.has_node(node_id):
            return result

        node_attributes = self.graph.graph.nodes[node_id].get('attributes', [])

        for attr in node_attributes:
            attr_name = attr.get('name')
            attr_value = attr.get('value')
            attr_provider = attr.get('provider', 'unknown')

            should_exclude = (
                any(excluded_key in attr_name or attr_name == excluded_key for excluded_key in self.EXCLUDED_KEYS) or
                not isinstance(attr_value, (str, int, float, bool)) or
                attr_value is None or
                isinstance(attr_value, bool) or
                (isinstance(attr_value, str) and (
                    len(attr_value) < 4 or
                    self.date_pattern.match(attr_value) or
                    attr_value.lower() in ['unknown', 'none', 'null', 'n/a', 'true', 'false', '0', '1']
                )) or
                (isinstance(attr_value, (int, float)) and (
                    attr_value == 0 or
                    attr_value == 1 or
                    abs(attr_value) > 1000000
                ))
            )

            if should_exclude:
                continue

            if attr_value not in self.correlation_index:
                self.correlation_index[attr_value] = {
                    'nodes': set(),
                    'sources': []
                }

            self.correlation_index[attr_value]['nodes'].add(node_id)

            source_info = {
                'node_id': node_id,
                'provider': attr_provider,
                'attribute': attr_name,
                'path': f"{attr_provider}_{attr_name}"
            }

            existing_sources = [s for s in self.correlation_index[attr_value]['sources']
                              if s['node_id'] == node_id and s['path'] == source_info['path']]
            if not existing_sources:
                self.correlation_index[attr_value]['sources'].append(source_info)

            if len(self.correlation_index[attr_value]['nodes']) > 1:
                self._create_correlation_relationships(attr_value, self.correlation_index[attr_value], result)
        return result

    def _create_correlation_relationships(self, value: Any, correlation_data: Dict[str, Any], result: ProviderResult):
        """
        Create correlation relationships and add them to the provider result.
        """
        correlation_node_id = f"corr_{hash(str(value)) & 0x7FFFFFFF}"
        nodes = correlation_data['nodes']
        sources = correlation_data['sources']

        # Add the correlation node as an attribute to the result
        result.add_attribute(
            target_node=correlation_node_id,
            name="correlation_value",
            value=value,
            attr_type=str(type(value)),
            provider=self.name,
            confidence=0.9,
            metadata={
                'correlated_nodes': list(nodes),
                'sources': sources,
            }
        )

        for source in sources:
            node_id = source['node_id']
            provider = source['provider']
            attribute = source['attribute']
            relationship_label = f"corr_{provider}_{attribute}"

            # Add the relationship to the result
            result.add_relationship(
                source_node=node_id,
                target_node=correlation_node_id,
                relationship_type=relationship_label,
                provider=self.name,
                confidence=0.9,
                raw_data={
                    'correlation_value': value,
                    'original_attribute': attribute,
                    'correlation_type': 'attribute_matching'
                }
            )