#!/usr/bin/env python3
"""
Ontology Cache Manager for AGS SPARQL Agent

Provides caching of ontology metadata to improve startup performance.
Cache is organized by graphmart URI in ~/.ags_cache/
"""

import json
import hashlib
import os
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, Any, Optional, Tuple
import logging

logger = logging.getLogger(__name__)

class OntologyCacheManager:
    """Manages caching of ontology metadata by graphmart URI."""
    
    def __init__(self, cache_ttl_hours: int = 24):
        """Initialize cache manager.
        
        Args:
            cache_ttl_hours: Cache time-to-live in hours (default 24)
        """
        self.cache_ttl = timedelta(hours=cache_ttl_hours)
        self.cache_root = Path.home() / ".ags_cache"
        self.cache_root.mkdir(exist_ok=True)
        
        # Create .gitignore in cache root
        gitignore_path = self.cache_root / ".gitignore"
        if not gitignore_path.exists():
            gitignore_path.write_text("# AGS Cache - automatically generated\n*\n")
    
    def _get_cache_dir(self, graphmart_uri: str) -> Path:
        """Get cache directory for a specific graphmart URI."""
        # Create a hash of the graphmart URI for the directory name
        uri_hash = hashlib.md5(graphmart_uri.encode()).hexdigest()[:16]
        cache_dir = self.cache_root / uri_hash
        cache_dir.mkdir(exist_ok=True)
        return cache_dir
    
    def _get_cache_paths(self, graphmart_uri: str) -> Tuple[Path, Path]:
        """Get cache file paths for a graphmart URI."""
        cache_dir = self._get_cache_dir(graphmart_uri)
        metadata_path = cache_dir / "ontology_metadata.json"
        info_path = cache_dir / "cache_info.json"
        return metadata_path, info_path
    
    def is_cache_valid(self, graphmart_uri: str) -> bool:
        """Check if cache exists and is still valid."""
        try:
            metadata_path, info_path = self._get_cache_paths(graphmart_uri)
            
            if not (metadata_path.exists() and info_path.exists()):
                return False
            
            # Check cache info
            with open(info_path, 'r') as f:
                cache_info = json.load(f)
            
            # Verify graphmart URI matches
            if cache_info.get('graphmart_uri') != graphmart_uri:
                return False
            
            # Check if cache is expired
            cache_time = datetime.fromisoformat(cache_info.get('timestamp', ''))
            if datetime.now() - cache_time > self.cache_ttl:
                logger.info(f"Cache expired for {graphmart_uri}")
                return False
            
            logger.info(f"Valid cache found for {graphmart_uri}")
            return True
            
        except Exception as e:
            logger.warning(f"Error checking cache validity: {e}")
            return False
    
    def load_cache(self, graphmart_uri: str) -> Optional[Dict[str, Any]]:
        """Load cached ontology metadata."""
        try:
            if not self.is_cache_valid(graphmart_uri):
                return None
            
            metadata_path, _ = self._get_cache_paths(graphmart_uri)
            
            with open(metadata_path, 'r') as f:
                cached_data = json.load(f)
            
            logger.info(f"Loaded ontology cache for {graphmart_uri}")
            return cached_data
            
        except Exception as e:
            logger.error(f"Error loading cache: {e}")
            return None
    
    def save_cache(self, graphmart_uri: str, ontology_data: Dict[str, Any]) -> bool:
        """Save ontology metadata to cache."""
        try:
            metadata_path, info_path = self._get_cache_paths(graphmart_uri)
            
            # Save metadata
            with open(metadata_path, 'w') as f:
                json.dump(ontology_data, f, indent=2, default=str)
            
            # Save cache info
            cache_info = {
                'graphmart_uri': graphmart_uri,
                'timestamp': datetime.now().isoformat(),
                'cache_version': '1.0',
                'total_classes': len(ontology_data.get('classes', {})),
                'total_properties': len(ontology_data.get('properties', {}))
            }
            
            with open(info_path, 'w') as f:
                json.dump(cache_info, f, indent=2)
            
            logger.info(f"Saved ontology cache for {graphmart_uri}")
            return True
            
        except Exception as e:
            logger.error(f"Error saving cache: {e}")
            return False
    
    def clear_cache(self, graphmart_uri: Optional[str] = None) -> bool:
        """Clear cache for specific graphmart URI or all caches."""
        try:
            if graphmart_uri:
                # Clear specific graphmart cache
                cache_dir = self._get_cache_dir(graphmart_uri)
                if cache_dir.exists():
                    import shutil
                    shutil.rmtree(cache_dir)
                    logger.info(f"Cleared cache for {graphmart_uri}")
            else:
                # Clear all caches
                if self.cache_root.exists():
                    import shutil
                    shutil.rmtree(self.cache_root)
                    self.cache_root.mkdir(exist_ok=True)
                    logger.info("Cleared all ontology caches")
            
            return True
            
        except Exception as e:
            logger.error(f"Error clearing cache: {e}")
            return False
    
    def get_cache_info(self, graphmart_uri: Optional[str] = None) -> Dict[str, Any]:
        """Get information about cached data."""
        try:
            if graphmart_uri:
                # Info for specific graphmart
                _, info_path = self._get_cache_paths(graphmart_uri)
                if info_path.exists():
                    with open(info_path, 'r') as f:
                        return json.load(f)
                else:
                    return {"status": "no_cache", "graphmart_uri": graphmart_uri}
            else:
                # Info for all caches
                cache_info = {"cache_root": str(self.cache_root), "graphmarts": []}
                
                if self.cache_root.exists():
                    for cache_dir in self.cache_root.iterdir():
                        if cache_dir.is_dir():
                            info_path = cache_dir / "cache_info.json"
                            if info_path.exists():
                                with open(info_path, 'r') as f:
                                    gm_info = json.load(f)
                                    gm_info["cache_dir"] = cache_dir.name
                                    cache_info["graphmarts"].append(gm_info)
                
                return cache_info
                
        except Exception as e:
            logger.error(f"Error getting cache info: {e}")
            return {"error": str(e)}

