"""
Data models for the AGS SPARQL Agent.

These models define the core data structures used throughout the agent,
following the patterns established in the GraphRAG pipeline.
"""

import os
import json
from datetime import datetime
from typing import List, Dict, Optional, Any, Union
from dataclasses import dataclass, asdict
from pydantic import BaseModel, Field


# --- Configuration Models ---

@dataclass
class GraphmartConfig:
    """Configuration for connecting to an AGS graphmart."""
    ags_server: str
    ags_port: int
    graphmart_uri: str
    username: str
    password: str
    
    @classmethod
    def from_env(cls) -> 'GraphmartConfig':
        """Create configuration from environment variables."""
        graphmart_uri = os.getenv('GRAPHMART_URI')
        if not graphmart_uri:
            raise ValueError("GRAPHMART_URI environment variable is required")
            
        # Extract server and port from URI if possible
        # Format: http://server:port/graphmart/id
        try:
            from urllib.parse import urlparse
            parsed = urlparse(graphmart_uri)
            server = parsed.hostname or os.getenv('ANZO_SERVER', 'localhost')
            port = parsed.port or int(os.getenv('ANZO_PORT', '9091'))
        except:
            server = os.getenv('ANZO_SERVER', 'localhost')
            port = int(os.getenv('ANZO_PORT', '9091'))
        
        return cls(
            ags_server=server,
            ags_port=port,
            graphmart_uri=graphmart_uri,
            username=os.getenv('ANZO_USERNAME', os.getenv('USER', 'admin')),
            password=os.getenv('ANZO_PASSWORD', '')
        )


# --- Ontology Models ---

@dataclass
class ClassInfo:
    """Information about an ontology class."""
    name: str
    uri: str
    description: Optional[str] = None
    properties: List[str] = None
    instance_count: Optional[int] = None
    sample_instances: List[str] = None
    ontology_uri: Optional[str] = None  # NEW: Which ontology this class belongs to
    
    def __post_init__(self):
        if self.properties is None:
            self.properties = []
        if self.sample_instances is None:
            self.sample_instances = []
    
    def has_instances(self) -> bool:
        """Returns True if this class has instances in the knowledge graph."""
        return self.instance_count is not None and self.instance_count > 0


@dataclass  
class PropertyInfo:
    """Information about an ontology property."""
    name: str
    uri: str
    description: Optional[str] = None
    domain_classes: List[str] = None
    range_classes: List[str] = None
    usage_frequency: Optional[int] = None
    ontology_uri: Optional[str] = None  # Which ontology this property belongs to
    is_object_property: Optional[bool] = None  # True = owl:ObjectProperty, False = owl:DatatypeProperty, None = unknown
    
    def __post_init__(self):
        if self.domain_classes is None:
            self.domain_classes = []
        if self.range_classes is None:
            self.range_classes = []
    
    def has_usage(self) -> bool:
        """Returns True if this property is actually used in the knowledge graph."""
        return self.usage_frequency is not None and self.usage_frequency > 0


@dataclass
class OntologyMetadata:
    """Metadata about an individual ontology."""
    uri: str
    title: Optional[str] = None           # from rdfs:label or dc:title
    description: Optional[str] = None     # from rdfs:comment
    version: Optional[str] = None         # from owl:versionInfo
    creator: Optional[str] = None         # from dc:creator
    derived_title: Optional[str] = None   # smart fallback from namespace analysis
    
    def get_display_title(self) -> str:
        """Return the best available title for display."""
        if self.title:
            return self.title
        elif self.derived_title:
            return self.derived_title
        else:
            # Extract name from URI as last resort
            return self._extract_name_from_uri()
    
    def _extract_name_from_uri(self) -> str:
        """Extract a reasonable name from the ontology URI."""
        if '#' in self.uri:
            return self.uri.split('#')[-1]
        elif '/' in self.uri:
            parts = self.uri.rstrip('/').split('/')
            return parts[-1] if parts[-1] else parts[-2] if len(parts) > 1 else "Unknown"
        else:
            return "Unknown"


