"""
Ontology Discovery - AGS Ontology Discovery and Management

This module handles discovering and managing ontology information from 
Altair Graph Studio (AGS) graphmarts, based on empirical findings from
production AGS investigation.

Core functionality:
1. Discover ontology URIs via linkedOntology and model properties
2. Extract complete ontology graphs as RDF triples
3. Parse classes, properties, and basic metadata
4. Build summaries for LLM context
"""

import os
import sys
import json
import asyncio
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime

from pyanzo import AnzoClient

# Debug function for MCP compliance - logs to stderr only
def debug_log(message: str):
    """Log debug messages to stderr to avoid MCP protocol interference."""
    # Respect ENABLE_AGENT_DEBUG environment variable for MCP compatibility
    if os.getenv('ENABLE_AGENT_DEBUG', 'false').lower() == 'true':
        print(f"[DEBUG] {message}", file=sys.stderr)

try:
    # Try relative import first (for package usage)
    from .models import (
        GraphmartConfig,
        OntologyInfo,
        ClassInfo,
        PropertyInfo,
        OntologySummary,
        OntologyMetadata,
        OntologyExtractionResult
    )
except ImportError:
    # Fall back to absolute import (for testing and standalone usage)
    from models import (
        GraphmartConfig,
        OntologyInfo,
        ClassInfo,
        PropertyInfo,
        OntologySummary,
        OntologyMetadata,
        OntologyExtractionResult
    )


