"""
Domain Memory System for Toxo.

This module implements persistent memory for domain-specific
interactions, patterns, and knowledge accumulation.
"""

import asyncio
import json
import numpy as np
from typing import Dict, Any, List, Optional, Tuple, Union
from dataclasses import dataclass, field, asdict
from datetime import datetime, timedelta
import uuid
from pathlib import Path
import sqlite3
import threading

from ..core.config import ToxoConfig
from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError


@dataclass
class MemoryItem:
    """A single memory item storing an interaction."""
    
    id: str
    input_text: str
    response: str
    domain: str
    task_type: str
    timestamp: datetime
    importance: float = 0.5
    context: Dict[str, Any] = field(default_factory=dict)
    embedding: Optional[np.ndarray] = None
    tags: List[str] = field(default_factory=list)
    access_count: int = 0
    last_accessed: Optional[datetime] = None
    
    def __post_init__(self):
        """Convert string timestamp to datetime if needed."""
        if isinstance(self.timestamp, str):
            self.timestamp = datetime.fromisoformat(self.timestamp)
        if isinstance(self.last_accessed, str):
            self.last_accessed = datetime.fromisoformat(self.last_accessed)
    
    @property
    def content(self) -> str:
        """Get the content of this memory item."""
        return f"Input: {self.input_text}\nResponse: {self.response}"
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for storage."""
        data = asdict(self)
        data['timestamp'] = self.timestamp.isoformat()
        if self.last_accessed:
            data['last_accessed'] = self.last_accessed.isoformat()
        if self.embedding is not None:
            data['embedding'] = self.embedding.tolist()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'MemoryItem':
        """Create from dictionary."""
        if 'embedding' in data and data['embedding'] is not None:
            data['embedding'] = np.array(data['embedding'])
        return cls(**data)


class DomainMemory:
    """
    Domain-specific memory system for Toxo layers.
    
    Provides persistent storage and retrieval of interactions,
    with semantic search and importance-based filtering.
    """
    
    def __init__(self, config: ToxoConfig, domain: str = None):
        """Initialize domain memory."""
        self.config = config
        self.domain = domain or "default"
        self.logger = get_logger(f"toxo.memory.{self.domain}")
        
        # Memory storage
        self.memories: Dict[str, MemoryItem] = {}
        self.embeddings_cache: Dict[str, np.ndarray] = {}
        
        # Database setup
        self.db_path = Path(config.data_dir) / "memory" / f"{self.domain}_memory.db"
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Threading lock for database operations
        self._db_lock = threading.Lock()
        
        # Initialize database
        self._init_database()
        
        # Load existing memories
        self._load_memories()
        
        self.logger.info(f"Domain memory initialized for '{self.domain}' with {len(self.memories)} items")
    
    def _init_database(self) -> None:
        """Initialize SQLite database."""
        with self._db_lock:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()
            
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS memories (
                    id TEXT PRIMARY KEY,
                    input_text TEXT NOT NULL,
                    response TEXT NOT NULL,
                    domain TEXT NOT NULL,
                    task_type TEXT NOT NULL,
                    timestamp TEXT NOT NULL,
                    importance REAL DEFAULT 0.5,
                    context TEXT,
                    embedding BLOB,
                    tags TEXT,
                    access_count INTEGER DEFAULT 0,
                    last_accessed TEXT
                )
            """)
            
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_domain ON memories(domain)
            """)
            
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_task_type ON memories(task_type)
            """)
            
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_importance ON memories(importance)
            """)
            
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_timestamp ON memories(timestamp)
            """)
            
            conn.commit()
            conn.close()
    
    def _load_memories(self) -> None:
        """Load memories from database."""
        with self._db_lock:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()
            
            cursor.execute("""
                SELECT id, input_text, response, domain, task_type, timestamp,
                       importance, context, embedding, tags, access_count, last_accessed
                FROM memories
                WHERE domain = ?
            """, (self.domain,))
            
            for row in cursor.fetchall():
                memory_data = {
                    'id': row[0],
                    'input_text': row[1],
                    'response': row[2],
                    'domain': row[3],
                    'task_type': row[4],
                    'timestamp': row[5],
                    'importance': row[6],
                    'context': json.loads(row[7]) if row[7] else {},
                    'embedding': np.frombuffer(row[8]) if row[8] else None,
                    'tags': json.loads(row[9]) if row[9] else [],
                    'access_count': row[10],
                    'last_accessed': row[11]
                }
                
                memory = MemoryItem.from_dict(memory_data)
                self.memories[memory.id] = memory
            
            conn.close()
    
    async def store_interaction(
        self,
        input_text: str,
        response: str,
        context: Optional[Dict[str, Any]] = None,
        importance: float = 0.5,
        task_type: str = "general",
        tags: Optional[List[str]] = None
    ) -> str:
        """
        Store a new interaction in memory.
        
        Args:
            input_text: The input text
            response: The response text
            context: Additional context information
            importance: Importance score (0.0 to 1.0)
            task_type: Type of task
            tags: Optional tags for categorization
            
        Returns:
            Memory item ID
        """
        memory_id = str(uuid.uuid4())
        
        memory = MemoryItem(
            id=memory_id,
            input_text=input_text,
            response=response,
            domain=self.domain,
            task_type=task_type,
            timestamp=datetime.now(),
            importance=max(0.0, min(1.0, importance)),
            context=context or {},
            tags=tags or []
        )
        
        # Store in memory
        self.memories[memory_id] = memory
        
        # Store in database
        await self._save_memory_to_db(memory)
        
        self.logger.debug(f"Stored interaction {memory_id} with importance {importance}")
        return memory_id
    
    async def _save_memory_to_db(self, memory: MemoryItem) -> None:
        """Save memory item to database."""
        with self._db_lock:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()
            
            embedding_blob = memory.embedding.tobytes() if memory.embedding is not None else None
            
            cursor.execute("""
                INSERT OR REPLACE INTO memories 
                (id, input_text, response, domain, task_type, timestamp,
                 importance, context, embedding, tags, access_count, last_accessed)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                memory.id,
                memory.input_text,
                memory.response,
                memory.domain,
                memory.task_type,
                memory.timestamp.isoformat(),
                memory.importance,
                json.dumps(memory.context),
                embedding_blob,
                json.dumps(memory.tags),
                memory.access_count,
                memory.last_accessed.isoformat() if memory.last_accessed else None
            ))
            
            conn.commit()
            conn.close()
    
    async def retrieve(
        self,
        query: str,
        limit: int = 10,
        min_importance: float = 0.0,
        task_type: Optional[str] = None,
        since: Optional[datetime] = None,
        tags: Optional[List[str]] = None
    ) -> List[MemoryItem]:
        """
        Retrieve relevant memories based on query.
        
        Args:
            query: Search query
            limit: Maximum number of results
            min_importance: Minimum importance threshold
            task_type: Filter by task type
            since: Only return memories after this date
            tags: Filter by tags
            
        Returns:
            List of relevant memory items
        """
        # Filter memories based on criteria
        candidates = []
        
        for memory in self.memories.values():
            # Apply filters
            if memory.importance < min_importance:
                continue
            
            if task_type and memory.task_type != task_type:
                continue
            
            if since and memory.timestamp < since:
                continue
            
            if tags and not any(tag in memory.tags for tag in tags):
                continue
            
            candidates.append(memory)
        
        # Sort by relevance (simple text matching for now)
        query_lower = query.lower()
        
        def relevance_score(memory: MemoryItem) -> float:
            score = 0.0
            
            # Text similarity (simple keyword matching)
            input_matches = sum(1 for word in query_lower.split() 
                              if word in memory.input_text.lower())
            response_matches = sum(1 for word in query_lower.split() 
                                 if word in memory.response.lower())
            
            score += (input_matches + response_matches) * 0.3
            score += memory.importance * 0.4
            
            # Recency bonus
            days_old = (datetime.now() - memory.timestamp).days
            recency_score = max(0, 1.0 - (days_old / 365.0))  # Decay over a year
            score += recency_score * 0.3
            
            return score
        
        # Sort by relevance and take top results
        candidates.sort(key=relevance_score, reverse=True)
        results = candidates[:limit]
        
        # Update access information
        for memory in results:
            memory.access_count += 1
            memory.last_accessed = datetime.now()
            await self._save_memory_to_db(memory)
        
        self.logger.debug(f"Retrieved {len(results)} memories for query: {query[:50]}...")
        return results
    
    async def get_memory(self, memory_id: str) -> Optional[MemoryItem]:
        """Get a specific memory by ID."""
        memory = self.memories.get(memory_id)
        if memory:
            memory.access_count += 1
            memory.last_accessed = datetime.now()
            await self._save_memory_to_db(memory)
        return memory
    
    async def update_memory(
        self,
        memory_id: str,
        updates: Dict[str, Any]
    ) -> bool:
        """Update a memory item."""
        if memory_id not in self.memories:
            return False
        
        memory = self.memories[memory_id]
        
        # Apply updates
        for key, value in updates.items():
            if hasattr(memory, key):
                setattr(memory, key, value)
        
        # Save to database
        await self._save_memory_to_db(memory)
        
        self.logger.debug(f"Updated memory {memory_id}")
        return True
    
    async def delete_memory(self, memory_id: str) -> bool:
        """Delete a memory item."""
        if memory_id not in self.memories:
            return False
        
        # Remove from memory
        del self.memories[memory_id]
        
        # Remove from database
        with self._db_lock:
            conn = sqlite3.connect(str(self.db_path))
            cursor = conn.cursor()
            cursor.execute("DELETE FROM memories WHERE id = ?", (memory_id,))
            conn.commit()
            conn.close()
        
        self.logger.debug(f"Deleted memory {memory_id}")
        return True
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get memory statistics."""
        if not self.memories:
            return {
                "total_memories": 0,
                "avg_importance": 0.0,
                "task_types": {},
                "total_access_count": 0
            }
        
        task_type_counts = {}
        total_access_count = 0
        total_importance = 0
        
        for memory in self.memories.values():
            task_type_counts[memory.task_type] = task_type_counts.get(memory.task_type, 0) + 1
            total_access_count += memory.access_count
            total_importance += memory.importance
        
        return {
            "total_memories": len(self.memories),
            "avg_importance": total_importance / len(self.memories),
            "task_types": task_type_counts,
            "total_access_count": total_access_count,
            "domain": self.domain
        }
    
    async def cleanup_old_memories(
        self,
        max_age_days: int = 365,
        min_importance: float = 0.1,
        max_memories: int = 10000
    ) -> int:
        """
        Clean up old or low-importance memories.
        
        Args:
            max_age_days: Maximum age in days
            min_importance: Minimum importance to keep
            max_memories: Maximum number of memories to keep
            
        Returns:
            Number of memories deleted
        """
        cutoff_date = datetime.now() - timedelta(days=max_age_days)
        
        # Find memories to delete
        to_delete = []
        
        for memory_id, memory in self.memories.items():
            # Delete if too old and low importance
            if (memory.timestamp < cutoff_date and 
                memory.importance < min_importance):
                to_delete.append(memory_id)
        
        # If still too many memories, delete oldest low-importance ones
        if len(self.memories) > max_memories:
            remaining_memories = [
                memory for memory_id, memory in self.memories.items()
                if memory_id not in to_delete
            ]
            
            # Sort by importance and age (keep important and recent)
            remaining_memories.sort(
                key=lambda m: (m.importance, m.timestamp),
                reverse=True
            )
            
            # Mark excess memories for deletion
            excess_count = len(remaining_memories) - max_memories
            if excess_count > 0:
                for memory in remaining_memories[-excess_count:]:
                    to_delete.append(memory.id)
        
        # Delete marked memories
        deleted_count = 0
        for memory_id in to_delete:
            if await self.delete_memory(memory_id):
                deleted_count += 1
        
        if deleted_count > 0:
            self.logger.info(f"Cleaned up {deleted_count} old memories")
        
        return deleted_count
    
    def get_state(self) -> Dict[str, Any]:
        """Get memory state for serialization."""
        return {
            "domain": self.domain,
            "memory_count": len(self.memories),
            "statistics": self.get_statistics(),
            "db_path": str(self.db_path)
        }
    
    def load_state(self, state: Dict[str, Any]) -> None:
        """Load memory state from serialization."""
        # State loading is handled by database persistence
        # This method is for compatibility with the package system
        self.logger.info(f"Memory state loaded: {state.get('memory_count', 0)} memories")
    
    def get_all_embeddings(self) -> Optional[np.ndarray]:
        """Get all embeddings as a single array."""
        embeddings = []
        for memory in self.memories.values():
            if memory.embedding is not None:
                embeddings.append(memory.embedding)
        
        if embeddings:
            return np.array(embeddings)
        return None
    
    async def get_recent_interactions(
        self,
        hours: int = 24,
        limit: int = 50
    ) -> List[MemoryItem]:
        """Get recent interactions within specified hours."""
        since = datetime.now() - timedelta(hours=hours)
        
        recent = [
            memory for memory in self.memories.values()
            if memory.timestamp >= since
        ]
        
        # Sort by timestamp, most recent first
        recent.sort(key=lambda m: m.timestamp, reverse=True)
        
        return recent[:limit]
    
    async def get_important_memories(
        self,
        min_importance: float = 0.7,
        limit: int = 100
    ) -> List[MemoryItem]:
        """Get highly important memories."""
        important = [
            memory for memory in self.memories.values()
            if memory.importance >= min_importance
        ]
        
        # Sort by importance, then by recency
        important.sort(key=lambda m: (m.importance, m.timestamp), reverse=True)
        
        return important[:limit]
    
    async def store_feedback(
        self,
        input_text: str,
        expected_output: str,
        actual_output: Optional[str] = None,
        feedback_type: str = "correction",
        analysis_score: Optional[float] = None,
        suggestions: Optional[List[str]] = None,
        iteration: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Store feedback for learning and improvement.
        
        Args:
            input_text: Original input text
            expected_output: Desired output
            actual_output: What was actually produced
            feedback_type: Type of feedback (correction, improvement, refinement)
            analysis_score: Quality score from self-analysis
            suggestions: List of improvement suggestions
            iteration: Training iteration number
            metadata: Additional metadata
            
        Returns:
            Memory item ID for the stored feedback
        """
        feedback_id = str(uuid.uuid4())
        
        # Create comprehensive feedback content
        feedback_content = f"FEEDBACK - {feedback_type.upper()}\n"
        feedback_content += f"Input: {input_text}\n"
        feedback_content += f"Expected: {expected_output}\n"
        if actual_output:
            feedback_content += f"Actual: {actual_output}\n"
        if analysis_score is not None:
            feedback_content += f"Quality Score: {analysis_score}/10\n"
        if suggestions:
            feedback_content += f"Suggestions: {'; '.join(suggestions)}\n"
        
        # Prepare feedback context
        feedback_context = {
            "feedback_type": feedback_type,
            "analysis_score": analysis_score,
            "suggestions": suggestions or [],
            "iteration": iteration,
            "timestamp": datetime.now().isoformat()
        }
        
        if metadata:
            feedback_context.update(metadata)
        
        # Store as a memory item with high importance for learning
        memory = MemoryItem(
            id=feedback_id,
            input_text=input_text,
            response=expected_output,  # Store expected as the "correct" response
            domain=self.domain,
            task_type="feedback_learning",
            timestamp=datetime.now(),
            importance=0.9,  # High importance for learning
            context=feedback_context,
            tags=["feedback", feedback_type, "learning"]
        )
        
        # Store in memory and database
        self.memories[feedback_id] = memory
        await self._save_memory_to_db(memory)
        
        self.logger.info(f"Stored {feedback_type} feedback {feedback_id} with score {analysis_score}")
        return feedback_id

    async def get_feedback_history(
        self,
        feedback_type: Optional[str] = None,
        limit: int = 50
    ) -> List[MemoryItem]:
        """
        Get feedback history for analysis and learning.
        
        Args:
            feedback_type: Filter by feedback type
            limit: Maximum number of feedback items
            
        Returns:
            List of feedback memory items
        """
        feedback_memories = []
        
        for memory in self.memories.values():
            if memory.task_type == "feedback_learning":
                if feedback_type is None or memory.context.get("feedback_type") == feedback_type:
                    feedback_memories.append(memory)
        
        # Sort by timestamp, most recent first
        feedback_memories.sort(key=lambda m: m.timestamp, reverse=True)
        
        return feedback_memories[:limit]

    async def get_learning_patterns(self) -> Dict[str, Any]:
        """
        Analyze feedback patterns for learning insights.
        
        Returns:
            Dictionary with learning pattern analysis
        """
        feedback_history = await self.get_feedback_history()
        
        if not feedback_history:
            return {"total_feedback": 0, "patterns": {}}
        
        # Analyze patterns
        feedback_types = {}
        score_progression = []
        common_suggestions = {}
        
        for memory in feedback_history:
            # Count feedback types
            fb_type = memory.context.get("feedback_type", "unknown")
            feedback_types[fb_type] = feedback_types.get(fb_type, 0) + 1
            
            # Track score progression
            score = memory.context.get("analysis_score")
            if score is not None:
                score_progression.append({
                    "timestamp": memory.timestamp.isoformat(),
                    "score": score,
                    "iteration": memory.context.get("iteration")
                })
            
            # Count common suggestions
            suggestions = memory.context.get("suggestions", [])
            for suggestion in suggestions:
                # Extract key phrases from suggestions
                key_phrase = suggestion[:50]  # First 50 chars as key
                common_suggestions[key_phrase] = common_suggestions.get(key_phrase, 0) + 1
        
        # Calculate improvement trend
        improvement_trend = "stable"
        if len(score_progression) >= 2:
            recent_scores = [item["score"] for item in score_progression[:5]]
            older_scores = [item["score"] for item in score_progression[-5:]]
            
            if len(recent_scores) > 0 and len(older_scores) > 0:
                recent_avg = sum(recent_scores) / len(recent_scores)
                older_avg = sum(older_scores) / len(older_scores)
                
                if recent_avg > older_avg + 0.5:
                    improvement_trend = "improving"
                elif recent_avg < older_avg - 0.5:
                    improvement_trend = "declining"
        
        return {
            "total_feedback": len(feedback_history),
            "feedback_types": feedback_types,
            "score_progression": score_progression,
            "common_suggestions": dict(sorted(common_suggestions.items(), key=lambda x: x[1], reverse=True)[:10]),
            "improvement_trend": improvement_trend,
            "latest_score": score_progression[0]["score"] if score_progression else None
        } 