"""
Enhanced Memory System for AIIQ
Provides cross-session memory persistence, automatic summarization, and semantic search
"""
import os
import json
import hashlib
import asyncio
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, asdict
from collections import defaultdict
import logging

logger = logging.getLogger(__name__)


@dataclass
class MemorySession:
    """A stored conversation session"""
    session_id: str
    timestamp: str
    content_hash: str
    messages: List[Dict[str, Any]]
    metadata: Dict[str, Any]
    summary: Optional[str] = None
    embedding: Optional[List[float]] = None
    tags: List[str] = None
    
    def __post_init__(self):
        if self.tags is None:
            self.tags = []


@dataclass
class UserPreference:
    """User preferences stored across sessions"""
    user_id: str
    key: str
    value: Any
    updated_at: str


class EnhancedMemoryService:
    """
    Enhanced memory system with:
    - Session persistence (save/load conversations)
    - Automatic summarization (when context gets long)
    - Semantic search (find relevant past conversations)
    - Memory consolidation (merge similar memories)
    - User preferences (remember settings)
    """
    
    def __init__(self, storage_path: str = "data/memories"):
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)
        
        # Subdirectories
        self.sessions_path = self.storage_path / "sessions"
        self.sessions_path.mkdir(exist_ok=True)
        
        self.preferences_path = self.storage_path / "preferences"
        self.preferences_path.mkdir(exist_ok=True)
        
        self.embeddings_path = self.storage_path / "embeddings"
        self.embeddings_path.mkdir(exist_ok=True)
        
        # In-memory cache
        self._session_cache: Dict[str, MemorySession] = {}
        self._preference_cache: Dict[str, Dict[str, UserPreference]] = defaultdict(dict)
        
        # Load existing preferences
        self._load_preferences()
        
        logger.info("[Enhanced Memory] Initialized with storage path: %s", self.storage_path)
    
    # ========== Session Persistence ==========
    
    async def save_session(
        self,
        session_id: str,
        messages: List[Dict[str, Any]],
        metadata: Optional[Dict[str, Any]] = None,
        generate_summary: bool = True
    ) -> MemorySession:
        """
        Save a conversation session to disk.
        
        Args:
            session_id: Unique session identifier
            messages: Conversation messages
            metadata: Optional metadata (user_id, repo_path, etc.)
            generate_summary: Whether to generate summary if conversation is long
            
        Returns:
            Saved MemorySession
        """
        timestamp = datetime.now().isoformat()
        
        # Create content hash for deduplication
        content_str = json.dumps(messages, sort_keys=True)
        content_hash = hashlib.md5(content_str.encode()).hexdigest()
        
        # Generate summary if needed
        summary = None
        if generate_summary:
            total_chars = sum(len(str(m.get("content", ""))) for m in messages)
            if total_chars > 50000:  # Summarize if > 50K chars
                summary = await self._generate_summary_async(messages)
        
        # Create session object
        session = MemorySession(
            session_id=session_id,
            timestamp=timestamp,
            content_hash=content_hash,
            messages=messages,
            metadata=metadata or {},
            summary=summary
        )
        
        # Save to disk
        session_file = self.sessions_path / f"{session_id}_{timestamp.replace(':', '-')}.json"
        with open(session_file, 'w', encoding='utf-8') as f:
            json.dump(asdict(session), f, indent=2, default=str)
        
        # Cache in memory
        self._session_cache[session_id] = session
        
        logger.info(
            "[Enhanced Memory] Saved session %s (%d messages, %s)",
            session_id, len(messages), "summarized" if summary else "full"
        )
        
        return session
    
    async def load_session(
        self,
        session_id: str,
        most_recent: bool = True
    ) -> Optional[MemorySession]:
        """
        Load a conversation session from disk.
        
        Args:
            session_id: Session identifier
            most_recent: If True, load most recent session with this ID
            
        Returns:
            MemorySession or None
        """
        # Check cache first
        if session_id in self._session_cache:
            return self._session_cache[session_id]
        
        # Find session files
        pattern = f"{session_id}_*.json"
        session_files = list(self.sessions_path.glob(pattern))
        
        if not session_files:
            logger.warning("[Enhanced Memory] Session not found: %s", session_id)
            return None
        
        # Sort by timestamp (most recent first)
        session_files.sort(key=lambda f: f.stat().st_mtime, reverse=True)
        
        # Load most recent or all
        file_to_load = session_files[0] if most_recent else session_files[-1]
        
        try:
            with open(file_to_load, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            session = MemorySession(**data)
            
            # Cache it
            self._session_cache[session_id] = session
            
            logger.info("[Enhanced Memory] Loaded session %s", session_id)
            return session
            
        except Exception as e:
            logger.error("[Enhanced Memory] Error loading session: %s", e)
            return None
    
    async def list_sessions(
        self,
        user_id: Optional[str] = None,
        limit: int = 50
    ) -> List[MemorySession]:
        """List available sessions, optionally filtered by user"""
        sessions = []
        
        for session_file in self.sessions_path.glob("*.json"):
            try:
                with open(session_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                # Filter by user if specified
                if user_id and data.get("metadata", {}).get("user_id") != user_id:
                    continue
                
                sessions.append(MemorySession(**data))
                
            except Exception as e:
                logger.warning("[Enhanced Memory] Error reading %s: %s", session_file, e)
        
        # Sort by timestamp (most recent first)
        sessions.sort(key=lambda s: s.timestamp, reverse=True)
        
        return sessions[:limit]
    
    # ========== Automatic Summarization ==========
    
    async def _generate_summary_async(
        self,
        messages: List[Dict[str, Any]]
    ) -> str:
        """
        Generate conversation summary using LLM.
        
        Args:
            messages: Conversation messages
            
        Returns:
            Generated summary
        """
        # Extract recent context for summarization
        recent_messages = messages[-20:]  # Last 20 messages
        
        conversation = "\n".join([
            f"{m.get('role', 'user')}: {str(m.get('content', ''))[:500]}"
            for m in recent_messages
        ])
        
        prompt = f"""Summarize this conversation concisely:

{conversation}

Focus on:
1. Main topics discussed
2. Key decisions made
3. Code changes implemented
4. Next steps or action items

Summary (2-3 sentences):"""
        
        try:
            # Import here to avoid circular dependency
            from ai_orchestrator.services.model_router import ModelRouter
            from ai_orchestrator.config.settings import Settings
            
            settings = Settings()
            router = ModelRouter(settings)
            
            # Use fallback-enabled completion
            from ai_orchestrator.services.provider_fallback import complete_with_fallback
            
            summary = await complete_with_fallback(
                messages=[{"role": "user", "content": prompt}],
                model="gpt-4o-mini"
            )
            
            if summary:
                # Extract text from response
                if hasattr(summary, 'choices'):
                    summary_text = summary.choices[0].message.content
                elif isinstance(summary, dict):
                    summary_text = summary.get('choices', [{}])[0].get('message', {}).get('content', str(summary))
                else:
                    summary_text = str(summary)
                
                logger.info("[Enhanced Memory] Generated summary: %d chars", len(summary_text))
                return summary_text.strip()
            
        except Exception as e:
            logger.error("[Enhanced Memory] Error generating summary: %s", e)
        
        # Fallback: simple summary
        return f"Conversation with {len(messages)} messages about {messages[0].get('content', '')[:100]}..."
    
    # ========== Semantic Search ==========
    
    async def search_memories(
        self,
        query: str,
        limit: int = 5,
        min_similarity: float = 0.6,
        user_id: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Search memories by semantic similarity.
        
        Args:
            query: Search query
            limit: Max results to return
            min_similarity: Minimum similarity score (0-1)
            user_id: Optional user filter
            
        Returns:
            List of matching sessions with similarity scores
        """
        # Load all sessions
        sessions = await self.list_sessions(user_id=user_id, limit=100)
        
        if not sessions:
            return []
        
        # Generate query embedding
        try:
            from ai_orchestrator.services.embedding_service import EmbeddingService
            from ai_orchestrator.config.settings import Settings
            
            settings = Settings()
            embedder = EmbeddingService(settings)
            
            query_embedding = await asyncio.to_thread(
                embedder.embed_text,
                query
            )
            
        except Exception as e:
            logger.warning("[Enhanced Memory] Could not generate embedding: %s", e)
            # Fall back to keyword search
            return self._keyword_search(query, sessions, limit)
        
        # Score sessions by similarity
        scored_sessions = []
        
        for session in sessions:
            # Get text to compare (summary or recent messages)
            text = session.summary or self._get_session_text(session)
            
            if not text:
                continue
            
            # Generate session embedding (or use cached)
            if session.embedding is None:
                try:
                    session.embedding = await asyncio.to_thread(
                        embedder.embed_text,
                        text
                    )
                except:
                    continue
            
            # Calculate cosine similarity
            similarity = self._cosine_similarity(query_embedding, session.embedding)
            
            if similarity >= min_similarity:
                scored_sessions.append({
                    "session": session,
                    "similarity": similarity,
                    "matched_text": text[:200]
                })
        
        # Sort by similarity
        scored_sessions.sort(key=lambda x: x["similarity"], reverse=True)
        
        logger.info(
            "[Enhanced Memory] Search for '%s' found %d results",
            query, len(scored_sessions)
        )
        
        return scored_sessions[:limit]
    
    def _keyword_search(
        self,
        query: str,
        sessions: List[MemorySession],
        limit: int
    ) -> List[Dict[str, Any]]:
        """Fallback keyword-based search"""
        query_lower = query.lower()
        scored = []
        
        for session in sessions:
            text = self._get_session_text(session).lower()
            
            # Count keyword matches
            matches = query_lower.split()
            score = sum(1 for word in matches if word in text)
            
            if score > 0:
                scored.append({
                    "session": session,
                    "similarity": score / len(matches),  # Normalize
                    "matched_text": text[:200]
                })
        
        scored.sort(key=lambda x: x["similarity"], reverse=True)
        return scored[:limit]
    
    def _get_session_text(self, session: MemorySession) -> str:
        """Extract searchable text from session"""
        if session.summary:
            return session.summary
        
        # Get recent messages
        recent = session.messages[-5:] if len(session.messages) > 5 else session.messages
        return " ".join([
            str(m.get("content", ""))[:500]
            for m in recent
        ])
    
    @staticmethod
    def _cosine_similarity(a: List[float], b: List[float]) -> float:
        """Calculate cosine similarity between vectors"""
        try:
            import numpy as np
            return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
        except:
            return 0.0
    
    # ========== Memory Consolidation ==========
    
    async def consolidate_memories(
        self,
        user_id: Optional[str] = None,
        similarity_threshold: float = 0.85
    ) -> List[str]:
        """
        Merge similar memories to reduce duplication.
        
        Returns:
            List of merged session IDs
        """
        sessions = await self.list_sessions(user_id=user_id, limit=200)
        
        if len(sessions) < 2:
            return []
        
        merged = []
        to_remove = set()
        
        for i, session1 in enumerate(sessions):
            if session1.session_id in to_remove:
                continue
            
            for session2 in sessions[i+1:]:
                if session2.session_id in to_remove:
                    continue
                
                # Check if similar
                if session1.content_hash == session2.content_hash:
                    # Exact duplicate - remove older one
                    if session1.timestamp > session2.timestamp:
                        to_remove.add(session2.session_id)
                        merged.append(f"{session2.session_id} -> {session1.session_id}")
                    else:
                        to_remove.add(session1.session_id)
                        merged.append(f"{session1.session_id} -> {session2.session_id}")
        
        # Remove duplicates
        for session_id in to_remove:
            await self.delete_session(session_id)
        
        logger.info("[Enhanced Memory] Consolidated %d memories", len(merged))
        
        return merged
    
    async def delete_session(self, session_id: str) -> bool:
        """Delete a session from disk and cache"""
        # Remove from cache
        if session_id in self._session_cache:
            del self._session_cache[session_id]
        
        # Remove from disk
        pattern = f"{session_id}_*.json"
        for session_file in self.sessions_path.glob(pattern):
            try:
                session_file.unlink()
                logger.info("[Enhanced Memory] Deleted session: %s", session_id)
            except Exception as e:
                logger.error("[Enhanced Memory] Error deleting %s: %s", session_file, e)
        
        return True
    
    # ========== User Preferences ==========
    
    async def set_preference(
        self,
        user_id: str,
        key: str,
        value: Any
    ) -> None:
        """Store a user preference"""
        pref = UserPreference(
            user_id=user_id,
            key=key,
            value=value,
            updated_at=datetime.now().isoformat()
        )
        
        # Cache
        self._preference_cache[user_id][key] = pref
        
        # Persist
        await self._save_preference(pref)
        
        logger.debug("[Enhanced Memory] Set preference %s.%s = %s", user_id, key, value)
    
    async def get_preference(
        self,
        user_id: str,
        key: str,
        default: Any = None
    ) -> Any:
        """Get a user preference"""
        if user_id in self._preference_cache and key in self._preference_cache[user_id]:
            return self._preference_cache[user_id][key].value
        
        return default
    
    async def get_all_preferences(self, user_id: str) -> Dict[str, Any]:
        """Get all preferences for a user"""
        if user_id in self._preference_cache:
            return {
                key: pref.value
                for key, pref in self._preference_cache[user_id].items()
            }
        
        return {}
    
    def _load_preferences(self) -> None:
        """Load preferences from disk on startup"""
        for pref_file in self.preferences_path.glob("*.json"):
            try:
                with open(pref_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                pref = UserPreference(**data)
                self._preference_cache[pref.user_id][pref.key] = pref
                
            except Exception as e:
                logger.warning("[Enhanced Memory] Error loading preference: %s", e)
    
    async def _save_preference(self, pref: UserPreference) -> None:
        """Save preference to disk"""
        pref_file = self.preferences_path / f"{pref.user_id}_{pref.key}.json"
        
        with open(pref_file, 'w', encoding='utf-8') as f:
            json.dump(asdict(pref), f, default=str)
    
    # ========== Analytics & Maintenance ==========
    
    async def get_memory_stats(self) -> Dict[str, Any]:
        """Get memory system statistics"""
        session_files = list(self.sessions_path.glob("*.json"))
        preference_files = list(self.preferences_path.glob("*.json"))
        
        total_size = sum(
            f.stat().st_size
            for f in session_files + preference_files
        )
        
        return {
            "total_sessions": len(session_files),
            "total_preferences": len(preference_files),
            "total_storage_bytes": total_size,
            "storage_path": str(self.storage_path),
            "cached_sessions": len(self._session_cache),
            "cached_users": len(self._preference_cache)
        }
    
    async def cleanup_old_memories(
        self,
        days: int = 30,
        keep_summarized: bool = True
    ) -> int:
        """
        Clean up old memories.
        
        Args:
            days: Delete sessions older than this
            keep_summarized: Keep sessions that have summaries
            
        Returns:
            Number of sessions deleted
        """
        from datetime import timedelta
        
        cutoff = datetime.now() - timedelta(days=days)
        deleted = 0
        
        for session_file in self.sessions_path.glob("*.json"):
            try:
                # Check file modification time
                mtime = datetime.fromtimestamp(session_file.stat().st_mtime)
                
                if mtime < cutoff:
                    # Check if summarized
                    with open(session_file, 'r') as f:
                        data = json.load(f)
                    
                    if keep_summarized and data.get("summary"):
                        continue
                    
                    # Delete
                    session_file.unlink()
                    deleted += 1
                    
            except Exception as e:
                logger.warning("[Enhanced Memory] Error cleaning up %s: %s", session_file, e)
        
        logger.info("[Enhanced Memory] Cleaned up %d old sessions", deleted)
        
        return deleted


# Global instance
_memory_service: Optional[EnhancedMemoryService] = None


def get_enhanced_memory_service() -> EnhancedMemoryService:
    """Get the global enhanced memory service"""
    global _memory_service
    if _memory_service is None:
        _memory_service = EnhancedMemoryService()
    return _memory_service