class OntologyDiscoveryEngine:
    """
    Simple, focused ontology discovery engine for AGS graphmarts.
    
    Based on empirical investigation of production AGS systems, this engine:
    1. Discovers ontology URIs via linkedOntology and model properties
    2. Extracts complete ontology graphs from each URI
    3. Parses basic class/property information for LLM context
    """
    
    def __init__(self, anzo_client: AnzoClient, config: GraphmartConfig, filter_empty: bool = True):
        """Initialize the ontology discovery engine.
        
        Args:
            anzo_client: The AnzoClient for SPARQL queries
            config: Graphmart configuration
            filter_empty: If True, filter out classes with no instances and properties with no usage
        """
        self.anzo_client = anzo_client
        self.config = config
        self.filter_empty = filter_empty
        
        # Cache directory for ontology data
        self.cache_dir = os.path.join(
            os.path.dirname(__file__), 
            '..', 
            'data',
            'ontology_cache'
        )
        os.makedirs(self.cache_dir, exist_ok=True)
    
    async def discover_graphmart_ontologies(self) -> OntologyInfo:
        """
        Main discovery method - finds and extracts all ontologies from the graphmart.
        
        Returns:
            OntologyInfo: Complete ontology information for LLM context
        """
        debug_log(f"🔍 Starting ontology discovery for: {self.config.graphmart_uri}")
        
        try:
            # Step 1: Find all ontology URIs (linkedOntology + model properties)
            ontology_refs = await self._discover_ontology_references()
            
            if not ontology_refs:
                debug_log("❌ No ontology references found")
                return self._create_empty_ontology_info()
            
            debug_log(f"✅ Found {len(ontology_refs)} ontology references")
            
            # Step 2: Extract schemas from each ontology
            all_classes = {}
            all_properties = {}
            
            for source_property, ontology_uri in ontology_refs:
                debug_log(f"🔄 Extracting ontology: {ontology_uri}")
                
                classes, properties = await self._extract_ontology_schema(ontology_uri)
                
                # Merge into combined results
                all_classes.update(classes)
                all_properties.update(properties)
                
                debug_log(f"  ✅ Found {len(classes)} classes, {len(properties)} properties")
            
            # Step 3: Add usage counts to determine which classes/properties are actually used
            debug_log("🔢 Adding usage counts to filter populated classes and properties...")
            await self._add_usage_counts(all_classes, all_properties)
            
            # Step 4: Optionally filter to only include populated classes and properties
            if self.filter_empty:
                final_classes, final_properties = self._filter_populated_only(all_classes, all_properties)
                debug_log(f"🔍 Filtered to populated only due to filter_empty=True")
            else:
                final_classes, final_properties = all_classes, all_properties
                debug_log(f"📋 Keeping all classes/properties due to filter_empty=False")
            
            # Step 5: Create summary from final classes/properties
            domain_summary = self._generate_domain_summary(final_classes, final_properties)
            
            # Step 6: Generate comprehensive LLM-friendly context from final data
            class_map, prop_map = self._create_compact_uri_mappings(final_classes, final_properties)
            llm_context = self._generate_llm_friendly_context(final_classes, final_properties, self._extract_namespaces(final_classes, final_properties))
            
            # Create reverse mappings for URI expansion
            reverse_class_map = {v: k for k, v in class_map.items()}
            reverse_property_map = {v: k for k, v in prop_map.items()}
            
            # Calculate total instances
            total_instances = sum(cls.instance_count or 0 for cls in final_classes.values())
            
            ontology_info = OntologyInfo(
                classes=final_classes,
                properties=final_properties,
                namespaces=self._extract_namespaces(final_classes, final_properties),
                domain_summary=domain_summary,
                llm_context=llm_context,
                compact_class_map=class_map,
                compact_property_map=prop_map,
                reverse_class_map=reverse_class_map,
                reverse_property_map=reverse_property_map,
                total_instances=total_instances
            )
            
            debug_log(f"🎯 Discovery complete: {len(final_classes)} final classes, {len(final_properties)} final properties")
            debug_log(f"📊 Total instances: {total_instances}")
            return ontology_info
            
        except Exception as e:
            debug_log(f"❌ Ontology discovery failed: {e}")
            return self._create_empty_ontology_info()
    
    async def _discover_ontology_references(self) -> List[Tuple[str, str]]:
        """
        Discover all ontology URIs using wildcard subject approach.
        Queries for all ontology-related properties within the GraphMart metadata graph,
        regardless of which subject is making the reference.
        
        Returns:
            List of (source_property, ontology_uri) tuples
        """
        ontology_refs = []
        
        debug_log("🔍 Discovering ontology references with wildcard subjects...")
        
        # Single query to find all ontology references with wildcard subjects
        discovery_query = f"""
        SELECT ?subject ?property ?ontologyUri
        WHERE {{
          GRAPH <{self.config.graphmart_uri}> {{
            ?subject ?property ?ontologyUri .
            FILTER(?property IN (
                <http://cambridgesemantics.com/ontologies/Graphmarts#linkedOntology>,
                <http://cambridgesemantics.com/ontologies/Graphmarts#model>,
                <http://cambridgesemantics.com/ontologies/Graphmarts#ontology>
            ))
          }}
        }}
        """
        
        try:
            result = self.anzo_client.query_journal(discovery_query)
            table = result.as_table_results().as_list()
            
            # Track what we find for logging
            property_counts = {}
            
            for row in table:
                if isinstance(row, list) and len(row) >= 3:
                    subject = str(row[0]) if row[0] else None
                    property_uri = str(row[1]) if row[1] else None
                    ontology_uri = str(row[2]) if row[2] else None
                    
                    if ontology_uri and property_uri:
                        # Extract property name for cleaner logging
                        property_name = property_uri.split('#')[-1] if '#' in property_uri else property_uri.split('/')[-1]
                        
                        ontology_refs.append((property_name, ontology_uri))
                        
                        # Track counts for summary
                        property_counts[property_name] = property_counts.get(property_name, 0) + 1
                        
                        debug_log(f"  📍 Found: {ontology_uri} (via {property_name})")
                        if subject:
                            debug_log(f"      Subject: {subject}")
            
            # Summary
            debug_log(f"✅ Found {len(ontology_refs)} ontology references")
            for prop, count in property_counts.items():
                debug_log(f"  - {prop}: {count} references")
        
        except Exception as e:
            debug_log(f"❌ Failed to discover ontology references: {e}")
        
        # NEW: Discover File-Backed Linked Data Set (LDS) ontology references
        debug_log("🔍 Discovering File-Backed LDS ontology references...")
        
        try:
            # Step 1: Find LoadDatasetStep objects that reference FileBackedLinkedDataSet
            lds_discovery_query = f"""
            SELECT ?loadStep ?ldsObject
            WHERE {{
              GRAPH <{self.config.graphmart_uri}> {{
                ?loadStep <http://cambridgesemantics.com/ontologies/Graphmarts#linkedDataset> ?ldsObject .
                FILTER(CONTAINS(STR(?ldsObject), "FileBackedLinkedDataSet"))
              }}
            }}
            """
            
            lds_result = self.anzo_client.query_journal(lds_discovery_query)
            lds_table = lds_result.as_table_results().as_list()
            
            # Step 2: For each LDS object, query for its ontology references
            lds_ontology_count = 0
            for row in lds_table:
                if isinstance(row, list) and len(row) >= 2:
                    load_step = str(row[0]) if row[0] else None
                    lds_object = str(row[1]) if row[1] else None
                    
                    if lds_object:
                        debug_log(f"  🔗 Found LDS object: {lds_object}")
                        
                        # Query the LDS object for ontology references
                        lds_ontology_query = f"""
                        SELECT ?ontologyUri
                        WHERE {{
                          GRAPH <{lds_object}> {{
                            ?subject <http://cambridgesemantics.com/ontologies/2009/05/LinkedData#ontology> ?ontologyUri .
                          }}
                        }}
                        """
                        
                        try:
                            ont_result = self.anzo_client.query_journal(lds_ontology_query)
                            ont_table = ont_result.as_table_results().as_list()
                            
                            for ont_row in ont_table:
                                if isinstance(ont_row, list) and len(ont_row) >= 1:
                                    ontology_uri = str(ont_row[0]) if ont_row[0] else None
                                    if ontology_uri:
                                        ontology_refs.append(("linkedDataset", ontology_uri))
                                        lds_ontology_count += 1
                                        debug_log(f"    📍 Found LDS ontology: {ontology_uri}")
                                        
                        except Exception as lds_ont_e:
                            debug_log(f"    ❌ Failed to query LDS ontologies for {lds_object}: {lds_ont_e}")
            
            if lds_ontology_count > 0:
                debug_log(f"✅ Found {lds_ontology_count} ontology references via File-Backed LDS")
            else:
                debug_log("ℹ️  No File-Backed LDS ontology references found")
                
        except Exception as lds_e:
            debug_log(f"❌ Failed to discover File-Backed LDS references: {lds_e}")
        
        # Filter out system ontologies that shouldn't be included in discovery
        system_ontologies = {
            "http://cambridgesemantics.com/ontologies/Graphmarts#Managed"  # AGS system metadata, always empty
        }
        
        original_count = len(ontology_refs)
        ontology_refs = [
            (prop, uri) for prop, uri in ontology_refs 
            if uri not in system_ontologies
        ]
        
        filtered_count = original_count - len(ontology_refs)
        if filtered_count > 0:
            debug_log(f"🚫 Filtered out {filtered_count} system ontologies from discovery (Graphmarts#Managed)")
        
        return ontology_refs
    
    async def _extract_ontology_metadata(self, ontology_uri: str) -> 'OntologyMetadata':
        """
        Extract metadata about the ontology itself.
        
        Args:
            ontology_uri: URI of the ontology graph
            
        Returns:
            OntologyMetadata object with title, description, etc.
        """
        try:
            # Query for ontology-level metadata
            metadata_query = f"""
            SELECT ?predicate ?object
            WHERE {{
              GRAPH <{ontology_uri}> {{
                <{ontology_uri}> ?predicate ?object
              }}
            }}
            """
            
            result = self.anzo_client.query_journal(metadata_query)
            table = result.as_table_results().as_list()
            
            # Extract metadata properties
            title = None
            description = None
            version = None
            creator = None
            
            for row in table:
                if len(row) >= 2:
                    predicate = str(row[0]) if row[0] else ""
                    obj = str(row[1]) if row[1] else ""
                    
                    if predicate == "http://www.w3.org/2000/01/rdf-schema#label":
                        title = obj
                    elif predicate == "http://purl.org/dc/terms/title":
                        title = obj  # Dublin Core title takes precedence
                    elif predicate == "http://www.w3.org/2000/01/rdf-schema#comment":
                        description = obj
                    elif predicate == "http://www.w3.org/2002/07/owl#versionInfo":
                        version = obj
                    elif predicate == "http://purl.org/dc/terms/creator":
                        creator = obj
            
            # Create metadata object
            metadata = OntologyMetadata(
                uri=ontology_uri,
                title=title,
                description=description,
                version=version,
                creator=creator
            )
            
            return metadata
            
        except Exception as e:
            debug_log(f"  ⚠️ Could not extract metadata for {ontology_uri}: {e}")
            return OntologyMetadata(uri=ontology_uri)

    async def _extract_ontology_schema(self, ontology_uri: str) -> Tuple[Dict[str, ClassInfo], Dict[str, PropertyInfo]]:
        """
        Extract complete schema from an ontology graph.
        
        Args:
            ontology_uri: URI of the ontology graph to extract
            
        Returns:
            Tuple of (classes_dict, properties_dict)
        """
        classes = {}
        properties = {}
        
        try:
            # First, count the number of classes in this ontology to apply filtering
            class_count_query = f"""
            SELECT (COUNT(DISTINCT ?class) AS ?classCount)
            WHERE {{
              GRAPH <{ontology_uri}> {{
                ?class a ?classType .
                FILTER(?classType IN (
                    <http://www.w3.org/2000/01/rdf-schema#Class>, 
                    <http://www.w3.org/2002/07/owl#Class>
                ))
              }}
            }}
            """
            
            debug_log(f"  🔍 Checking class count for ontology: {ontology_uri}")
            count_result = self.anzo_client.query_journal(class_count_query)
            count_table = count_result.as_table_results().as_list()
            
            class_count = 0
            if count_table and len(count_table) > 0 and count_table[0]:
                class_count = int(count_table[0][0]) if count_table[0][0] else 0
            
            debug_log(f"  📊 Ontology has {class_count} classes")
            
            # Skip ontologies with more than 200 classes
            if class_count > 200:
                debug_log(f"  ⏭️  Skipping large ontology {ontology_uri} ({class_count} classes > 200 threshold)")
                return {}, {}
            
            # Query all triples in the ontology graph
            ontology_query = f"""
            SELECT ?subject ?predicate ?object
            WHERE {{
              GRAPH <{ontology_uri}> {{
                ?subject ?predicate ?object
              }}
            }}
            """
            
            debug_log(f"  🔄 Querying ontology graph: {ontology_uri}")
            result = self.anzo_client.query_journal(ontology_query)
            table = result.as_table_results().as_list()
            
            debug_log(f"  📊 Found {len(table)} triples in ontology")
            
            # Parse triples to extract classes and properties
            class_uris = set()
            property_uris = set()
            class_labels = {}
            property_labels = {}
            class_comments = {}
            property_comments = {}
            property_domains = {}
            property_ranges = {}
            object_property_uris = set()  # Track owl:ObjectProperty
            datatype_property_uris = set()  # Track owl:DatatypeProperty
            
            for row in table:
                if len(row) >= 3:
                    subject = str(row[0]) if row[0] else ""
                    predicate = str(row[1]) if row[1] else ""
                    obj = str(row[2]) if row[2] else ""
                    
                    # Find classes (subjects that are of type Class)
                    if predicate == "http://www.w3.org/1999/02/22-rdf-syntax-ns#type":
                        if obj in ["http://www.w3.org/2000/01/rdf-schema#Class", 
                                  "http://www.w3.org/2002/07/owl#Class"]:
                            class_uris.add(subject)
                        elif obj == "http://www.w3.org/2002/07/owl#ObjectProperty":
                            property_uris.add(subject)
                            object_property_uris.add(subject)
                        elif obj in ["http://www.w3.org/1999/02/22-rdf-syntax-ns#Property",
                                    "http://www.w3.org/2002/07/owl#DatatypeProperty"]:
                            property_uris.add(subject)
                            datatype_property_uris.add(subject)
                    
                    # Collect labels
                    elif predicate == "http://www.w3.org/2000/01/rdf-schema#label":
                        if subject in class_uris or any(subject.startswith(prefix) for prefix in ["http://", "https://"]):
                            class_labels[subject] = obj
                        property_labels[subject] = obj
                    
                    # Collect comments/descriptions
                    elif predicate == "http://www.w3.org/2000/01/rdf-schema#comment":
                        class_comments[subject] = obj
                        property_comments[subject] = obj
                    
                    # Collect property domains
                    elif predicate == "http://www.w3.org/2000/01/rdf-schema#domain":
                        if subject not in property_domains:
                            property_domains[subject] = []
                        property_domains[subject].append(obj)
                    
                    # Collect property ranges
                    elif predicate == "http://www.w3.org/2000/01/rdf-schema#range":
                        if subject not in property_ranges:
                            property_ranges[subject] = []
                        property_ranges[subject].append(obj)
            
            # Create ClassInfo objects (counts will be added later)
            for class_uri in class_uris:
                # Skip invalid URIs during schema extraction
                if not self._is_valid_uri(class_uri):
                    debug_log(f"  ⚠️ Skipping invalid class URI during extraction: {class_uri}")
                    continue
                    
                class_name = class_labels.get(class_uri) or self._extract_name_from_uri(class_uri)
                classes[class_uri] = ClassInfo(
                    name=class_name,
                    uri=class_uri,
                    description=class_comments.get(class_uri),
                    instance_count=None,  # Will be populated by _add_usage_counts
                    sample_instances=[],
                    ontology_uri=ontology_uri  # NEW: Track which ontology this class belongs to
                )
            
            # Create PropertyInfo objects (counts will be added later)
            for prop_uri in property_uris:
                # Skip invalid URIs during schema extraction
                if not self._is_valid_uri(prop_uri):
                    debug_log(f"  ⚠️ Skipping invalid property URI during extraction: {prop_uri}")
                    continue
                    
                prop_name = property_labels.get(prop_uri) or self._extract_name_from_uri(prop_uri)
                # Determine property type: ObjectProperty, DatatypeProperty, or unknown
                is_obj = True if prop_uri in object_property_uris else (False if prop_uri in datatype_property_uris else None)
                properties[prop_uri] = PropertyInfo(
                    name=prop_name,
                    uri=prop_uri,
                    description=property_comments.get(prop_uri),
                    domain_classes=property_domains.get(prop_uri, []),
                    range_classes=property_ranges.get(prop_uri, []),
                    usage_frequency=None,  # Will be populated by _add_usage_counts
                    ontology_uri=ontology_uri,  # Track which ontology this property belongs to
                    is_object_property=is_obj  # True=ObjectProperty, False=DatatypeProperty, None=unknown
                )
            
            debug_log(f"  ✅ Extracted {len(classes)} valid classes and {len(properties)} valid properties")
            return classes, properties
            
        except Exception as e:
            debug_log(f"  ❌ Failed to extract schema from {ontology_uri}: {e}")
            return {}, {}

    async def extract_ontology_with_metadata(self, ontology_uri: str) -> OntologyExtractionResult:
        """
        Extract complete ontology information including metadata, classes, and properties.
        
        Args:
            ontology_uri: URI of the ontology graph to extract
            
        Returns:
            OntologyExtractionResult with metadata and schema
        """
        try:
            debug_log(f"  🔄 Querying ontology graph: {ontology_uri}")
            
            # Extract metadata and schema in parallel for efficiency
            metadata_task = self._extract_ontology_metadata(ontology_uri)
            schema_task = self._extract_ontology_schema(ontology_uri)
            
            metadata, (classes, properties) = await asyncio.gather(metadata_task, schema_task)
            
            # Derive title from content if none found in metadata
            if not metadata.title:
                metadata.derived_title = self._derive_title_from_content(classes, properties, ontology_uri)
            
            debug_log(f"  ✅ Found {len(classes)} classes, {len(properties)} properties")
            
            return OntologyExtractionResult(
                metadata=metadata,
                classes=classes,
                properties=properties
            )
            
        except Exception as e:
            debug_log(f"  ❌ Failed to extract ontology {ontology_uri}: {e}")
            return OntologyExtractionResult(
                metadata=OntologyMetadata(uri=ontology_uri),
                classes={},
                properties={}
            )

    def _derive_title_from_content(self, classes: Dict[str, ClassInfo], properties: Dict[str, PropertyInfo], ontology_uri: str) -> str:
        """
        Derive a meaningful title from the ontology content.
        
        Args:
            classes: Dictionary of class information
            properties: Dictionary of property information
            ontology_uri: The ontology URI
            
        Returns:
            Derived title string
        """
        # Extract namespace themes from classes
        namespace_themes = {}
        for class_uri in classes.keys():
            if '#' in class_uri:
                namespace = class_uri.split('#')[0]
                if '/ontologies/' in namespace:
                    theme = namespace.split('/ontologies/')[-1]
                    namespace_themes[theme] = namespace_themes.get(theme, 0) + 1
                elif '/' in namespace:
                    theme = namespace.split('/')[-1]
                    namespace_themes[theme] = namespace_themes.get(theme, 0) + 1
        
        # Find the most common theme
        if namespace_themes:
            dominant_theme = max(namespace_themes.items(), key=lambda x: x[1])[0]
            return f"{dominant_theme.title()} Ontology"
        
        # Fallback to URI-based extraction
        if 'CommsLinking' in ontology_uri:
            return "Communications Linking Ontology"
        elif 'Communications' in ontology_uri or any('Communications' in cls for cls in classes.keys()):
            return "Communications Ontology" 
        elif 'Recipes' in ontology_uri or any('Recipes' in cls for cls in classes.keys()):
            return "Recipes Ontology"
        else:
            # Last resort: extract from URI
            if '#' in ontology_uri:
                return ontology_uri.split('#')[-1]
            elif '/' in ontology_uri:
                parts = ontology_uri.rstrip('/').split('/')
                return parts[-1] if parts[-1] else parts[-2] if len(parts) > 1 else "Unknown Ontology"
            else:
                return "Unknown Ontology"
    
    def _extract_name_from_uri(self, uri: str) -> str:
        """Extract local name from URI."""
        if "#" in uri:
            return uri.split("#")[-1]
        elif "/" in uri:
            return uri.split("/")[-1]
        else:
            return uri
    
    async def _add_usage_counts(self, classes: Dict[str, ClassInfo], properties: Dict[str, PropertyInfo]) -> None:
        """
        OPTIMIZED: Add instance counts for classes and usage frequencies for properties.
        Uses one query per class to get both class instances and all property usage counts.
        
        Args:
            classes: Dictionary of class information to update
            properties: Dictionary of property information to update
        """
        debug_log(f"🔢 OPTIMIZED: Adding usage counts for {len(classes)} classes and {len(properties)} properties")
        
        # Group properties by their domain classes for efficient querying
        class_properties = {}
        for prop_uri, prop_info in properties.items():
            if prop_info.domain_classes:
                for domain_class in prop_info.domain_classes:
                    if domain_class not in class_properties:
                        class_properties[domain_class] = []
                    class_properties[domain_class].append(prop_uri)
        
        valid_class_count = 0
        invalid_class_count = 0
        total_properties_counted = 0
        
        # Process each class with its properties in one optimized query
        for class_uri, class_info in classes.items():
            # Skip invalid URIs
            if not self._is_valid_uri(class_uri):
                debug_log(f"  ⚠️ Skipping invalid URI: {class_uri}")
                class_info.instance_count = 0
                invalid_class_count += 1
                continue
                
            try:
                debug_log(f"  🔍 Starting count query for class: {class_info.name} ({class_uri})")
                
                # Get properties for this class
                class_props = class_properties.get(class_uri, [])
                valid_class_props = [uri for uri in class_props if self._is_valid_uri(uri)]
                
                debug_log(f"  📝 Found {len(valid_class_props)} properties for this class")
                
                # Build optimized query using UNION of subqueries with BIND for type identification
                union_clauses = []
                
                # First subquery: class count
                union_clauses.append(f"""
                {{
                    SELECT ?type ?count WHERE {{
                        {{
                            SELECT (COUNT(?instance) AS ?count) WHERE {{
                                ?instance a <{class_uri}> .
                            }}
                        }}
                        BIND("class" AS ?type)
                    }}
                }}""")
                
                # Additional subqueries: one for each property count
                for i, prop_uri in enumerate(valid_class_props):
                    union_clauses.append(f"""
                {{
                    SELECT ?type ?count WHERE {{
                        {{
                            SELECT (COUNT(?instance) AS ?count) WHERE {{
                                ?instance a <{class_uri}> .
                                ?instance <{prop_uri}> ?value .
                            }}
                        }}
                        BIND("prop_{i}" AS ?type)
                    }}
                }}""")
                
                optimized_query = f"""
                SELECT ?type ?count
                WHERE {{
                    {' UNION '.join(union_clauses)}
                }}
                ORDER BY ?type
                """
                
                debug_log(f"  🚀 Executing query for class: {class_info.name}")
                
                # Execute with 30 second timeout
                try:
                    result = await asyncio.wait_for(
                        asyncio.get_event_loop().run_in_executor(
                            None,
                            lambda: self.anzo_client.query_graphmart(self.config.graphmart_uri, optimized_query)
                        ),
                        timeout=30
                    )
                    debug_log(f"  ✅ Query completed for class: {class_info.name}")
                except asyncio.TimeoutError:
                    debug_log(f"  ⏰ Query timed out for class: {class_info.name}")
                    continue  # Skip this class and continue with others
                table = result.as_table_results().as_list()
                
                if table and len(table) > 0:
                    # Parse UNION results - each row has ?type and ?count
                    class_count = 0
                    property_counts = {}
                    
                    for row in table:
                        if len(row) >= 2:
                            type_str = str(row[0]) if row[0] else ""
                            count_val = int(row[1]) if row[1] else 0
                            
                            if type_str == "class":
                                class_count = count_val
                            elif type_str.startswith("prop_"):
                                prop_index = int(type_str.split("_")[1])
                                if prop_index < len(valid_class_props):
                                    prop_uri = valid_class_props[prop_index]
                                    property_counts[prop_uri] = count_val
                    
                    # Set class count
                    class_info.instance_count = class_count
                    valid_class_count += 1
                    
                    if class_count > 0:
                        debug_log(f"  📊 {class_info.name}: {class_count} instances")
                    
                    # Set property counts
                    for prop_uri, prop_count in property_counts.items():
                        if prop_uri in properties:
                            properties[prop_uri].usage_frequency = prop_count
                            total_properties_counted += 1
                            if prop_count > 0:
                                debug_log(f"    🔗 {properties[prop_uri].name}: {prop_count} uses")
                else:
                    class_info.instance_count = 0
                    valid_class_count += 1
                    
            except Exception as e:
                debug_log(f"  ❌ Failed to count for {class_info.name}: {e}")
                class_info.instance_count = 0
        
        # Handle properties not associated with any domain class (standalone properties)
        orphan_properties = [uri for uri in properties.keys() 
                           if not any(uri in class_properties.get(cls, []) for cls in class_properties)]
        
        debug_log(f"🔗 Processing {len(orphan_properties)} orphan properties...")
        valid_orphan_count = 0
        
        for prop_uri in orphan_properties:
            if not self._is_valid_uri(prop_uri):
                properties[prop_uri].usage_frequency = 0
                continue
                
            try:
                usage_query = f"""
                SELECT (COUNT(*) AS ?count)
                WHERE {{
                  ?subject <{prop_uri}> ?object .
                }}
                """
                
                # Execute with 30 second timeout
                try:
                    result = await asyncio.wait_for(
                        asyncio.get_event_loop().run_in_executor(
                            None,
                            lambda: self.anzo_client.query_graphmart(self.config.graphmart_uri, usage_query)
                        ),
                        timeout=30
                    )
                except asyncio.TimeoutError:
                    debug_log(f"  ⏰ Usage query timed out for property: {prop_uri}")
                    properties[prop_uri].usage_frequency = 0
                    continue
                table = result.as_table_results().as_list()
                
                if table and len(table) > 0 and table[0]:
                    count = int(table[0][0]) if table[0][0] else 0
                    properties[prop_uri].usage_frequency = count
                    valid_orphan_count += 1
                    total_properties_counted += 1
                    
                    if count > 0:
                        debug_log(f"  � {properties[prop_uri].name}: {count} uses (orphan)")
                else:
                    properties[prop_uri].usage_frequency = 0
                    
            except Exception as e:
                debug_log(f"  ❌ Failed to count orphan property {properties[prop_uri].name}: {e}")
                properties[prop_uri].usage_frequency = 0
        
        # Log summary
        populated_classes = sum(1 for cls in classes.values() if cls.has_instances())
        populated_properties = sum(1 for prop in properties.values() if prop.has_usage())
        
        debug_log(f"✅ OPTIMIZED usage counts complete:")
        debug_log(f"  - Valid classes processed: {valid_class_count}, Invalid URIs skipped: {invalid_class_count}")
        debug_log(f"  - Properties counted via class queries: {total_properties_counted - valid_orphan_count}")
        debug_log(f"  - Orphan properties processed: {valid_orphan_count}")
        debug_log(f"  - Classes with instances: {populated_classes}/{len(classes)}")
        debug_log(f"  - Properties with usage: {populated_properties}/{len(properties)}")
        debug_log(f"  - Total queries executed: ~{valid_class_count + valid_orphan_count} (vs {len(classes) + len(properties)} in old approach)")
    
    def _is_valid_uri(self, uri: str) -> bool:
        """
        Check if a URI is valid for SPARQL queries.
        
        Args:
            uri: The URI to validate
            
        Returns:
            True if the URI is valid, False otherwise
        """
        if not uri:
            return False
            
        # Check for proper URI scheme (http://, https://, urn:, etc.)
        if ':' not in uri:
            return False
            
        # Check for common valid URI patterns
        if (uri.startswith('http://') or 
            uri.startswith('https://') or 
            uri.startswith('urn:') or 
            uri.startswith('file:') or
            uri.startswith('ftp:')):
            return True
            
        # Check for other schemes with colon
        parts = uri.split(':', 1)
        if len(parts) == 2 and len(parts[0]) > 0:
            # Has a scheme, likely valid
            return True
            
        return False
    
    def _filter_populated_only(self, classes: Dict[str, ClassInfo], properties: Dict[str, PropertyInfo]) -> Tuple[Dict[str, ClassInfo], Dict[str, PropertyInfo]]:
        """
        Filter to only include classes with instances and properties with usage.
        
        Args:
            classes: Dictionary of all class information
            properties: Dictionary of all property information
            
        Returns:
            Tuple of (filtered_classes, filtered_properties)
        """
        filtered_classes = {uri: info for uri, info in classes.items() if info.has_instances()}
        filtered_properties = {uri: info for uri, info in properties.items() if info.has_usage()}
        
        debug_log(f"🔍 Filtered to populated only:")
        debug_log(f"  - Classes: {len(filtered_classes)}/{len(classes)} have instances")
        debug_log(f"  - Properties: {len(filtered_properties)}/{len(properties)} have usage")
        
        return filtered_classes, filtered_properties
    
    def _extract_namespaces(self, classes: Dict[str, ClassInfo], properties: Dict[str, PropertyInfo]) -> Dict[str, str]:
        """Extract common namespaces from URIs."""
        namespaces = {}
        all_uris = list(classes.keys()) + list(properties.keys())
        
        namespace_counts = {}
        for uri in all_uris:
            if "#" in uri:
                namespace = uri.split("#")[0] + "#"
            elif "/" in uri:
                parts = uri.split("/")
                namespace = "/".join(parts[:-1]) + "/"
            else:
                continue
            
            namespace_counts[namespace] = namespace_counts.get(namespace, 0) + 1
        
        # Keep namespaces used more than once
        for namespace, count in namespace_counts.items():
            if count > 1:
                # Generate appropriate prefix
                if 'www.w3.org/1999/02/22-rdf-syntax-ns' in namespace:
                    prefix = 'rdf'
                elif 'www.w3.org/2000/01/rdf-schema' in namespace:
                    prefix = 'rdfs'
                elif 'www.w3.org/2002/07/owl' in namespace:
                    prefix = 'owl'
                elif 'siemens.com' in namespace:
                    prefix = 'siemens'
                elif 'cambridgesemantics.com' in namespace:
                    prefix = 'cs'
                else:
                    # Generic prefix
                    prefix = f"ns{len(namespaces) + 1}"
                
                namespaces[prefix] = namespace
        
        return namespaces
    
    def _generate_domain_summary(self, classes: Dict[str, ClassInfo], properties: Dict[str, PropertyInfo]) -> str:
        """Generate a human-readable domain summary."""
        if not classes and not properties:
            return "Empty ontology - no classes or properties found."
        
        summary_parts = []
        
        if classes:
            class_names = [cls.name for cls in classes.values()]
            summary_parts.append(f"This knowledge graph contains {len(classes)} classes including: {', '.join(class_names[:5])}")
            if len(class_names) > 5:
                summary_parts[-1] += f" and {len(class_names) - 5} others"
        
        if properties:
            prop_names = [prop.name for prop in properties.values()]
            summary_parts.append(f"There are {len(properties)} properties including: {', '.join(prop_names[:5])}")
            if len(prop_names) > 5:
                summary_parts[-1] += f" and {len(prop_names) - 5} others"
        
        # Try to infer domain
        domain_hints = []
        for class_name in [cls.name.lower() for cls in classes.values()]:
            if any(term in class_name for term in ['recipe', 'material', 'product']):
                domain_hints.append("manufacturing/recipes")
            elif any(term in class_name for term in ['email', 'communication', 'message']):
                domain_hints.append("communications")
            elif any(term in class_name for term in ['experiment', 'lab', 'test']):
                domain_hints.append("laboratory/research")
        
        if domain_hints:
            summary_parts.append(f"The domain appears to be related to: {', '.join(set(domain_hints))}")
        
        return ". ".join(summary_parts) + "."
    
    def _generate_llm_friendly_context(
        self, 
        classes: Dict[str, ClassInfo], 
        properties: Dict[str, PropertyInfo],
        namespaces: Dict[str, str]
    ) -> str:
        """
        Generate a compact, hierarchical, LLM-friendly ontology context.
        
        Groups properties under their domain classes with ranges, similar to ontology_prompt.txt format.
        Uses compact URI mapping (c:0, p:0 style) and hierarchical organization by ontology source.
        """
        context_parts = []
        
        # Create compact URI mappings
        class_map, prop_map = self._create_compact_uri_mappings(classes, properties)
        
        # Header with compact stats
        context_parts.append("=== COMPACT ONTOLOGY CONTEXT ===")
        context_parts.append(f"Classes: {len(classes)} | Properties: {len(properties)}")
        context_parts.append("")
        
        # Compact domain overview
        domain_terms = self._extract_domain_keywords(classes, properties)
        context_parts.append(f"DOMAIN: {', '.join(domain_terms[:8])}")
        context_parts.append("")
        
        # Organize by ontology source (hierarchical)
        ontology_groups = self._group_by_ontology_source(classes, properties)
        
        for ontology_base, (ont_classes, ont_properties) in ontology_groups.items():
            if not ont_classes and not ont_properties:
                continue
                
            context_parts.append(f"# {ontology_base} Ontology Classes")
            
            # Group properties by their domain classes (including cross-ontology properties)
            class_properties = self._group_properties_by_domain_class(ont_classes, ont_properties, properties, classes)
            
            # Show each class with its properties
            for class_uri in sorted(ont_classes, key=lambda x: classes[x].name):
                class_info = classes[class_uri]
                compact_ref = class_map[class_uri]
                
                # Class header: c:0 ClassName; class
                context_parts.append(f"{compact_ref} {class_info.name}; class")
                
                # Properties for this class
                if class_uri in class_properties:
                    for prop_uri in sorted(class_properties[class_uri], key=lambda x: properties[x].name):
                        prop_info = properties[prop_uri]
                        prop_compact_ref = prop_map[prop_uri]
                        
                        # Property line with range: p:0 PropertyName; range c:1 ,functional
                        prop_line = f"\t{prop_compact_ref} {prop_info.name};"
                        
                        # Add range information
                        if prop_info.range_classes:
                            range_refs = []
                            for range_uri in prop_info.range_classes[:3]:  # Max 3 ranges
                                if range_uri in class_map:
                                    range_refs.append(class_map[range_uri])
                                else:
                                    # Try to infer type from URI or name
                                    range_type = self._infer_property_type(prop_info, range_uri)
                                    range_refs.append(range_type)
                            
                            if range_refs:
                                prop_line += f" range {' '.join(range_refs)}"
                        else:
                            # Infer type from property name/context
                            prop_type = self._infer_property_type(prop_info)
                            if prop_type:
                                prop_line += f" {prop_type}"
                        
                        # Add functional/cardinality hints if available
                        if hasattr(prop_info, 'is_functional') and prop_info.is_functional:
                            prop_line += " ,functional"
                        
                        # Add cross-ontology annotation if property is from different ontology
                        prop_ontology_base = self._extract_ontology_base(prop_uri)
                        if prop_ontology_base != ontology_base:
                            prop_line += f" [from {prop_ontology_base}]"
                        
                        context_parts.append(prop_line)
                
                context_parts.append("")  # Blank line after each class
            
            # Properties without clear domain classes (if any)
            orphaned_props = [p for p in ont_properties if not any(p in class_properties.get(c, []) for c in ont_classes)]
            if orphaned_props:
                context_parts.append("# Additional Properties")
                for prop_uri in sorted(orphaned_props, key=lambda x: properties[x].name):
                    prop_info = properties[prop_uri]
                    prop_compact_ref = prop_map[prop_uri]
                    prop_type = self._infer_property_type(prop_info)
                    context_parts.append(f"{prop_compact_ref} {prop_info.name}; {prop_type or 'str'}")
                context_parts.append("")
        
        # Compact URI reference mapping (abbreviated)
        context_parts.append("# COMPACT REFERENCES")
        context_parts.append("@prefix c: <c:>")
        context_parts.append("@prefix p: <p:>")
        context_parts.append("")
        
        # Show sample mappings for reference
        context_parts.append("# Sample class mappings:")
        for class_uri, compact_ref in list(class_map.items())[:5]:
            class_name = classes[class_uri].name
            context_parts.append(f"# {compact_ref} = {class_name}")
        if len(class_map) > 5:
            context_parts.append(f"# ... ({len(class_map)-5} more classes)")
        
        context_parts.append("# Sample property mappings:")
        for prop_uri, compact_ref in list(prop_map.items())[:5]:
            prop_name = properties[prop_uri].name
            context_parts.append(f"# {compact_ref} = {prop_name}")
        if len(prop_map) > 5:
            context_parts.append(f"# ... ({len(prop_map)-5} more properties)")
        
        return "\n".join(context_parts)
    
    def _get_namespace_from_uri(self, uri: str, namespaces: Dict[str, str]) -> str:
        """Get the namespace prefix for a URI, or 'other' if not found."""
        for prefix, namespace_uri in namespaces.items():
            if uri.startswith(namespace_uri):
                return prefix
        return "other"
    
    def _create_empty_ontology_info(self) -> OntologyInfo:
        """Create empty ontology info for error cases."""
        return OntologyInfo(
            classes={},
            properties={},
            namespaces={},
            domain_summary="No ontology information available - discovery failed or no ontologies found.",
            llm_context="No ontology context available for SPARQL generation.",
            total_instances=0
        )
    
    def _create_compact_uri_mappings(
        self, 
        classes: Dict[str, ClassInfo], 
        properties: Dict[str, PropertyInfo]
    ) -> Tuple[Dict[str, str], Dict[str, str]]:
        """
        Create compact URI mappings (c:0, c:1, p:0, p:1 style) for classes and properties.
        
        Returns:
            Tuple of (class_map, property_map) where keys are full URIs and values are compact refs
        """
        class_map = {}
        prop_map = {}
        
        # Sort by name for consistent ordering
        sorted_classes = sorted(classes.items(), key=lambda x: x[1].name)
        sorted_props = sorted(properties.items(), key=lambda x: x[1].name)
        
        # Create compact class references
        for i, (class_uri, class_info) in enumerate(sorted_classes):
            class_map[class_uri] = f"c:{i}"
        
        # Create compact property references  
        for i, (prop_uri, prop_info) in enumerate(sorted_props):
            prop_map[prop_uri] = f"p:{i}"
            
        return class_map, prop_map
    
    def _group_by_ontology_source(
        self, 
        classes: Dict[str, ClassInfo], 
        properties: Dict[str, PropertyInfo]
    ) -> Dict[str, Tuple[List[str], List[str]]]:
        """
        Group classes and properties by their ontology source/base URI.
        
        Returns:
            Dict mapping ontology base names to (class_uris, property_uris) tuples
        """
        ontology_groups = {}
        
        # Group classes
        for class_uri, class_info in classes.items():
            ontology_base = self._extract_ontology_base(class_uri)
            if ontology_base not in ontology_groups:
                ontology_groups[ontology_base] = ([], [])
            ontology_groups[ontology_base][0].append(class_uri)
        
        # Group properties
        for prop_uri, prop_info in properties.items():
            ontology_base = self._extract_ontology_base(prop_uri)
            if ontology_base not in ontology_groups:
                ontology_groups[ontology_base] = ([], [])
            ontology_groups[ontology_base][1].append(prop_uri)
            
        return ontology_groups
    
    def _extract_ontology_base(self, uri: str) -> str:
        """
        Extract a meaningful ontology base name from a URI.
        
        Examples:
        - http://siemens.com/ontologies/Recipes#Material -> Recipes
        - http://siemens.com/ontologies/ELN#... -> ELN
        """
        if "siemens.com/ontologies/" in uri:
            # Extract the ontology name between /ontologies/ and #
            parts = uri.split("/ontologies/")[1].split("#")[0]
            return parts
        elif "#" in uri:
            # Generic case - use the part before #
            return uri.split("#")[0].split("/")[-1]
        else:
            # Fallback - use domain
            return uri.split("/")[2] if "/" in uri else "unknown"
    
    def _extract_domain_keywords(
        self, 
        classes: Dict[str, ClassInfo], 
        properties: Dict[str, PropertyInfo]
    ) -> List[str]:
        """
        Extract key domain terms from class and property names for compact domain summary.
        """
        all_names = []
        
        # Collect class names
        for class_info in classes.values():
            all_names.append(class_info.name)
        
        # Collect property names (sample to avoid overwhelming)
        prop_names = [prop_info.name for prop_info in properties.values()]
        all_names.extend(prop_names[:20])  # Sample first 20 properties
        
        # Extract meaningful terms (filter out common words)
        terms = set()
        common_words = {'the', 'and', 'or', 'of', 'to', 'in', 'for', 'with', 'by', 'a', 'an'}
        
        for name in all_names:
            # Split on spaces and common separators
            words = name.replace('_', ' ').replace('-', ' ').split()
            for word in words:
                if len(word) > 2 and word.lower() not in common_words:
                    terms.add(word.lower())
        
        # Return most relevant terms (alphabetically sorted for consistency)
        return sorted(list(terms))[:10]
    
    def _group_properties_by_domain_class(
        self, 
        ont_classes: List[str], 
        ont_properties: List[str], 
        all_properties: Dict[str, PropertyInfo],
        all_classes: Dict[str, ClassInfo] = None
    ) -> Dict[str, List[str]]:
        """
        Group properties by their domain classes, supporting cross-ontology relationships.
        
        This method now supports linking properties that have domains in different ontologies.
        For example, a RecipeLinking property with domain in Email ontology will be grouped
        under the Email domain class when processing the Email ontology.
        
        Args:
            ont_classes: Classes in the current ontology being processed
            ont_properties: Properties in the current ontology being processed  
            all_properties: All properties from all ontologies
            all_classes: All classes from all ontologies (for cross-ontology lookup)
        
        Returns:
            Dict mapping class URIs to lists of property URIs that have that class as domain
        """
        class_properties = {}
        
        # Process properties from the current ontology
        for prop_uri in ont_properties:
            prop_info = all_properties[prop_uri]
            
            # Find which classes this property has as domain
            for domain_uri in prop_info.domain_classes:
                if domain_uri in ont_classes:
                    if domain_uri not in class_properties:
                        class_properties[domain_uri] = []
                    class_properties[domain_uri].append(prop_uri)
        
        # Also find cross-ontology properties that have domains in this ontology
        if all_classes is not None:
            for prop_uri, prop_info in all_properties.items():
                # Skip properties already processed from current ontology
                if prop_uri in ont_properties:
                    continue
                    
                # Check if this cross-ontology property has a domain in current ontology
                for domain_uri in prop_info.domain_classes:
                    if domain_uri in ont_classes:
                        if domain_uri not in class_properties:
                            class_properties[domain_uri] = []
                        class_properties[domain_uri].append(prop_uri)
        
        return class_properties
    
    def _infer_property_type(self, prop_info: PropertyInfo, range_uri: str = None) -> str:
        """
        Infer the property type/range from property info or range URI.
        
        Returns types like 'str', 'int', 'date', 'number', etc.
        """
        if range_uri:
            # Try to infer from range URI
            range_lower = range_uri.lower()
            if 'string' in range_lower or 'literal' in range_lower:
                return 'str'
            elif 'integer' in range_lower or 'int' in range_lower:
                return 'int'
            elif 'decimal' in range_lower or 'float' in range_lower or 'double' in range_lower:
                return 'number'
            elif 'date' in range_lower:
                return 'date'
            elif 'time' in range_lower:
                return 'time'
            elif 'boolean' in range_lower:
                return 'bool'
        
        # Infer from property name
        prop_name_lower = prop_info.name.lower()
        
        # Date/time patterns
        if any(word in prop_name_lower for word in ['date', 'time', 'timestamp', 'created', 'updated']):
            if 'time' in prop_name_lower and 'date' not in prop_name_lower:
                return 'time'
            return 'date'
        
        # Numeric patterns
        if any(word in prop_name_lower for word in ['id', 'count', 'number', 'quantity', 'seq', 'revision']):
            return 'int'
        
        if any(word in prop_name_lower for word in ['mass', 'weight', 'amount', 'composition', 'percentage']):
            return 'number'
        
        # Default to string
        return 'str'
