"""
SPARQL Query Engine - Generates SPARQL queries from natural language questions

This module handles the core intelligence of converting natural language questions
into executable SPARQL queries using the discovered ontology context and LLM reasoning.

Key functionality:
1. Natural language question analysis
2. SPARQL query generation using ontology context
3. Query optimization and validation
4. Iterative improvement based on execution results
5. Clear error reporting when query generation fails

Note: The system does not generate fallback queries to avoid misleading results.
When query generation fails, clear error information is provided instead.
"""

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

# 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)

from datetime import datetime

from .models import (
    GraphmartConfig,
    OntologyInfo,
    ClassInfo,
    PropertyInfo,
    SPARQLQuery,
    QueryResult,
    AgentMemory
)


class SPARQLQueryEngine:
    """
    Core engine for generating SPARQL queries from natural language questions.
    
    Uses LLM reasoning combined with ontology context to generate, validate,
    and iteratively improve SPARQL queries for AGS graphmarts.
    """
    
    def __init__(self, anzo_client, openai_client, config: GraphmartConfig):
        """Initialize the query engine."""
        self.anzo_client = anzo_client
        self.openai_client = openai_client
        self.config = config
        
        # Query generation settings
        self.max_iterations = 3
        self.query_timeout = 30
        self.confidence_threshold = 0.6
        
        # Prompt management (templates disabled to avoid bias)
        # self.query_templates = self._load_query_templates()  # Disabled to avoid template bias
        self.system_prompts = self._load_system_prompts()
    
    def _load_query_templates(self) -> Dict[str, str]:
        """Load SPARQL query templates for common patterns.
        
        DISABLED: Templates are disabled to avoid bias in LLM query generation.
        The LLM should generate queries purely based on ontology context.
        """
        # Templates disabled to avoid bias - return empty dict
        return {}
    
    def _load_system_prompts(self) -> Dict[str, str]:
        """Load system prompts for LLM interactions."""
        return {
            "sparql_generation": """
You are an expert SPARQL query generator for Altair Graph Studio (AGS) knowledge graphs.

Your task is to convert natural language questions into precise SPARQL queries using the provided ontology context.

ONTOLOGY CONTEXT:
{ontology_context}

GRAPHMART URI: {graphmart_uri}

IMPORTANT RULES:
1. DO NOT use GRAPH clauses - the graphmart URI is already the query context via pyanzo
2. Use only classes and properties from the provided ontology context
3. Always include appropriate LIMIT clauses (default 20)
4. Use OPTIONAL for properties that might not exist
5. Include rdfs:label when possible for readable results
6. Use DISTINCT to avoid duplicates
7. For counting queries, use COUNT(DISTINCT ?variable)

COMMON PATTERNS:
- Finding instances: ?x rdf:type <ClassURI>
- Getting labels: OPTIONAL {{ ?x rdfs:label ?label }}
- Property relationships: ?subject <PropertyURI> ?object
- Text search: FILTER(CONTAINS(LCASE(STR(?object)), "text"))

Generate a valid SPARQL query that answers the question accurately.
""",

            "query_analysis": """
You are analyzing a SPARQL query execution result to understand if it successfully answered the user's question.

ORIGINAL QUESTION: {question}
SPARQL QUERY: {sparql_query}
EXECUTION RESULT: {execution_result}
ERROR (if any): {error_message}

Analyze whether:
1. The query executed successfully
2. The results answer the original question
3. The query can be improved

Provide:
- success: true/false
- answer_quality: "excellent" / "good" / "partial" / "poor"
- suggestions: list of specific improvements
- natural_answer: human-readable answer based on results
""",

            "query_improvement": """
You are a SPARQL expert fixing a broken query. You must analyze the error and generate a DIFFERENT, corrected structured query.

ORIGINAL QUESTION: {question}

FAILED QUERY:
{previous_query}

ERROR MESSAGE: {error_message}

ONTOLOGY CONTEXT:
{ontology_context}

ANALYSIS INSTRUCTIONS:
1. Read the error message carefully to understand the specific problem
2. Examine the failed query to identify the syntax or logic issue
3. IMPORTANT: Generate a DIFFERENT query than the previous attempt - do NOT return the same query

Common fixes for SPARQL errors:
   - "variable X is projected but not used" → Use COUNT(*) or define the variable in all UNION branches
   - "No results returned" → Try these fixes:
     • Make some relationships OPTIONAL: OPTIONAL {{ ?x property ?y }}
     • Simplify complex joins - start with just one class/property combination
     • Remove or relax FILTER conditions that might be too restrictive
     • Try alternative property paths or class relationships
     • Use broader patterns instead of specific ones
     • Check if the data actually contains the expected relationships
   - "Type mismatch" → Verify property domain/range compatibility

For counting across UNION branches, use these patterns:
- COUNT(*) to count all pattern matches
- COUNT(?variable) where the variable exists in all branches
- Use a common variable name across all UNION branches

SPECIFIC STRATEGIES for "No results returned":
1. If the query tries to join multiple classes, simplify to just one class first
2. If using specific property values in FILTER, make them OPTIONAL or remove filters
3. If expecting exact relationships, try making some relationships OPTIONAL
4. If using complex patterns, break down into simpler patterns

Generate a corrected SPARQL query that:
1. Is DIFFERENT from the previous failed query
2. Fixes the specific error mentioned above
3. Answers the original question correctly
4. Uses proper SPARQL syntax with compact references (c:X, p:X)
5. Includes appropriate ontology classes and properties

CRITICAL: Your response must be a DIFFERENT query than the previous attempt. Do not return identical queries.
"""
        }
    
    async def generate_sparql_query(
        self, 
        question: str, 
        ontology_info: OntologyInfo,
        agent_memory: Optional[AgentMemory] = None
    ) -> SPARQLQuery:
        """
        Generate a SPARQL query from a natural language question.
        
        Args:
            question: Natural language question
            ontology_info: Discovered ontology context
            agent_memory: Previous successful patterns (optional)
            
        Returns:
            SPARQLQuery object with generated query and metadata
        """
        debug_log(f"🔍 Generating SPARQL for: {question}")
        
        try:
            # Step 1: Analyze question and build context
            ontology_context = self._build_ontology_context(ontology_info, question)
            
            # Step 2: Check for similar patterns in memory
            similar_patterns = self._find_similar_patterns(question, agent_memory) if agent_memory else []
            
            # Step 3: Generate complete SPARQL query with metadata
            sparql_query = await self._generate_initial_query(
                question, ontology_context, similar_patterns
            )
            
            # Step 4: Validate and optimize query text
            validated_compact_query = await self._validate_query_syntax(sparql_query.latest_compact_query)
            
            # Step 5: Expand compact URI references to full URIs
            expanded_query_text = self._expand_compact_uris_in_sparql(validated_compact_query, ontology_info)
            
            # Update the query object with both compact and expanded versions
            sparql_query.latest_compact_query = validated_compact_query
            sparql_query.latest_expanded_query = expanded_query_text
            
            return sparql_query
            
        except Exception as e:
            debug_log(f"❌ Query generation failed: {e}")
            # Return clear error information instead of attempting to generate queries
            error_details = str(e)
            return SPARQLQuery(
                latest_compact_query=f"# ERROR: SPARQL query generation failed\n# Reason: {error_details}\n# No query could be generated for this question.",
                latest_expanded_query=f"# ERROR: SPARQL query generation failed\n# Reason: {error_details}\n# No query could be generated for this question.",
                original_question=question,
                confidence=0.0,
                ontology_classes_used=[],
                ontology_properties_used=[],
                query_type="error",
                explanation=f"Complete query generation failure: {error_details}. This indicates a system-level issue that prevented any query generation attempt."
            )
    
    def _build_ontology_context(self, ontology_info: OntologyInfo, question: str) -> str:
        """
        Use the pre-generated comprehensive LLM-friendly ontology context.
        
        This provides ALL classes and properties to the LLM, letting it decide
        what's relevant rather than pre-filtering based on simple keyword matching.
        """
        # Use the pre-generated comprehensive context
        if ontology_info.llm_context:
            return ontology_info.llm_context
        
        # Fallback: generate basic context if llm_context is missing
        return self._generate_fallback_context(ontology_info)
    
    def _generate_fallback_context(self, ontology_info: OntologyInfo) -> str:
        """Generate basic ontology context when pre-generated context is unavailable."""
        context_parts = []
        
        context_parts.append("DOMAIN SUMMARY:")
        context_parts.append(ontology_info.domain_summary or "No domain summary available")
        context_parts.append("")
        
        if ontology_info.classes:
            context_parts.append(f"AVAILABLE CLASSES ({len(ontology_info.classes)} total):")
            for class_uri, class_info in list(ontology_info.classes.items())[:20]:  # First 20
                context_parts.append(f"- {class_info.name}: <{class_uri}>")
            if len(ontology_info.classes) > 20:
                context_parts.append(f"... and {len(ontology_info.classes) - 20} more classes")
            context_parts.append("")
        
        if ontology_info.properties:
            context_parts.append(f"AVAILABLE PROPERTIES ({len(ontology_info.properties)} total):")
            for prop_uri, prop_info in list(ontology_info.properties.items())[:25]:  # First 25
                context_parts.append(f"- {prop_info.name}: <{prop_uri}>")
            if len(ontology_info.properties) > 25:
                context_parts.append(f"... and {len(ontology_info.properties) - 25} more properties")
            context_parts.append("")
        
        if ontology_info.namespaces:
            context_parts.append("NAMESPACES:")
            for prefix, uri in ontology_info.namespaces.items():
                context_parts.append(f"- {prefix}: <{uri}>")
        
        return "\n".join(context_parts)
    
    def _find_similar_patterns(self, question: str, agent_memory: AgentMemory) -> List[Dict]:
        """Find similar successful query patterns from memory."""
        if not agent_memory or not agent_memory.successful_queries:
            return []
        
        similar_patterns = []
        question_lower = question.lower()
        
        for memory_item in agent_memory.successful_queries:
            # Simple similarity based on shared words
            memory_question = memory_item.get('question', '').lower()
            shared_words = set(question_lower.split()) & set(memory_question.split())
            
            if len(shared_words) >= 2:  # At least 2 shared words
                similar_patterns.append({
                    'question': memory_item.get('question'),
                    'sparql': memory_item.get('sparql'),
                    'similarity': len(shared_words)
                })
        
        # Sort by similarity and return top 3
        similar_patterns.sort(key=lambda x: x['similarity'], reverse=True)
        return similar_patterns[:3]
    
    async def _generate_initial_query(
        self, 
        question: str, 
        ontology_context: str, 
        similar_patterns: List[Dict]
    ) -> SPARQLQuery:
        """Generate initial SPARQL query using LLM with structured response."""
        
        # Build enhanced system prompt that instructs the LLM to generate structured output
        system_prompt = """You are an expert SPARQL query generator for knowledge graphs. You will receive an ontology context and a natural language question, and must generate a complete SPARQL query with metadata.

Your response must include:
1. latest_compact_query: A valid SPARQL query using compact references (c:X, p:X)
2. latest_expanded_query: Will be populated automatically - leave as empty string
3. original_question: The exact question asked
4. confidence: Your confidence level (0.0-1.0) in the query's correctness
5. ontology_classes_used: List of class compact references (c:0, c:1, etc.) used in the query
6. ontology_properties_used: List of property compact references (p:0, p:1, etc.) used in the query  
7. query_type: Type of query (count, select, ask, construct, describe)
8. explanation: Brief explanation of what the query does and why it was constructed this way

IMPORTANT SPARQL Guidelines:
- CRITICAL: In your SPARQL query text, you MUST use ONLY the compact URI references (c:0, c:1, p:0, p:1, etc.) from the ontology context
- NEVER write full URIs like <http://...> in the query text - always use the compact references like c:0, p:0
- Example: Write "?product a c:6 ; p:17 ?name" NOT "?product a <http://siemens.com/ontologies/Recipes#Product> ; <http://siemens.com/ontologies/Recipes#name> ?name"
- The compact references will be automatically expanded to full URIs during processing
- DO NOT use GRAPH clauses - when querying graphmarts via pyanzo, the graphmart URI is already the context
- Include proper SPARQL structure with SELECT, WHERE clauses (but NO GRAPH clause inside WHERE)
- Add FILTER clauses for text searches when appropriate
- Use standard SPARQL patterns: ?variable a c:X ; p:Y ?value .
- Prioritize matches over broad patterns
- Always use fuzzy matching techniques when looking for matches (e.g., FILTER(CONTAINS(LCASE(STR(?object)), "text")))
- Be conservative with confidence scores unless very certain
- Before writing the query, write out the triple patterns and relationships to make sure you get this correct, but do not return this thinking."""

        # Build context-specific prompt
        user_prompt = f"""Generate a SPARQL query for the following question using the provided ontology context.

**Question:** {question}

**Ontology Context:**
{ontology_context}

**Graphmart URI:** {self.config.graphmart_uri}

**IMPORTANT: Your SPARQL query MUST use compact references like this example:**
CORRECT: SELECT ?name WHERE {{ ?product a c:6 ; p:17 ?name . }}
WRONG: SELECT ?name WHERE {{ ?product a <http://siemens.com/ontologies/Recipes#Product> ; <http://siemens.com/ontologies/Recipes#name> ?name . }}

Use ONLY the compact references (c:X, p:X) from the ontology context above.
"""
        
        # Add similar patterns if available
        if similar_patterns:
            user_prompt += "\n**Similar Successful Patterns:**\n"
            for pattern in similar_patterns:
                user_prompt += f"Q: {pattern['question']}\nSPARQL: {pattern['sparql']}\n\n"
        
        user_prompt += "\nGenerate the SPARQL query with complete metadata now."
        
        # Generate structured query
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        try:
            # Use instructor for structured response (synchronous call)
            sparql_response = self._call_openai_structured_sync(messages, SPARQLQuery)
            debug_log(f"✅ LLM generated structured SPARQL with confidence: {sparql_response.confidence}")
            return sparql_response
            
        except Exception as e:
            debug_log(f"❌ LLM query generation failed: {e}")
            # Return clear error instead of misleading fallback query
            error_details = str(e)
            return SPARQLQuery(
                latest_compact_query=f"# ERROR: Unable to generate SPARQL query\n# Reason: {error_details}\n# Please check your question and try again.",
                latest_expanded_query=f"# ERROR: Unable to generate SPARQL query\n# Reason: {error_details}\n# Please check your question and try again.",
                original_question=question,
                confidence=0.0,  # Zero confidence for errors
                ontology_classes_used=[],
                ontology_properties_used=[],
                query_type="error",
                explanation=f"Query generation failed: {error_details}. The system was unable to convert your question into a valid SPARQL query. Please rephrase your question or check that it relates to the available ontology."
            )
    
    async def _call_openai(self, messages: List[Dict]) -> str:
        """Call OpenAI API with error handling (legacy string response)."""
        try:
            response = await self.openai_client.chat.completions.create(
                model="gpt-4o",
                messages=messages,
                temperature=0.1,
                max_tokens=1000
            )
            return response.choices[0].message.content
        except Exception as e:
            raise Exception(f"OpenAI API call failed: {e}")
    
    async def _call_openai_structured(self, messages: List[Dict], response_model) -> Any:
        """Call OpenAI API with instructor for structured responses."""
        try:
            # Instructor calls are synchronous, not async
            response = self.openai_client.chat.completions.create(
                model="gpt-4o",
                response_model=response_model,
                messages=messages,
                temperature=0.1,
                max_tokens=2000
            )
            return response
        except Exception as e:
            raise Exception(f"OpenAI API call failed: {e}")
    
    def _call_openai_structured_sync(self, messages: List[Dict], response_model) -> Any:
        """Call OpenAI API with instructor for structured responses (synchronous)."""
        try:
            response = self.openai_client.chat.completions.create(
                model="gpt-4o",
                response_model=response_model,
                messages=messages,
                temperature=0.1,
                max_tokens=2000
            )
            return response
        except Exception as e:
            raise Exception(f"OpenAI API call failed: {e}")
    
    def _extract_sparql_from_response(self, response: str) -> str:
        """Extract SPARQL query from LLM response."""
        # Look for SPARQL query in code blocks or after keywords
        patterns = [
            r'```sparql\s*(.*?)\s*```',
            r'```\s*(SELECT.*?)\s*```',
            r'(SELECT.*?)(?:\n\n|\Z)',
        ]
        
        for pattern in patterns:
            match = re.search(pattern, response, re.DOTALL | re.IGNORECASE)
            if match:
                return match.group(1).strip()
        
        # If no pattern matches, return the whole response
        return response.strip()
    
    async def _validate_query_syntax(self, sparql_query: str) -> str:
        """Validate SPARQL query syntax and fix common issues."""
        # Basic syntax validation and cleanup
        query = sparql_query.strip()
        
        # DO NOT add GRAPH clauses - pyanzo handles the graphmart context automatically
        # The previous code that added GRAPH clauses has been removed to follow our design
        
        # Ensure LIMIT clause exists for SELECT queries
        if query.upper().startswith('SELECT') and 'LIMIT' not in query.upper():
            query += '\nLIMIT 20'
        
        return query
    
    def _extract_classes_from_query(self, sparql_query: str) -> List[str]:
        """Extract ontology classes used in the query."""
        # Look for rdf:type patterns
        patterns = [
            r'\?[^}\s]+\s+rdf:type\s+<([^>]+)>',
            r'\?[^}\s]+\s+a\s+<([^>]+)>',
        ]
        
        classes = []
        for pattern in patterns:
            matches = re.findall(pattern, sparql_query, re.IGNORECASE)
            classes.extend(matches)
        
        return list(set(classes))
    
    def _extract_properties_from_query(self, sparql_query: str) -> List[str]:
        """Extract ontology properties used in the query."""
        # Look for property patterns (excluding rdf:type and common predicates)
        pattern = r'<([^>]+)>'
        matches = re.findall(pattern, sparql_query)
        
        # Filter out common RDF/RDFS/OWL predicates and classes
        common_predicates = {
            'http://www.w3.org/1999/02/22-rdf-syntax-ns#type',
            'http://www.w3.org/2000/01/rdf-schema#label',
            'http://www.w3.org/2000/01/rdf-schema#comment',
        }
        
        properties = []
        for match in matches:
            if match not in common_predicates and 'Class' not in match:
                properties.append(match)
        
        return list(set(properties))
    
    def _classify_query_type(self, sparql_query: str) -> str:
        """Classify the type of SPARQL query."""
        query_upper = sparql_query.upper()
        
        if query_upper.startswith('SELECT'):
            if 'COUNT' in query_upper:
                return 'count'
            elif 'DISTINCT' in query_upper:
                return 'select_distinct'
            else:
                return 'select'
        elif query_upper.startswith('ASK'):
            return 'ask'
        elif query_upper.startswith('CONSTRUCT'):
            return 'construct'
        elif query_upper.startswith('DESCRIBE'):
            return 'describe'
        else:
            return 'unknown'
    
    async def execute_sparql_query(self, sparql_query: SPARQLQuery) -> QueryResult:
        """
        Execute a SPARQL query and return structured results.
        
        Args:
            sparql_query: SPARQLQuery object to execute
            
        Returns:
            QueryResult with execution results and metadata
        """
        debug_log(f"🔄 Executing SPARQL query...")
        
        start_time = datetime.now()
        timeout_seconds = 30  # Hard-coded 30 second timeout
        
        try:
            # Execute query against graphmart data with timeout
            result = await asyncio.wait_for(
                asyncio.get_event_loop().run_in_executor(
                    None,
                    lambda: self.anzo_client.query_graphmart(
                        self.config.graphmart_uri, 
                        sparql_query.latest_expanded_query
                    )
                ),
                timeout=timeout_seconds
            )
            
            table_results = result.as_table_results().as_list()
            
            # Extract column names from the SPARQL query
            column_names = self._extract_column_names(sparql_query.latest_expanded_query)
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            return QueryResult(
                success=True,
                data=table_results,
                column_names=column_names,
                sparql_query=sparql_query.latest_expanded_query,
                execution_time=execution_time,
                error_message=None
            )
            
        except asyncio.TimeoutError:
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # Extract column names even for timed out queries
            column_names = self._extract_column_names(sparql_query.latest_expanded_query)
            
            debug_log(f"⏰ Query timed out after {timeout_seconds} seconds")
            
            return QueryResult(
                success=False,
                data=[],
                column_names=column_names,
                sparql_query=sparql_query.latest_expanded_query,
                execution_time=execution_time,
                error_message=f"Query timed out after {timeout_seconds} seconds. The query may be too complex or the dataset too large. Try simplifying the query or adding more specific filters."
            )
            
        except Exception as e:
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # Extract column names even for failed queries
            column_names = self._extract_column_names(sparql_query.latest_expanded_query)
            
            return QueryResult(
                success=False,
                data=[],
                column_names=column_names,
                sparql_query=sparql_query.latest_expanded_query,
                execution_time=execution_time,
                error_message=str(e)
            )
    
    def _extract_column_names(self, sparql_query: str) -> List[str]:
        """Extract column names from SELECT query."""
        # Simple extraction of SELECT variables
        select_match = re.search(r'SELECT\s+(.*?)\s+WHERE', sparql_query, re.IGNORECASE | re.DOTALL)
        if not select_match:
            return []
        
        select_clause = select_match.group(1).strip()
        
        # Handle COUNT queries specially
        if re.search(r'COUNT\s*\(', select_clause, re.IGNORECASE):
            return ['count']
        
        # Extract regular variables
        variables = re.findall(r'\?(\w+)', select_clause)
        return variables
    
    async def improve_query_based_on_results(
        self, 
        original_query: SPARQLQuery,
        query_result: QueryResult,
        ontology_info: OntologyInfo
    ) -> Optional[SPARQLQuery]:
        """
        Improve a query based on execution results and errors.
        
        Args:
            original_query: The original query that was executed
            query_result: Results from executing the query
            ontology_info: Current ontology context
            
        Returns:
            Improved SPARQLQuery or None if no improvement possible
        """
        if query_result.success and len(query_result.data) > 0:
            # Query was successful, no improvement needed
            return None
        
        debug_log(f"🔧 Attempting to improve query...")
        
        try:
            # Build improvement prompt with enhanced error context
            ontology_context = self._build_ontology_context(ontology_info, original_query.original_question)
            
            prompt = self.system_prompts["query_improvement"].format(
                question=original_query.original_question,
                previous_query=original_query.latest_compact_query,  # ← Use compact query for improvement
                error_message=query_result.error_message or "No results returned",
                ontology_context=ontology_context
            )
            
            messages = [
                {"role": "system", "content": prompt},
                {"role": "user", "content": f"""The previous query failed with: {query_result.error_message or 'No results'}.

CRITICAL REQUIREMENT: Generate a DIFFERENT, improved SPARQL query that will work. 

Previous query was:
{original_query.latest_compact_query}

If the error is "No results returned", try these specific changes:
1. Make some relationships OPTIONAL (wrap in OPTIONAL {{ ... }})
2. Remove or simplify complex FILTER conditions  
3. Try different property combinations from the ontology
4. Simplify joins - use fewer classes in the pattern
5. Remove GROUP BY/ORDER BY and just do basic SELECT first

Your new query MUST be syntactically different from the previous attempt."""}
            ]
            
            # Use structured API call (same as initial query generation)
            improved_response = self._call_openai_structured_sync(messages, SPARQLQuery)
            
            # Normalize queries for comparison (remove extra whitespace)
            original_normalized = ' '.join(original_query.latest_compact_query.split())
            improved_normalized = ' '.join(improved_response.latest_compact_query.split())
            
            # Check if the improved query is actually different
            if improved_normalized == original_normalized:
                debug_log(f"⚠️ LLM returned identical query (normalized) - skipping improvement")
                debug_log(f"   Original: {original_normalized[:100]}...")
                debug_log(f"   Improved: {improved_normalized[:100]}...")
                return None
            
            debug_log(f"✅ LLM generated different query:")
            debug_log(f"   Original: {original_normalized[:100]}...")
            debug_log(f"   Improved: {improved_normalized[:100]}...")
            
            # Expand compact URIs in the improved query
            expanded_query = self._expand_compact_uris_in_sparql(
                improved_response.latest_compact_query, ontology_info
            )
            
            # Add this improvement iteration to the original query's history
            original_query.add_iteration(
                compact_query=improved_response.latest_compact_query,
                expanded_query=expanded_query,
                reason=f"Fixed: {query_result.error_message or 'No results returned'}"
            )
            
            # Lower confidence for improved queries
            original_query.confidence = max(0.3, original_query.confidence - 0.2)
            
            debug_log(f"✅ Query improved with confidence {original_query.confidence:.2f}")
            return original_query  # ← Return the updated original query, not a new one
            
        except Exception as e:
            debug_log(f"❌ Query improvement failed: {e}")
            return None
    
    def _expand_compact_uris_in_sparql(self, sparql_query: str, ontology_info: OntologyInfo) -> str:
        """
        Expand compact URI references (c:0, p:0) in SPARQL query to full URIs.
        
        Args:
            sparql_query: SPARQL query with compact URI references
            ontology_info: Ontology info containing compact URI mappings
            
        Returns:
            SPARQL query with full URIs
        """
        if not sparql_query or not ontology_info:
            return sparql_query
            
        expanded_query = sparql_query
        
        try:
            # Expand class references (c:0, c:1, etc.)
            if ontology_info.compact_class_map:
                # Create reverse mapping: compact_ref -> full_uri
                class_reverse_map = {v: k for k, v in ontology_info.compact_class_map.items()}
                
                # Sort by compact reference length (longest first) to prevent partial matches
                # This ensures c:31 is processed before c:3, avoiding malformed URIs
                sorted_class_items = sorted(class_reverse_map.items(), key=lambda x: len(x[0]), reverse=True)
                
                for compact_ref, full_uri in sorted_class_items:
                    # Replace compact references with full URIs wrapped in angle brackets
                    expanded_query = expanded_query.replace(compact_ref, f"<{full_uri}>")
            
            # Expand property references (p:0, p:1, etc.)
            if ontology_info.compact_property_map:
                # Create reverse mapping: compact_ref -> full_uri
                prop_reverse_map = {v: k for k, v in ontology_info.compact_property_map.items()}
                
                # Sort by compact reference length (longest first) to prevent partial matches
                # This ensures p:123 is processed before p:12, avoiding malformed URIs
                sorted_prop_items = sorted(prop_reverse_map.items(), key=lambda x: len(x[0]), reverse=True)
                
                for compact_ref, full_uri in sorted_prop_items:
                    # Replace compact references with full URIs wrapped in angle brackets
                    expanded_query = expanded_query.replace(compact_ref, f"<{full_uri}>")
                    
            debug_log(f"✅ Expanded {len(ontology_info.compact_class_map or {})} class refs and {len(ontology_info.compact_property_map or {})} property refs")
            return expanded_query
            
        except Exception as e:
            debug_log(f"⚠️ Error expanding compact URIs: {e}")
            return sparql_query  # Return original query if expansion fails