@dataclass
class OntologyExtractionResult:
    """Result of extracting schema information from a single ontology."""
    metadata: OntologyMetadata
    classes: Dict[str, ClassInfo]
    properties: Dict[str, PropertyInfo]


@dataclass
class OntologyInfo:
    """Complete ontology information for a graphmart."""
    classes: Dict[str, ClassInfo]
    properties: Dict[str, PropertyInfo]
    namespaces: Dict[str, str]
    domain_summary: str = ""
    total_instances: int = 0
    last_updated: str = ""
    # NEW: Pre-generated LLM-friendly context
    llm_context: str = ""
    # NEW: Compact URI mappings for efficient SPARQL generation
    compact_class_map: Dict[str, str] = None  # full_uri -> compact_ref (c:0, c:1, etc.)
    compact_property_map: Dict[str, str] = None  # full_uri -> compact_ref (p:0, p:1, etc.)
    reverse_class_map: Dict[str, str] = None  # compact_ref -> full_uri
    reverse_property_map: Dict[str, str] = None  # compact_ref -> full_uri
    
    def __post_init__(self):
        if not self.last_updated:
            self.last_updated = datetime.now().isoformat()
        if self.compact_class_map is None:
            self.compact_class_map = {}
        if self.compact_property_map is None:
            self.compact_property_map = {}
        if self.reverse_class_map is None:
            self.reverse_class_map = {}
        if self.reverse_property_map is None:
            self.reverse_property_map = {}
    
    def get_full_class_uri(self, compact_ref: str) -> Optional[str]:
        """Get full URI from compact class reference (e.g., c:1 -> full URI)."""
        return self.reverse_class_map.get(compact_ref)
    
    def get_full_property_uri(self, compact_ref: str) -> Optional[str]:
        """Get full URI from compact property reference (e.g., p:151 -> full URI)."""
        return self.reverse_property_map.get(compact_ref)
    
    def get_compact_class_ref(self, full_uri: str) -> Optional[str]:
        """Get compact reference from full class URI."""
        return self.compact_class_map.get(full_uri)
    
    def get_compact_property_ref(self, full_uri: str) -> Optional[str]:
        """Get compact reference from full property URI."""
        return self.compact_property_map.get(full_uri)
    
    def get_classes_with_instances(self) -> Dict[str, ClassInfo]:
        """Get only classes that have instances in the knowledge graph."""
        return {uri: cls_info for uri, cls_info in self.classes.items() 
                if cls_info.has_instances()}
    
    def get_properties_with_usage(self) -> Dict[str, PropertyInfo]:
        """Get only properties that are actually used in the knowledge graph."""
        return {uri: prop_info for uri, prop_info in self.properties.items() 
                if prop_info.has_usage()}
    
    def get_populated_summary(self) -> str:
        """Get a summary focusing only on classes and properties with data."""
        populated_classes = self.get_classes_with_instances()
        populated_properties = self.get_properties_with_usage()
        
        total_instances = sum(cls.instance_count or 0 for cls in populated_classes.values())
        total_property_usage = sum(prop.usage_frequency or 0 for prop in populated_properties.values())
        
        return (f"Knowledge graph with {len(populated_classes)} populated classes "
                f"({total_instances} total instances) and {len(populated_properties)} used properties "
                f"({total_property_usage} total property assertions). "
                f"Domain focus: {self.domain_summary}")


@dataclass
class OntologySummary:
    """High-level summary of a graphmart for discovery."""
    graphmart_uri: str
    domains: List[str]
    key_classes: List[ClassInfo]
    key_properties: List[PropertyInfo]
    sample_questions: List[str]
    total_triples: Optional[int] = None
    
    def __post_init__(self):
        if not self.domains:
            self.domains = []
        if not self.sample_questions:
            self.sample_questions = []


# --- Query Models (from GraphRAG pipeline) ---

