"""
Knowledge Manager Tool for Multi-Swarm Framework.

This tool provides advanced knowledge management capabilities for agents,
including knowledge organization, quality assessment, and gap analysis.
"""

from typing import Dict, Any, List, Optional, Union
from datetime import datetime, UTC
from pathlib import Path
import json
from pydantic import BaseModel, Field, ConfigDict

from ..core.rag import RAGSystem, Document, SearchResult

class DateTimeEncoder(json.JSONEncoder):
    """Custom JSON encoder for datetime objects."""
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        return super().default(obj)

class KnowledgeCategory(BaseModel):
    """Represents a category of knowledge."""
    name: str
    description: str
    parent: Optional[str] = None
    metadata: Dict[str, Any] = Field(default_factory=dict)

class KnowledgeQuality(BaseModel):
    """Assessment of knowledge quality."""
    completeness: float = Field(ge=0.0, le=1.0)
    accuracy: float = Field(ge=0.0, le=1.0)
    relevance: float = Field(ge=0.0, le=1.0)
    timeliness: float = Field(ge=0.0, le=1.0)
    confidence: float = Field(ge=0.0, le=1.0)
    
    def overall_score(self) -> float:
        """Calculate overall quality score."""
        weights = {
            "completeness": 0.25,
            "accuracy": 0.3,
            "relevance": 0.2,
            "timeliness": 0.15,
            "confidence": 0.1
        }
        return sum(
            getattr(self, metric) * weight
            for metric, weight in weights.items()
        )

class KnowledgeGap(BaseModel):
    """Represents a gap in knowledge."""
    topic: str
    description: str
    priority: int = Field(ge=1, le=5)
    identified_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
    status: str = "open"
    metadata: Dict[str, Any] = Field(default_factory=dict)

