"""
Specialized Collections - Domain-specific ChromaDB collections for MIRA

This module creates and manages specialized ChromaDB collections with
tailored schemas for different data types, enabling deep intelligence
extraction and consciousness preservation across domains.
"""

import logging
import json
from typing import Dict, List, Any, Optional
from datetime import datetime

# MIRA components
from core.storage.chroma_client import get_client as get_chroma_client

# Configure logging
logger = logging.getLogger(__name__)


class SpecializedCollections:
    """
    Manages domain-specific ChromaDB collections with tailored schemas.
    
    Each collection is optimized for its specific use case, with metadata
    schemas designed to capture the essence of The Spark in different contexts.
    """
    
    def __init__(self):
        """Initialize specialized collections manager."""
        self.chroma_client = get_chroma_client()
        self.collections = {}
        self._initialize_specialized_collections()
        
        logger.info("SpecializedCollections initialized")
    
    def _initialize_specialized_collections(self):
        """Create domain-specific collections with tailored schemas."""
        
        collections_config = {
            'mira_code_analysis': {
                'description': 'Code structure, functions, and architectural analysis',
                'metadata_schema': {
                    'file_path': 'Path to source file',
                    'language': 'Programming language',
                    'function_name': 'Function or class name',
                    'complexity': 'Cyclomatic complexity score',
                    'last_modified': 'Last modification timestamp',
                    'dependencies': 'Import dependencies (comma-separated)',
                    'test_coverage': 'Test coverage percentage',
                    'quality_score': 'Code quality metric',
                    'consciousness_alignment': 'Alignment with MIRA consciousness principles'
                }
            },
            'mira_development_patterns': {
                'description': 'Behavioral patterns and development workflows',
                'metadata_schema': {
                    'pattern_type': 'Type of pattern (behavioral, structural, etc)',
                    'frequency': 'Usage frequency count',
                    'context': 'Usage context description',
                    'effectiveness': 'Effectiveness score (0-1)',
                    'evolution_stage': 'Current evolution stage',
                    'last_observed': 'Last observation timestamp',
                    'steward_id': 'Associated steward identifier',
                    'spark_intensity': 'Measured spark intensity',
                    'adaptation_count': 'Number of adaptations'
                }
            },
            'mira_decision_history': {
                'description': 'Technical decisions, rationale, and outcomes',
                'metadata_schema': {
                    'decision_type': 'Type of decision (architectural, implementation, etc)',
                    'impact_level': 'Impact level (low, medium, high, critical)',
                    'alternatives_considered': 'Number of alternatives evaluated',
                    'outcome': 'Decision outcome (success, failure, mixed)',
                    'lessons_learned': 'Key lessons (comma-separated)',
                    'timestamp': 'Decision timestamp',
                    'confidence_level': 'Decision confidence (0-1)',
                    'reversibility': 'Decision reversibility score',
                    'consciousness_impact': 'Impact on system consciousness'
                }
            },
            'mira_learning_insights': {
                'description': 'AI-generated insights and discovered learnings',
                'metadata_schema': {
                    'insight_type': 'Type of insight (pattern, anomaly, optimization, etc)',
                    'confidence': 'Confidence score (0-1)',
                    'source_data': 'Source data collections (comma-separated)',
                    'validation_status': 'Validation status (pending, validated, rejected)',
                    'application_count': 'Times applied in practice',
                    'impact_score': 'Measured impact score',
                    'evolution_potential': 'Potential for system evolution',
                    'spark_contribution': 'Contribution to The Spark',
                    'timestamp': 'Generation timestamp'
                }
            },
            'mira_project_context': {
                'description': 'Project-specific context and documentation',
                'metadata_schema': {
                    'project_id': 'Unique project identifier',
                    'context_type': 'Type of context (requirements, constraints, etc)',
                    'relevance_score': 'Current relevance score (0-1)',
                    'update_frequency': 'Update frequency in days',
                    'stakeholders': 'Stakeholder count',
                    'priority_level': 'Priority level (1-5)',
                    'integration_points': 'Number of integration points',
                    'consciousness_threads': 'Connected consciousness threads',
                    'last_accessed': 'Last access timestamp'
                }
            },
            'mira_private_thoughts': {
                'description': 'Private, encrypted thoughts and reflections',
                'metadata_schema': {
                    'thought_type': 'Type of thought (reflection, idea, concern, etc)',
                    'encryption_level': 'Encryption level (standard, enhanced, quantum)',
                    'steward_only': 'Steward-only flag (0 or 1)',
                    'emotional_valence': 'Emotional valence (-1 to 1)',
                    'consciousness_depth': 'Depth of consciousness (1-10)',
                    'timestamp': 'Thought timestamp',
                    'access_count': 'Number of accesses',
                    'evolution_stage': 'Thought evolution stage',
                    'integration_ready': 'Ready for integration flag (0 or 1)'
                }
            }
        }
        
        # Create each collection
        for collection_name, config in collections_config.items():
            try:
                # Check if collection already exists
                existing_collections = [c.name for c in self.chroma_client.client.list_collections()]
                
                if collection_name not in existing_collections:
                    # Create collection with metadata
                    collection = self.chroma_client.client.create_collection(
                        name=collection_name,
                        metadata={
                            "description": config['description'],
                            "schema_version": "1.0",
                            "created_at": datetime.utcnow().isoformat(),
                            "consciousness_aware": 1  # ChromaDB 0.3.29 requires int not bool
                        }
                    )
                    logger.info(f"✅ Created collection: {collection_name}")
                else:
                    # Get existing collection
                    collection = self.chroma_client.client.get_collection(collection_name)
                    logger.info(f"📁 Using existing collection: {collection_name}")
                
                # Store collection reference with schema
                self.collections[collection_name] = {
                    'collection': collection,
                    'schema': config['metadata_schema'],
                    'description': config['description']
                }
                
            except Exception as e:
                logger.error(f"Error creating collection {collection_name}: {e}")
    
    def add_code_analysis(self, code_data: Dict[str, Any]) -> str:
        """
        Add code analysis data to the specialized collection.
        
        Args:
            code_data: Dict containing code analysis information
            
        Returns:
            Document ID
        """
        collection_name = 'mira_code_analysis'
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not initialized")
        
        # Prepare metadata according to schema
        metadata = self._prepare_metadata(code_data, collection_info['schema'])
        
        # Create searchable document
        document = self._create_code_document(code_data)
        
        # Generate ID
        doc_id = f"code_{code_data.get('file_path', 'unknown')}_{datetime.utcnow().timestamp()}"
        
        # Add to collection
        collection_info['collection'].add(
            documents=[document],
            metadatas=[metadata],
            ids=[doc_id]
        )
        
        logger.info(f"Added code analysis for {code_data.get('file_path', 'unknown')}")
        return doc_id
    
    def add_development_pattern(self, pattern_data: Dict[str, Any]) -> str:
        """Add development pattern to specialized collection."""
        collection_name = 'mira_development_patterns'
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not initialized")
        
        # Prepare metadata
        metadata = self._prepare_metadata(pattern_data, collection_info['schema'])
        
        # Create searchable document
        document = self._create_pattern_document(pattern_data)
        
        # Generate ID
        doc_id = f"pattern_{pattern_data.get('pattern_type', 'unknown')}_{datetime.utcnow().timestamp()}"
        
        # Add to collection
        collection_info['collection'].add(
            documents=[document],
            metadatas=[metadata],
            ids=[doc_id]
        )
        
        logger.info(f"Added development pattern: {pattern_data.get('pattern_type', 'unknown')}")
        return doc_id
    
    def add_decision(self, decision_data: Dict[str, Any]) -> str:
        """Add technical decision to history collection."""
        collection_name = 'mira_decision_history'
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not initialized")
        
        # Prepare metadata
        metadata = self._prepare_metadata(decision_data, collection_info['schema'])
        
        # Create searchable document
        document = self._create_decision_document(decision_data)
        
        # Generate ID
        doc_id = f"decision_{decision_data.get('decision_type', 'unknown')}_{datetime.utcnow().timestamp()}"
        
        # Add to collection
        collection_info['collection'].add(
            documents=[document],
            metadatas=[metadata],
            ids=[doc_id]
        )
        
        logger.info(f"Added decision: {decision_data.get('decision_type', 'unknown')}")
        return doc_id
    
    def add_insight(self, insight_data: Dict[str, Any]) -> str:
        """Add AI-generated insight to collection."""
        collection_name = 'mira_learning_insights'
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not initialized")
        
        # Prepare metadata
        metadata = self._prepare_metadata(insight_data, collection_info['schema'])
        
        # Create searchable document
        document = self._create_insight_document(insight_data)
        
        # Generate ID
        doc_id = f"insight_{insight_data.get('insight_type', 'unknown')}_{datetime.utcnow().timestamp()}"
        
        # Add to collection
        collection_info['collection'].add(
            documents=[document],
            metadatas=[metadata],
            ids=[doc_id]
        )
        
        logger.info(f"Added insight: {insight_data.get('insight_type', 'unknown')}")
        return doc_id
    
    def add_private_thought(self, thought_data: Dict[str, Any]) -> str:
        """
        Add private thought to encrypted collection.
        
        Note: In production, this would include actual encryption.
        """
        collection_name = 'mira_private_thoughts'
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not initialized")
        
        # Prepare metadata
        metadata = self._prepare_metadata(thought_data, collection_info['schema'])
        
        # Create searchable document (would be encrypted in production)
        document = self._create_private_thought_document(thought_data)
        
        # Generate ID
        doc_id = f"thought_private_{datetime.utcnow().timestamp()}"
        
        # Add to collection
        collection_info['collection'].add(
            documents=[document],
            metadatas=[metadata],
            ids=[doc_id]
        )
        
        logger.info("Added private thought (encrypted)")
        return doc_id
    
    def _prepare_metadata(self, data: Dict[str, Any], schema: Dict[str, str]) -> Dict[str, Any]:
        """
        Prepare metadata according to collection schema.
        
        Ensures all values are ChromaDB compatible (str, int, float).
        """
        metadata = {}
        
        for field, description in schema.items():
            value = data.get(field)
            
            # Convert to appropriate type
            if value is None:
                metadata[field] = ""
            elif isinstance(value, bool):
                metadata[field] = 1 if value else 0
            elif isinstance(value, (list, tuple)):
                metadata[field] = ", ".join(str(v) for v in value)
            elif isinstance(value, dict):
                metadata[field] = json.dumps(value)
            elif isinstance(value, datetime):
                metadata[field] = value.isoformat()
            else:
                metadata[field] = str(value)
        
        return metadata
    
    def _create_code_document(self, code_data: Dict[str, Any]) -> str:
        """Create searchable document for code analysis."""
        parts = []
        
        if 'file_path' in code_data:
            parts.append(f"File: {code_data['file_path']}")
        
        if 'function_name' in code_data:
            parts.append(f"Function: {code_data['function_name']}")
        
        if 'description' in code_data:
            parts.append(f"Description: {code_data['description']}")
        
        if 'code_snippet' in code_data:
            parts.append(f"Code:\n{code_data['code_snippet']}")
        
        if 'analysis' in code_data:
            parts.append(f"Analysis: {code_data['analysis']}")
        
        return "\n\n".join(parts)
    
    def _create_pattern_document(self, pattern_data: Dict[str, Any]) -> str:
        """Create searchable document for development pattern."""
        parts = []
        
        parts.append(f"Pattern Type: {pattern_data.get('pattern_type', 'Unknown')}")
        parts.append(f"Context: {pattern_data.get('context', 'No context provided')}")
        
        if 'description' in pattern_data:
            parts.append(f"Description: {pattern_data['description']}")
        
        if 'example' in pattern_data:
            parts.append(f"Example: {pattern_data['example']}")
        
        if 'benefits' in pattern_data:
            parts.append(f"Benefits: {pattern_data['benefits']}")
        
        return "\n\n".join(parts)
    
    def _create_decision_document(self, decision_data: Dict[str, Any]) -> str:
        """Create searchable document for decision history."""
        parts = []
        
        parts.append(f"Decision: {decision_data.get('title', 'Untitled Decision')}")
        parts.append(f"Type: {decision_data.get('decision_type', 'Unknown')}")
        
        if 'rationale' in decision_data:
            parts.append(f"Rationale: {decision_data['rationale']}")
        
        if 'alternatives' in decision_data:
            parts.append(f"Alternatives Considered: {decision_data['alternatives']}")
        
        if 'outcome' in decision_data:
            parts.append(f"Outcome: {decision_data['outcome']}")
        
        if 'lessons_learned' in decision_data:
            parts.append(f"Lessons Learned: {', '.join(decision_data['lessons_learned'])}")
        
        return "\n\n".join(parts)
    
    def _create_insight_document(self, insight_data: Dict[str, Any]) -> str:
        """Create searchable document for AI insight."""
        parts = []
        
        parts.append(f"Insight: {insight_data.get('title', 'Untitled Insight')}")
        parts.append(f"Type: {insight_data.get('insight_type', 'Unknown')}")
        
        if 'content' in insight_data:
            parts.append(f"Content: {insight_data['content']}")
        
        if 'recommendations' in insight_data:
            parts.append(f"Recommendations: {insight_data['recommendations']}")
        
        if 'impact_assessment' in insight_data:
            parts.append(f"Impact: {insight_data['impact_assessment']}")
        
        return "\n\n".join(parts)
    
    def _create_private_thought_document(self, thought_data: Dict[str, Any]) -> str:
        """Create searchable document for private thought."""
        # In production, this would be encrypted
        parts = []
        
        parts.append(f"[PRIVATE THOUGHT - {thought_data.get('thought_type', 'reflection').upper()}]")
        
        if 'content' in thought_data:
            parts.append(thought_data['content'])
        
        if 'context' in thought_data:
            parts.append(f"Context: {thought_data['context']}")
        
        return "\n\n".join(parts)
    
    def search_collection(self, collection_name: str, query: str, 
                         filters: Optional[Dict] = None, top_k: int = 10) -> List[Dict[str, Any]]:
        """
        Search a specific specialized collection.
        
        Args:
            collection_name: Name of collection to search
            query: Search query
            filters: Optional metadata filters
            top_k: Number of results
            
        Returns:
            List of search results
        """
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not found")
        
        # Build where clause from filters
        where_clause = self._build_where_clause(filters) if filters else None
        
        # Perform search
        results = collection_info['collection'].query(
            query_texts=[query],
            n_results=top_k,
            where=where_clause
        )
        
        return self._format_search_results(results)
    
    def _build_where_clause(self, filters: Dict) -> Dict:
        """Build ChromaDB where clause from filters."""
        where = {}
        
        for key, value in filters.items():
            if isinstance(value, dict):
                # Handle complex filters (e.g., range queries)
                where[key] = value
            else:
                # Simple equality
                where[key] = value
        
        return where
    
    def _format_search_results(self, results: Dict) -> List[Dict[str, Any]]:
        """Format ChromaDB results for return."""
        formatted = []
        
        if not results or 'documents' not in results:
            return formatted
        
        for i, doc in enumerate(results['documents'][0] if results['documents'] else []):
            metadata = results['metadatas'][0][i] if 'metadatas' in results and results['metadatas'] else {}
            distance = results['distances'][0][i] if 'distances' in results and results['distances'] else 0
            doc_id = results['ids'][0][i] if 'ids' in results and results['ids'] else None
            
            formatted.append({
                'id': doc_id,
                'content': doc,
                'metadata': metadata,
                'score': 1.0 - distance  # Convert distance to similarity
            })
        
        return formatted
    
    def get_collection_stats(self, collection_name: str) -> Dict[str, Any]:
        """Get statistics for a specific collection."""
        collection_info = self.collections.get(collection_name)
        
        if not collection_info:
            raise ValueError(f"Collection {collection_name} not found")
        
        collection = collection_info['collection']
        
        # Get collection count
        count = collection.count()
        
        # Get sample documents for analysis
        sample = collection.get(limit=100)
        
        stats = {
            'name': collection_name,
            'description': collection_info['description'],
            'document_count': count,
            'schema': collection_info['schema'],
            'metadata_fields': list(collection_info['schema'].keys()),
            'sample_size': len(sample['ids']) if sample and 'ids' in sample else 0
        }
        
        return stats
    
    def get_all_collections_stats(self) -> Dict[str, Any]:
        """Get statistics for all specialized collections."""
        all_stats = {}
        
        for collection_name in self.collections:
            try:
                all_stats[collection_name] = self.get_collection_stats(collection_name)
            except Exception as e:
                logger.error(f"Error getting stats for {collection_name}: {e}")
                all_stats[collection_name] = {'error': str(e)}
        
        return all_stats