class SPARQLQuery(BaseModel):
    """
    SPARQL query with explicit compact/expanded tracking and iteration history.
    
    The LLM always works with compact references (c:X, p:X).
    The execution engine always works with expanded URIs (<http://...>).
    """
    
    # Current/Latest versions (what gets used)
    latest_compact_query: str           # ← LLM-generated with c:X, p:X
    latest_expanded_query: str = ""     # ← Execution-ready with <http://...>
    
    # Query metadata  
    query_type: str = "select"
    confidence: float = 0.0
    explanation: str = ""
    
    # Ontology references used (compact format for clarity)
    ontology_classes_used: List[str] = Field(default_factory=list)     # ["c:3", "c:6"] 
    ontology_properties_used: List[str] = Field(default_factory=list)  # ["p:9", "p:17"]
    
    # Iteration history (for debugging and analysis)
    compact_query_history: List[str] = Field(default_factory=list)     # All compact versions tried
    expanded_query_history: List[str] = Field(default_factory=list)    # All expanded versions tried
    improvement_reasons: List[str] = Field(default_factory=list)       # Why each improvement was made
    
    # Context
    original_question: str = ""
    attempt_number: int = 1
    generated_at: str = Field(default_factory=lambda: datetime.now().isoformat())
    
    def add_iteration(self, compact_query: str, expanded_query: str, reason: str = ""):
        """Add a new iteration to the history and update latest versions."""
        # Store previous versions in history
        if self.latest_compact_query:
            self.compact_query_history.append(self.latest_compact_query)
        if self.latest_expanded_query:
            self.expanded_query_history.append(self.latest_expanded_query)
        if reason:
            self.improvement_reasons.append(reason)
        
        # Update latest versions
        self.latest_compact_query = compact_query
        self.latest_expanded_query = expanded_query
        self.attempt_number = len(self.compact_query_history) + 1
    
    @property
    def total_attempts(self) -> int:
        """Total number of query attempts (including current)."""
        return len(self.compact_query_history) + 1
    
    @property
    def first_compact_query(self) -> str:
        """Get the very first compact query attempted."""
        if self.compact_query_history:
            return self.compact_query_history[0]
        return self.latest_compact_query
    
    @property
    def first_expanded_query(self) -> str:
        """Get the very first expanded query attempted."""
        if self.expanded_query_history:
            return self.expanded_query_history[0]
        return self.latest_expanded_query


class QueryResult(BaseModel):
    """Result from executing a SPARQL query."""
    success: bool = Field(..., description="Whether the query executed successfully")
    data: List[Any] = Field(default_factory=list, description="Query result data - list of rows (each row is a list of values)")
    column_names: List[str] = Field(default_factory=list, description="Column names for the result data")
    error_message: Optional[str] = Field(None, description="Error message if query failed")
    sparql_query: str = Field(..., description="The SPARQL query that was executed")
    execution_time: Optional[float] = Field(None, description="Query execution time in seconds")


# --- Memory Models ---

@dataclass
class QueryPattern:
    """A successful query pattern for learning."""
    question_type: str
    sparql_template: str
    success_count: int = 1
    last_used: str = ""
    
    def __post_init__(self):
        if not self.last_used:
            self.last_used = datetime.now().isoformat()


@dataclass
class ErrorPattern:
    """A failed query pattern to avoid."""
    error_type: str
    sparql_pattern: str
    error_message: str
    occurrence_count: int = 1
    last_seen: str = ""
    
    def __post_init__(self):
        if not self.last_seen:
            self.last_seen = datetime.now().isoformat()


@dataclass
class QueryTemplate:
    """A reusable query template for a domain."""
    name: str
    description: str
    sparql_template: str
    parameters: List[str]
    domain: str = "general"
    usage_count: int = 0