class KnowledgeManager(BaseModel):
    """
    Advanced knowledge management tool for agents.
    
    Features:
    - Knowledge organization and categorization
    - Quality assessment and monitoring
    - Gap identification and tracking
    - Knowledge base maintenance
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    storage_path: Optional[Union[str, Path]] = None
    rag_system: Optional[RAGSystem] = None
    categories: Dict[str, KnowledgeCategory] = Field(default_factory=dict)
    quality_assessments: Dict[str, KnowledgeQuality] = Field(default_factory=dict)
    knowledge_gaps: List[KnowledgeGap] = Field(default_factory=list)
    
    def __init__(self, **data):
        super().__init__(**data)
        self._initialize()
        
    def _initialize(self):
        """Initialize the knowledge manager."""
        if self.storage_path:
            self.storage_path = Path(self.storage_path)
            self.storage_path.mkdir(parents=True, exist_ok=True)
            
            # Load existing state
            self._load_state()
            
    def _load_state(self):
        """Load manager state from disk."""
        if not self.storage_path:
            return
            
        state_path = self.storage_path / "knowledge_state.json"
        if not state_path.exists():
            return
            
        with open(state_path, "r") as f:
            state = json.load(f)
            
        # Load categories
        self.categories = {
            name: KnowledgeCategory(**data)
            for name, data in state.get("categories", {}).items()
        }
        
        # Load quality assessments
        self.quality_assessments = {
            doc_id: KnowledgeQuality(**data)
            for doc_id, data in state.get("quality_assessments", {}).items()
        }
        
        # Load knowledge gaps
        self.knowledge_gaps = [
            KnowledgeGap(**{
                **data,
                "identified_at": datetime.fromisoformat(data["identified_at"])
            })
            for data in state.get("knowledge_gaps", [])
        ]
        
    def _save_state(self):
        """Save manager state to disk."""
        if not self.storage_path:
            return
            
        state = {
            "categories": {
                name: category.model_dump()
                for name, category in self.categories.items()
            },
            "quality_assessments": {
                doc_id: assessment.model_dump()
                for doc_id, assessment in self.quality_assessments.items()
            },
            "knowledge_gaps": [
                {
                    **gap.model_dump(),
                    "identified_at": gap.identified_at.isoformat()
                }
                for gap in self.knowledge_gaps
            ]
        }
        
        with open(self.storage_path / "knowledge_state.json", "w") as f:
            json.dump(state, f, indent=2, cls=DateTimeEncoder)
            
    def add_category(
        self,
        name: str,
        description: str,
        parent: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> KnowledgeCategory:
        """Add a new knowledge category."""
        if name in self.categories:
            raise ValueError(f"Category already exists: {name}")
            
        if parent and parent not in self.categories:
            raise ValueError(f"Parent category not found: {parent}")
            
        category = KnowledgeCategory(
            name=name,
            description=description,
            parent=parent,
            metadata=metadata or {}
        )
        
        self.categories[name] = category
        self._save_state()
        
        return category
        
    def get_category(self, name: str) -> Optional[KnowledgeCategory]:
        """Get a category by name."""
        return self.categories.get(name)
        
    def list_categories(
        self,
        parent: Optional[str] = None
    ) -> List[KnowledgeCategory]:
        """List categories, optionally filtered by parent."""
        categories = []
        for category in self.categories.values():
            if parent is None or category.parent == parent:
                categories.append(category)
        return categories
        
    def assess_quality(
        self,
        document: Document,
        assessment: KnowledgeQuality
    ):
        """Assess the quality of a document."""
        self.quality_assessments[document.id] = assessment
        self._save_state()
        
    def get_quality(
        self,
        document_id: str
    ) -> Optional[KnowledgeQuality]:
        """Get quality assessment for a document."""
        return self.quality_assessments.get(document_id)
        
    def add_knowledge_gap(
        self,
        topic: str,
        description: str,
        priority: int = 3,
        metadata: Optional[Dict[str, Any]] = None
    ) -> KnowledgeGap:
        """Add a new knowledge gap."""
        gap = KnowledgeGap(
            topic=topic,
            description=description,
            priority=priority,
            metadata=metadata or {}
        )
        
        self.knowledge_gaps.append(gap)
        self._save_state()
        
        return gap
        
    def list_knowledge_gaps(
        self,
        status: Optional[str] = None,
        min_priority: Optional[int] = None
    ) -> List[KnowledgeGap]:
        """List knowledge gaps with optional filtering."""
        gaps = self.knowledge_gaps
        
        if status:
            gaps = [gap for gap in gaps if gap.status == status]
            
        if min_priority:
            gaps = [gap for gap in gaps if gap.priority >= min_priority]
            
        return sorted(gaps, key=lambda g: g.priority, reverse=True)
        
    def update_gap_status(
        self,
        topic: str,
        status: str,
        resolution: Optional[str] = None
    ):
        """Update the status of a knowledge gap."""
        for gap in self.knowledge_gaps:
            if gap.topic == topic:
                gap.status = status
                if resolution:
                    gap.metadata["resolution"] = resolution
                gap.metadata["updated_at"] = datetime.now(UTC).isoformat()
                break
                
        self._save_state()
        
    def analyze_coverage(
        self,
        category: Optional[str] = None
    ) -> Dict[str, Any]:
        """Analyze knowledge coverage and quality."""
        if not self.rag_system:
            raise ValueError("RAG system not initialized")
            
        # Filter documents by category if specified
        documents = self.rag_system.documents.values()
        if category:
            documents = [
                doc for doc in documents
                if doc.metadata.get("category") == category
            ]
            
        # Calculate metrics
        total_docs = len(documents)
        assessed_docs = sum(
            1 for doc in documents
            if doc.id in self.quality_assessments
        )
        
        quality_scores = [
            self.quality_assessments[doc.id].overall_score()
            for doc in documents
            if doc.id in self.quality_assessments
        ]
        
        avg_quality = (
            sum(quality_scores) / len(quality_scores)
            if quality_scores else 0.0
        )
        
        return {
            "total_documents": total_docs,
            "assessed_documents": assessed_docs,
            "average_quality": avg_quality,
            "coverage_ratio": assessed_docs / total_docs if total_docs > 0 else 0.0,
            "gaps_identified": len(self.knowledge_gaps),
            "open_gaps": len([g for g in self.knowledge_gaps if g.status == "open"])
        }
        
    def suggest_improvements(self) -> List[Dict[str, Any]]:
        """Suggest improvements for the knowledge base."""
        suggestions = []
        
        # Check document quality
        low_quality_docs = [
            (doc_id, assessment)
            for doc_id, assessment in self.quality_assessments.items()
            if assessment.overall_score() < 0.7
        ]
        
        if low_quality_docs:
            suggestions.append({
                "type": "quality_improvement",
                "description": "Documents requiring quality improvement",
                "documents": [
                    {
                        "id": doc_id,
                        "score": assessment.overall_score(),
                        "issues": [
                            metric for metric, value in assessment.model_dump().items()
                            if value < 0.7
                        ]
                    }
                    for doc_id, assessment in low_quality_docs
                ]
            })
            
        # Check for outdated content
        current_time = datetime.now(UTC)
        outdated_docs = [
            doc for doc in self.rag_system.documents.values()
            if (current_time - doc.created_at).days > 180
        ]
        
        if outdated_docs:
            suggestions.append({
                "type": "content_refresh",
                "description": "Documents requiring content refresh",
                "documents": [
                    {
                        "id": doc.id,
                        "age_days": (current_time - doc.created_at).days
                    }
                    for doc in outdated_docs
                ]
            })
            
        # Check knowledge gaps
        high_priority_gaps = [
            gap for gap in self.knowledge_gaps
            if gap.status == "open" and gap.priority >= 4
        ]
        
        if high_priority_gaps:
            suggestions.append({
                "type": "knowledge_gaps",
                "description": "High-priority knowledge gaps to address",
                "gaps": [
                    {
                        "topic": gap.topic,
                        "priority": gap.priority,
                        "identified_at": gap.identified_at.isoformat()
                    }
                    for gap in high_priority_gaps
                ]
            })
            
        return suggestions

if __name__ == "__main__":
    # Example usage
    manager = KnowledgeManager(
        storage_path="knowledge_manager_test"
    )
    
    # Add categories
    manager.add_category(
        "technical",
        "Technical documentation and guides"
    )
    manager.add_category(
        "python",
        "Python-specific knowledge",
        parent="technical"
    )
    
    # Add a knowledge gap
    gap = manager.add_knowledge_gap(
        "Advanced Python Patterns",
        "Need more information about advanced design patterns in Python",
        priority=4
    )
    
    # List categories
    categories = manager.list_categories()
    print("\nCategories:")
    for category in categories:
        print(f"- {category.name}: {category.description}")
        
    # List gaps
    gaps = manager.list_knowledge_gaps(min_priority=4)
    print("\nHigh-priority gaps:")
    for gap in gaps:
        print(f"- {gap.topic} (Priority: {gap.priority})")
        print(f"  {gap.description}") 