def serialize_ontology_info(ontology_info) -> Dict[str, Any]:
    """Serialize OntologyInfo object for caching."""
    serialized = {
        "classes": {},
        "properties": {},
        "namespaces": ontology_info.namespaces,
        "domain_summary": ontology_info.domain_summary,
        "total_instances": ontology_info.total_instances,
        "last_updated": ontology_info.last_updated,
        "llm_context": ontology_info.llm_context,
        "compact_class_map": ontology_info.compact_class_map,
        "compact_property_map": ontology_info.compact_property_map,
        "reverse_class_map": ontology_info.reverse_class_map,
        "reverse_property_map": ontology_info.reverse_property_map,
        "serialization_timestamp": datetime.now().isoformat()
    }
    
    # Serialize classes
    for uri, class_info in ontology_info.classes.items():
        serialized["classes"][uri] = {
            "uri": class_info.uri,
            "name": class_info.name,
            "description": class_info.description,
            "properties": class_info.properties,
            "instance_count": class_info.instance_count,
            "sample_instances": class_info.sample_instances,
            "ontology_uri": class_info.ontology_uri
        }
    
    # Serialize properties
    for uri, prop_info in ontology_info.properties.items():
        serialized["properties"][uri] = {
            "uri": prop_info.uri,
            "name": prop_info.name,
            "description": prop_info.description,
            "domain_classes": list(prop_info.domain_classes),
            "range_classes": list(prop_info.range_classes),
            "usage_frequency": prop_info.usage_frequency,
            "ontology_uri": prop_info.ontology_uri,
            "is_object_property": prop_info.is_object_property
        }
    
    return serialized

def deserialize_ontology_info(cached_data: Dict[str, Any]):
    """Deserialize cached data back to OntologyInfo object."""
    from models import OntologyInfo, ClassInfo, PropertyInfo
    
    # Reconstruct classes
    classes = {}
    for uri, class_data in cached_data.get("classes", {}).items():
        classes[uri] = ClassInfo(
            uri=class_data["uri"],
            name=class_data["name"],
            description=class_data.get("description"),
            properties=class_data.get("properties", []),
            instance_count=class_data.get("instance_count"),
            sample_instances=class_data.get("sample_instances", []),
            ontology_uri=class_data.get("ontology_uri")
        )
    
    # Reconstruct properties
    properties = {}
    for uri, prop_data in cached_data.get("properties", {}).items():
        properties[uri] = PropertyInfo(
            uri=prop_data["uri"],
            name=prop_data["name"],
            description=prop_data.get("description"),
            domain_classes=prop_data.get("domain_classes", []),
            range_classes=prop_data.get("range_classes", []),
            usage_frequency=prop_data.get("usage_frequency"),
            ontology_uri=prop_data.get("ontology_uri"),
            is_object_property=prop_data.get("is_object_property")
        )
    
    # Create OntologyInfo object
    ontology_info = OntologyInfo(
        classes=classes,
        properties=properties,
        namespaces=cached_data.get("namespaces", {}),
        domain_summary=cached_data.get("domain_summary", ""),
        total_instances=cached_data.get("total_instances", 0),
        last_updated=cached_data.get("last_updated", ""),
        llm_context=cached_data.get("llm_context", ""),
        compact_class_map=cached_data.get("compact_class_map", {}),
        compact_property_map=cached_data.get("compact_property_map", {}),
        reverse_class_map=cached_data.get("reverse_class_map", {}),
        reverse_property_map=cached_data.get("reverse_property_map", {})
    )
    
    return ontology_info