@dataclass
class AgentMemory:
    """Persistent memory for the agent's learning."""
    successful_queries: List[QueryPattern]
    failed_patterns: List[ErrorPattern]
    domain_templates: List[QueryTemplate]
    popular_classes: Dict[str, int]
    popular_properties: Dict[str, int]
    last_updated: str = ""
    
    def __post_init__(self):
        if not self.successful_queries:
            self.successful_queries = []
        if not self.failed_patterns:
            self.failed_patterns = []
        if not self.domain_templates:
            self.domain_templates = []
        if not self.popular_classes:
            self.popular_classes = {}
        if not self.popular_properties:
            self.popular_properties = {}
        if not self.last_updated:
            self.last_updated = datetime.now().isoformat()
    
    def save_to_file(self, filepath: str):
        """Save memory to JSON file."""
        with open(filepath, 'w') as f:
            json.dump(asdict(self), f, indent=2)
    
    @classmethod
    def load_from_file(cls, filepath: str) -> 'AgentMemory':
        """Load memory from JSON file."""
        if not os.path.exists(filepath):
            return cls([], [], [], {}, {})
        
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        # Convert dictionaries back to dataclass instances
        successful_queries = [QueryPattern(**q) for q in data.get('successful_queries', [])]
        failed_patterns = [ErrorPattern(**e) for e in data.get('failed_patterns', [])]
        domain_templates = [QueryTemplate(**t) for t in data.get('domain_templates', [])]
        
        return cls(
            successful_queries=successful_queries,
            failed_patterns=failed_patterns,
            domain_templates=domain_templates,
            popular_classes=data.get('popular_classes', {}),
            popular_properties=data.get('popular_properties', {}),
            last_updated=data.get('last_updated', '')
        )


# --- MCP Response Models ---

class ExecutionDetails(BaseModel):
    """Detailed execution information for debugging/testing."""
    final_query_result: Optional['QueryResult'] = Field(None, description="Final query execution result")
    final_sparql_query: Optional['SPARQLQuery'] = Field(None, description="Query with complete iteration history")
    retry_attempts: int = Field(0, description="Number of retry attempts made")
    total_execution_time: float = Field(0.0, description="Total time including retries")


class QuestionResponse(BaseModel):
    """Response from the ask_question tool."""
    answer: str = Field(..., description="Natural language answer to the question")
    confidence: float = Field(..., description="Confidence score (0.0-1.0)")
    sources_used: int = Field(..., description="Number of data sources used")
    reasoning_steps: Optional[List[str]] = Field(None, description="Steps taken to find the answer")
    sparql_generated: Optional[str] = Field(None, description="Generated SPARQL query")
    execution_details: Optional[ExecutionDetails] = Field(None, description="Detailed execution information for testing")


class KnowledgeExploration(BaseModel):
    """Response from the explore_knowledge tool."""
    domain_summary: str = Field(..., description="High-level description of the knowledge graph")
    key_classes: List[Dict[str, Any]] = Field(..., description="Important classes with usage stats")
    key_properties: List[Dict[str, Any]] = Field(..., description="Important properties with usage stats")
    sample_questions: List[str] = Field(..., description="Example questions that can be asked")
    total_classes: int = Field(..., description="Total number of classes available")
    total_properties: int = Field(..., description="Total number of properties available")


class ExamplesResponse(BaseModel):
    """Response from the get_examples tool."""
    sample_questions: List[Dict[str, str]] = Field(..., description="Sample questions by complexity")
    data_examples: List[Dict[str, Any]] = Field(..., description="Representative data instances")
    query_patterns: List[Dict[str, str]] = Field(..., description="Common SPARQL patterns")


class AgentInsights(BaseModel):
    """Response from the get_agent_insights tool."""
    query_success_rate: float = Field(..., description="Overall query success rate")
    avg_iterations_per_query: float = Field(..., description="Average iterations needed per successful query")
    most_popular_classes: List[Dict[str, int]] = Field(..., description="Most frequently used classes")
    most_popular_properties: List[Dict[str, int]] = Field(..., description="Most frequently used properties")
    learning_trends: Dict[str, Any] = Field(..., description="Agent learning and improvement trends")
    memory_stats: Dict[str, int] = Field(..., description="Memory usage statistics")
