"""
Graph data model for DNSRecon using NetworkX.
Manages in-memory graph storage with confidence scoring and forensic metadata.
"""

import json
import threading
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple, Set
from enum import Enum
from datetime import timezone

import networkx as nx


class NodeType(Enum):
    """Enumeration of supported node types."""
    DOMAIN = "domain"
    IP = "ip"
    CERTIFICATE = "certificate"
    ASN = "asn"
    LARGE_ENTITY = "large_entity"


class RelationshipType(Enum):
    """Enumeration of supported relationship types with confidence scores."""
    SAN_CERTIFICATE = ("san", 0.9)           # Certificate SAN relationships
    A_RECORD = ("a_record", 0.8)             # A/AAAA record relationships
    CNAME_RECORD = ("cname", 0.8)            # CNAME relationships
    PASSIVE_DNS = ("passive_dns", 0.6)       # Passive DNS relationships
    ASN_MEMBERSHIP = ("asn", 0.7)            # ASN relationships
    MX_RECORD = ("mx_record", 0.7)           # MX record relationships
    NS_RECORD = ("ns_record", 0.7)           # NS record relationships
    
    def __init__(self, relationship_name: str, default_confidence: float):
        self.relationship_name = relationship_name
        self.default_confidence = default_confidence


class GraphManager:
    """
    Thread-safe graph manager for DNSRecon infrastructure mapping.
    Uses NetworkX for in-memory graph storage with confidence scoring.
    """
    
    def __init__(self):
        """Initialize empty directed graph."""
        self.graph = nx.DiGraph()
        # self.lock = threading.Lock()
        self.creation_time = datetime.now(timezone.utc).isoformat()
        self.last_modified = self.creation_time
        
    def add_node(self, node_id: str, node_type: NodeType, 
                 metadata: Optional[Dict[str, Any]] = None) -> bool:
        """
        Add a node to the graph.
        
        Args:
            node_id: Unique identifier for the node
            node_type: Type of the node (Domain, IP, Certificate, ASN)
            metadata: Additional metadata for the node
            
        Returns:
            bool: True if node was added, False if it already exists
        """
        if self.graph.has_node(node_id):
            # Update metadata if node exists
            existing_metadata = self.graph.nodes[node_id].get('metadata', {})
            if metadata:
                existing_metadata.update(metadata)
                self.graph.nodes[node_id]['metadata'] = existing_metadata
            return False
        
        node_attributes = {
            'type': node_type.value,
            'added_timestamp': datetime.now(timezone.utc).isoformat(),
            'metadata': metadata or {}
        }
        
        self.graph.add_node(node_id, **node_attributes)
        self.last_modified = datetime.now(timezone.utc).isoformat()
        return True
    
    def add_edge(self, source_id: str, target_id: str, 
                 relationship_type: RelationshipType,
                 confidence_score: Optional[float] = None,
                 source_provider: str = "unknown",
                 raw_data: Optional[Dict[str, Any]] = None) -> bool:
        """
        Add an edge between two nodes.
        
        Args:
            source_id: Source node identifier
            target_id: Target node identifier  
            relationship_type: Type of relationship
            confidence_score: Custom confidence score (overrides default)
            source_provider: Provider that discovered this relationship
            raw_data: Raw data from provider response
            
        Returns:
            bool: True if edge was added, False if it already exists
        """
        #with self.lock:
        # Ensure both nodes exist
        if not self.graph.has_node(source_id) or not self.graph.has_node(target_id):
            # If the target node is a subdomain, it should be added.
            # The scanner will handle this logic.
            pass

        # Check if edge already exists
        if self.graph.has_edge(source_id, target_id):
            # Update confidence score if new score is higher
            existing_confidence = self.graph.edges[source_id, target_id]['confidence_score']
            new_confidence = confidence_score or relationship_type.default_confidence
            
            if new_confidence > existing_confidence:
                self.graph.edges[source_id, target_id]['confidence_score'] = new_confidence
                self.graph.edges[source_id, target_id]['updated_timestamp'] = datetime.now(timezone.utc).isoformat()
                self.graph.edges[source_id, target_id]['updated_by'] = source_provider
            
            return False
        
        edge_attributes = {
            'relationship_type': relationship_type.relationship_name,
            'confidence_score': confidence_score or relationship_type.default_confidence,
            'source_provider': source_provider,
            'discovery_timestamp': datetime.now(timezone.utc).isoformat(),
            'raw_data': raw_data or {}
        }
        
        self.graph.add_edge(source_id, target_id, **edge_attributes)
        self.last_modified = datetime.now(timezone.utc).isoformat()
        return True

    def get_node_count(self) -> int:
        """Get total number of nodes in the graph."""
        #with self.lock:
        return self.graph.number_of_nodes()
    
    def get_edge_count(self) -> int:
        """Get total number of edges in the graph."""
        #with self.lock:
        return self.graph.number_of_edges()
    
    def get_nodes_by_type(self, node_type: NodeType) -> List[str]:
        """
        Get all nodes of a specific type.
        
        Args:
            node_type: Type of nodes to retrieve
            
        Returns:
            List of node identifiers
        """
        #with self.lock:
        return [
            node_id for node_id, attributes in self.graph.nodes(data=True)
            if attributes.get('type') == node_type.value
        ]
    
    def get_neighbors(self, node_id: str) -> List[str]:
        """
        Get all neighboring nodes (both incoming and outgoing).
        
        Args:
            node_id: Node identifier
            
        Returns:
            List of neighboring node identifiers
        """
        #with self.lock:
        if not self.graph.has_node(node_id):
            return []
        
        predecessors = list(self.graph.predecessors(node_id))
        successors = list(self.graph.successors(node_id))
        return list(set(predecessors + successors))
    
    def get_high_confidence_edges(self, min_confidence: float = 0.8) -> List[Tuple[str, str, Dict]]:
        """
        Get edges with confidence score above threshold.
        
        Args:
            min_confidence: Minimum confidence threshold
            
        Returns:
            List of tuples (source, target, attributes)
        """
        #with self.lock:
        return [
            (source, target, attributes)
            for source, target, attributes in self.graph.edges(data=True)
            if attributes.get('confidence_score', 0) >= min_confidence
        ]
    
    def get_graph_data(self) -> Dict[str, Any]:
        """
        Export graph data for visualization.
        
        Returns:
            Dictionary containing nodes and edges for frontend visualization
        """
        #with self.lock:
        nodes = []
        edges = []
        
        # Format nodes for visualization
        for node_id, attributes in self.graph.nodes(data=True):
            node_data = {
                'id': node_id,
                'label': node_id,
                'type': attributes.get('type', 'unknown'),
                'metadata': attributes.get('metadata', {}),
                'added_timestamp': attributes.get('added_timestamp')
            }
            
            # Color coding by type - now returns color objects for enhanced visualization
            type_colors = {
                'domain': {
                    'background': '#00ff41',
                    'border': '#00aa2e',
                    'highlight': {'background': '#44ff75', 'border': '#00ff41'},
                    'hover': {'background': '#22ff63', 'border': '#00cc35'}
                },
                'ip': {
                    'background': '#ff9900',
                    'border': '#cc7700',
                    'highlight': {'background': '#ffbb44', 'border': '#ff9900'},
                    'hover': {'background': '#ffaa22', 'border': '#dd8800'}
                },
                'certificate': {
                    'background': '#c7c7c7',
                    'border': '#999999',
                    'highlight': {'background': '#e0e0e0', 'border': '#c7c7c7'},
                    'hover': {'background': '#d4d4d4', 'border': '#aaaaaa'}
                },
                'asn': {
                    'background': '#00aaff',
                    'border': '#0088cc',
                    'highlight': {'background': '#44ccff', 'border': '#00aaff'},
                    'hover': {'background': '#22bbff', 'border': '#0099dd'}
                }
            }
            
            node_color_config = type_colors.get(attributes.get('type', 'unknown'), type_colors['domain'])
            node_data['color'] = node_color_config
            
            # Pass the has_valid_cert metadata for styling
            if 'metadata' in attributes and 'has_valid_cert' in attributes['metadata']:
                node_data['has_valid_cert'] = attributes['metadata']['has_valid_cert']

            nodes.append(node_data)
        
        # Format edges for visualization
        for source, target, attributes in self.graph.edges(data=True):
            edge_data = {
                'from': source,
                'to': target,
                'label': attributes.get('relationship_type', ''),
                'confidence_score': attributes.get('confidence_score', 0),
                'source_provider': attributes.get('source_provider', ''),
                'discovery_timestamp': attributes.get('discovery_timestamp')
            }
            
            # Enhanced edge styling based on confidence
            confidence = attributes.get('confidence_score', 0)
            if confidence >= 0.8:
                edge_data['color'] = {
                    'color': '#00ff41',
                    'highlight': '#44ff75',
                    'hover': '#22ff63',
                    'inherit': False
                }
                edge_data['width'] = 4
            elif confidence >= 0.6:
                edge_data['color'] = {
                    'color': '#ff9900',
                    'highlight': '#ffbb44', 
                    'hover': '#ffaa22',
                    'inherit': False
                }
                edge_data['width'] = 3
            else:
                edge_data['color'] = {
                    'color': '#666666',
                    'highlight': '#888888',
                    'hover': '#777777',
                    'inherit': False
                }
                edge_data['width'] = 2
            
            # Add dashed line for low confidence
            if confidence < 0.6:
                edge_data['dashes'] = [5, 5]
            
            edges.append(edge_data)
        
        return {
            'nodes': nodes,
            'edges': edges,
            'statistics': {
                'node_count': len(nodes),
                'edge_count': len(edges),
                'creation_time': self.creation_time,
                'last_modified': self.last_modified
            }
        }
    
    def export_json(self) -> Dict[str, Any]:
        """
        Export complete graph data as JSON for download.
        
        Returns:
            Dictionary containing complete graph data with metadata
        """
        #with self.lock:
        # Get basic graph data
        graph_data = self.get_graph_data()
        
        # Add comprehensive metadata
        export_data = {
            'export_metadata': {
                'export_timestamp': datetime.now(timezone.utc).isoformat(),
                'graph_creation_time': self.creation_time,
                'last_modified': self.last_modified,
                'total_nodes': self.graph.number_of_nodes(),
                'total_edges': self.graph.number_of_edges(),
                'graph_format': 'dnsrecon_v1'
            },
            'nodes': graph_data['nodes'],
            'edges': graph_data['edges'],
            'node_types': [node_type.value for node_type in NodeType],
            'relationship_types': [
                {
                    'name': rel_type.relationship_name,
                    'default_confidence': rel_type.default_confidence
                }
                for rel_type in RelationshipType
            ],
            'confidence_distribution': self._get_confidence_distribution()
        }
        
        return export_data
    
    def _get_confidence_distribution(self) -> Dict[str, int]:
        """Get distribution of confidence scores."""
        distribution = {'high': 0, 'medium': 0, 'low': 0}
        
        for _, _, attributes in self.graph.edges(data=True):
            confidence = attributes.get('confidence_score', 0)
            if confidence >= 0.8:
                distribution['high'] += 1
            elif confidence >= 0.6:
                distribution['medium'] += 1
            else:
                distribution['low'] += 1
        
        return distribution
    
    def get_statistics(self) -> Dict[str, Any]:
        """
        Get comprehensive graph statistics.
        
        Returns:
            Dictionary containing various graph metrics
        """
        #with self.lock:
        stats = {
            'basic_metrics': {
                'total_nodes': self.graph.number_of_nodes(),
                'total_edges': self.graph.number_of_edges(),
                'creation_time': self.creation_time,
                'last_modified': self.last_modified
            },
            'node_type_distribution': {},
            'relationship_type_distribution': {},
            'confidence_distribution': self._get_confidence_distribution(),
            'provider_distribution': {}
        }
        
        # Node type distribution
        for node_type in NodeType:
            count = len(self.get_nodes_by_type(node_type))
            stats['node_type_distribution'][node_type.value] = count
        
        # Relationship type distribution
        for _, _, attributes in self.graph.edges(data=True):
            rel_type = attributes.get('relationship_type', 'unknown')
            stats['relationship_type_distribution'][rel_type] = \
                stats['relationship_type_distribution'].get(rel_type, 0) + 1
        
        # Provider distribution
        for _, _, attributes in self.graph.edges(data=True):
            provider = attributes.get('source_provider', 'unknown')
            stats['provider_distribution'][provider] = \
                stats['provider_distribution'].get(provider, 0) + 1
        
        return stats
    
    def clear(self) -> None:
        """Clear all nodes and edges from the graph."""
        #with self.lock:
        self.graph.clear()
        self.creation_time = datetime.now(timezone.utc).isoformat()
        self.last_modified = self.creation_time