"""
S3-based memory backend for cloud deployments.

Uses S3 for storage and in-memory vectors for semantic search.
"""

import boto3
import json
import hashlib
import os
import numpy as np
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any, Optional
from .metadata import MemoryMetadata


class S3MemoryBackend:
    """
    Cloud memory backend using S3 for persistence and in-memory vectors for search.

    Storage structure:
    - MEMORY.md: Long-term memory file
    - logs/YYYY-MM-DD.md: Daily log files
    - vectors.json: Embeddings for semantic search
    """

    def __init__(
        self,
        org_id: str,
        project_name: str,
        workspace: str,
        agent_id: str,
        scope: str = "project",
        embedding_provider: str = "openai",
        embedding_model: str = "text-embedding-3-small"
    ):
        """
        Initialize S3 memory backend.

        Args:
            org_id: Organization ID from execution context
            project_name: Project name from deployment
            workspace: Workspace name (agent name without UUID)
            agent_id: Full agent ID with UUID for attribution
            scope: "project" or "global"
            embedding_provider: Provider for embeddings (only "openai" supported)
            embedding_model: Model name for embeddings
        """
        self.s3 = boto3.client('s3')
        self.bucket = f"daita-memory-{os.getenv('AWS_REGION', 'us-east-1')}"
        self.org_id = org_id
        self.project_name = project_name
        self.workspace = workspace
        self.agent_id = agent_id
        self.scope = scope

        # Build S3 prefix based on scope
        if scope == "global":
            self.prefix = f"orgs/{org_id}/global/workspaces/{workspace}/"
        else:
            self.prefix = f"orgs/{org_id}/projects/{project_name}/workspaces/{workspace}/"

        # Initialize storage components
        from .s3_storage import S3FileStorage
        self.storage = S3FileStorage(self.s3, self.bucket, self.prefix, agent_id)

        # Initialize embedding provider
        self.embedding_provider = embedding_provider
        self.embedding_model = embedding_model

        if embedding_provider == "openai":
            from openai import AsyncOpenAI
            api_key = os.getenv('OPENAI_API_KEY')
            if not api_key:
                raise ValueError("OPENAI_API_KEY environment variable required for cloud memory embeddings")
            self.embedder = AsyncOpenAI(api_key=api_key)
        else:
            raise NotImplementedError(f"Embedding provider {embedding_provider} not supported")

        # Load vectors into memory on initialization (lazy loading for performance)
        self.vectors = None
        self._vectors_loaded = False
        self._vectors_modified = False
        self._pending_vector_updates = 0  # Track number of updates since last flush

        # Memory count tracking for metrics
        self._pending_memory_count_increment = 0
        self._pending_memory_count_decrement = 0

    def _load_vectors_from_s3(self) -> Dict[str, Any]:
        """
        Load vectors.json from S3 into memory.

        Returns:
            Dict containing vectors and metadata
        """
        try:
            response = self.s3.get_object(
                Bucket=self.bucket,
                Key=f"{self.prefix}vectors.json"
            )
            data = json.loads(response['Body'].read().decode('utf-8'))
            print(f"Cloud Memory: Loaded {len(data.get('vectors', {}))} vectors from S3")
            return data
        except self.s3.exceptions.NoSuchKey:
            print("Cloud Memory: No existing vectors found, starting fresh")
            return {
                "vectors": {},
                "metadata": {
                    "created_at": datetime.now().isoformat(),
                    "workspace": self.workspace
                }
            }
        except Exception as e:
            print(f"Cloud Memory: Error loading vectors from S3: {e}")
            return {"vectors": {}, "metadata": {}}

    def _save_vectors_to_s3(self, force: bool = False):
        """
        Save vectors.json back to S3 (batched for performance).

        Batching strategy:
        - Only flush when: agent stops, explicit flush(), or 50+ pending updates
        - Reduces S3 PUT requests by up to 50x
        - Keeps vectors in memory during agent lifecycle

        Args:
            force: Force immediate save regardless of batch size
        """
        if not self._vectors_modified:
            return

        # Batch threshold: flush every 50 updates or when forced
        if not force and self._pending_vector_updates < 50:
            return

        try:
            self.s3.put_object(
                Bucket=self.bucket,
                Key=f"{self.prefix}vectors.json",
                Body=json.dumps(self.vectors, indent=2).encode('utf-8'),
                ContentType='application/json'
            )
            print(f"Cloud Memory: Flushed {self._pending_vector_updates} vector updates to S3")
            self._vectors_modified = False
            self._pending_vector_updates = 0
        except Exception as e:
            print(f"Cloud Memory: Error saving vectors to S3: {e}")
            raise

    async def flush(self):
        """Force immediate flush of pending vectors to S3."""
        self._save_vectors_to_s3(force=True)

    def get_pending_metrics(self) -> dict:
        """
        Return accumulated memory metrics for the Lambda to forward to the API.

        Returns:
            Dict with memory_count_delta and memory_retrieval_count
        """
        return {
            "memory_count_delta": self._pending_memory_count_increment - self._pending_memory_count_decrement,
            "memory_retrieval_count": getattr(self, '_retrieval_count', 0)
        }

    @property
    def search(self):
        """Expose self as the search interface (duck-type compatibility with local backend)."""
        return self

    async def remember(
        self,
        content: str,
        category: Optional[str] = None,
        metadata: Optional[MemoryMetadata] = None
    ) -> Dict[str, Any]:
        """
        Store a memory: appends to daily log and indexes in vectors.
        MEMORY.md is managed exclusively by the curator.

        Args:
            content: Content to remember
            category: Optional category tag
            metadata: Optional MemoryMetadata (created with defaults if not provided)

        Returns:
            Result dict with status and chunk_id
        """
        await self.storage.append_to_daily_log(content, category)

        if metadata is None:
            metadata = MemoryMetadata(
                content=content,
                importance=0.5,
                source='agent_inferred',
                category=category
            )

        chunk_id = await self.store_chunk(content, metadata)

        return {
            "status": "success",
            "message": "Memory stored and indexed",
            "chunk_id": chunk_id,
            "indexed": True
        }

    async def store_chunk(
        self,
        content: str,
        metadata: MemoryMetadata,
        chunk_id: Optional[str] = None
    ) -> str:
        """
        Store a chunk directly into the in-memory vector store.
        Primary write path for both the remember tool and the curator.

        Returns:
            chunk_id of the stored chunk
        """
        if not self._vectors_loaded:
            self.vectors = self._load_vectors_from_s3()
            self._vectors_loaded = True

        if chunk_id is None:
            chunk_id = hashlib.md5(
                f"direct:{content[:100]}:{metadata.created_at.isoformat()}".encode()
            ).hexdigest()

        # Incremental: skip if already exists
        if chunk_id in self.vectors.get('vectors', {}):
            return chunk_id

        embedding = await self._generate_embedding(content)

        self.vectors['vectors'][chunk_id] = {
            'content': content,
            'embedding': embedding,
            'metadata': metadata.to_dict(),
            'agent_id': self.agent_id
        }
        self._vectors_modified = True
        self._pending_vector_updates += 1
        self._pending_memory_count_increment += 1

        self._save_vectors_to_s3(force=False)
        return chunk_id

    def track_access(self, chunk_ids: List[str]):
        """
        Increment access_count and update last_accessed for recalled chunks.
        No-op if vectors not yet loaded.
        """
        if not self._vectors_loaded or not chunk_ids:
            return

        now = datetime.now().isoformat()
        for chunk_id in chunk_ids:
            vec_data = self.vectors.get('vectors', {}).get(chunk_id)
            if vec_data:
                meta = vec_data.get('metadata', {})
                meta['access_count'] = meta.get('access_count', 0) + 1
                meta['last_accessed'] = now
                vec_data['metadata'] = meta
                self._vectors_modified = True

    async def update_memory(
        self,
        query: str,
        new_content: str,
        importance: float = 0.5
    ) -> Dict[str, Any]:
        """
        Find memories matching query, delete them, and store updated content.
        """
        if not self._vectors_loaded:
            self.vectors = self._load_vectors_from_s3()
            self._vectors_loaded = True

        matches = await self.recall(query, limit=5, score_threshold=0.7)
        if not matches:
            return {"status": "not_found", "updated": 0, "message": "No matching memories found"}

        chunk_ids = [m['chunk_id'] for m in matches if m.get('chunk_id')]
        for cid in chunk_ids:
            self.vectors['vectors'].pop(cid, None)
        self._vectors_modified = True
        self._pending_memory_count_decrement += len(chunk_ids)

        metadata = MemoryMetadata(
            content=new_content,
            importance=importance,
            source='agent_inferred'
        )
        await self.store_chunk(new_content, metadata)

        self._save_vectors_to_s3(force=False)
        return {
            "status": "success",
            "updated": len(chunk_ids),
            "message": f"Replaced {len(chunk_ids)} memories with updated content"
        }

    async def regenerate_memory_md(self, min_importance: float = 0.4):
        """
        Regenerate MEMORY.md in S3 from the in-memory vector store.
        Groups by category, sorted by importance. Curator-only write path.
        """
        if not self._vectors_loaded:
            self.vectors = self._load_vectors_from_s3()
            self._vectors_loaded = True

        by_category: Dict[str, List] = {}
        for vec_data in self.vectors.get('vectors', {}).values():
            meta = vec_data.get('metadata', {})
            importance = meta.get('importance', 0.5)
            if importance < min_importance:
                continue
            category = meta.get('category') or 'General'
            by_category.setdefault(category, []).append((importance, vec_data['content'].strip()))

        lines = ["# Long-Term Memory\n"]
        for category in sorted(by_category.keys()):
            entries = sorted(by_category[category], reverse=True)
            lines.append(f"\n## {category.title()}\n")
            for _, content in entries:
                lines.append(f"- {content}\n")

        self.s3.put_object(
            Bucket=self.bucket,
            Key=f"{self.prefix}MEMORY.md",
            Body=''.join(lines).encode('utf-8'),
            ContentType='text/markdown'
        )

    def _apply_score_adjustments(
        self,
        base_score: float,
        metadata_dict: Optional[Dict]
    ) -> float:
        """
        Apply importance boost and temporal decay to a base relevance score.
        """
        if not metadata_dict:
            return base_score

        importance = metadata_dict.get('importance', 0.5)
        importance_boost = (importance - 0.5) * 0.2

        decay = 1.0
        if not metadata_dict.get('pinned', False):
            created_at_str = metadata_dict.get('created_at')
            if created_at_str:
                try:
                    created_at = datetime.fromisoformat(created_at_str)
                    age_days = (datetime.now() - created_at).days
                    decay = max(0.7, 1.0 - (age_days / 365) * 0.3)
                except Exception:
                    pass

        return min((base_score + importance_boost) * decay, 1.0)

    async def recall(
        self,
        query: str,
        limit: int = 5,
        score_threshold: float = 0.6,
        strategy: str = "hybrid",
        min_importance: Optional[float] = None,
        max_importance: Optional[float] = None
    ) -> List[Dict[str, Any]]:
        """
        Search memories using hybrid semantic + keyword search.

        Args:
            query: Search query (natural language)
            limit: Maximum number of results
            score_threshold: Minimum relevance score (0-1, default 0.6)
            strategy: Retrieval strategy
                - "hybrid": Combine semantic + keyword (default, recommended)
                - "semantic": Pure vector similarity
                - "keyword": Pure BM25 keyword matching
            min_importance: Filter to memories with importance >= this value
            max_importance: Filter to memories with importance <= this value

        Returns:
            List of matching memory excerpts with scores and score breakdown
        """
        # Lazy load vectors if not already loaded
        if not self._vectors_loaded:
            self.vectors = self._load_vectors_from_s3()
            self._vectors_loaded = True

        # If no vectors, return empty results
        if not self.vectors.get('vectors'):
            return []

        # Route to appropriate strategy
        if strategy == "semantic":
            results = await self._semantic_recall(query, limit, score_threshold)
        elif strategy == "keyword":
            results = await self._keyword_recall(query, limit, score_threshold)
        else:  # hybrid (default)
            results = await self._hybrid_recall(query, limit, score_threshold)

        # Apply importance filtering
        if min_importance is not None or max_importance is not None:
            results = self._filter_by_importance(results, min_importance, max_importance)

        return results

    async def _semantic_recall(
        self,
        query: str,
        limit: int,
        score_threshold: float
    ) -> List[Dict[str, Any]]:
        """
        Pure semantic search using vector embeddings.

        Args:
            query: Search query
            limit: Maximum results
            score_threshold: Minimum score

        Returns:
            List of matching memories
        """
        query_embedding = await self._generate_embedding(query)

        results = []
        for chunk_id, vector_data in self.vectors.get('vectors', {}).items():
            stored_embedding = vector_data['embedding']
            similarity = self._cosine_similarity(query_embedding, stored_embedding)

            metadata_dict = vector_data.get('metadata', {})
            if not metadata_dict:
                metadata_dict = {
                    'content': vector_data['content'],
                    'importance': 0.5,
                    'source': 'agent_inferred',
                    'category': vector_data.get('section'),
                    'created_at': vector_data.get('created_at'),
                    'pinned': False,
                    'access_count': 0
                }

            adjusted = self._apply_score_adjustments(similarity, metadata_dict)

            if adjusted >= score_threshold:
                results.append({
                    'content': vector_data['content'],
                    'relevance_score': float(adjusted),
                    'score': float(adjusted),
                    'chunk_id': chunk_id,
                    'agent_id': vector_data.get('agent_id'),
                    'metadata': metadata_dict
                })

        results.sort(key=lambda x: x['relevance_score'], reverse=True)
        return results[:limit]

    async def _keyword_recall(
        self,
        query: str,
        limit: int,
        score_threshold: float
    ) -> List[Dict[str, Any]]:
        """
        Pure keyword search using BM25.

        Args:
            query: Search query
            limit: Maximum results
            score_threshold: Minimum score

        Returns:
            List of matching memories
        """
        from .keyword_search import BM25Scorer
        from .text_utils import extract_keywords

        keywords = extract_keywords(query)
        documents = [v['content'] for v in self.vectors['vectors'].values()]
        bm25 = BM25Scorer(documents)

        results = []
        for chunk_id, vector_data in self.vectors['vectors'].items():
            content = vector_data['content']
            keyword_raw_score = bm25.score(keywords, content)
            keyword_score = bm25.normalize_score(keyword_raw_score)

            metadata_dict = vector_data.get('metadata', {})
            if not metadata_dict:
                metadata_dict = {
                    'content': content,
                    'importance': 0.5,
                    'source': 'agent_inferred',
                    'category': vector_data.get('section'),
                    'created_at': vector_data.get('created_at'),
                    'pinned': False,
                    'access_count': 0
                }

            adjusted = self._apply_score_adjustments(keyword_score, metadata_dict)

            if adjusted >= score_threshold:
                results.append({
                    'content': content,
                    'relevance_score': float(adjusted),
                    'score': float(adjusted),
                    'chunk_id': chunk_id,
                    'agent_id': vector_data.get('agent_id'),
                    'metadata': metadata_dict
                })

        results.sort(key=lambda x: x['relevance_score'], reverse=True)
        return results[:limit]

    async def _hybrid_recall(
        self,
        query: str,
        limit: int,
        score_threshold: float
    ) -> List[Dict[str, Any]]:
        """
        Hybrid search combining semantic and keyword retrieval.

        Uses weighted combination of semantic (60%) and keyword (40%) scores,
        with bonuses for exact phrase matches and consensus.

        Args:
            query: Search query
            limit: Maximum results
            score_threshold: Minimum score

        Returns:
            List of matching memories with score breakdowns
        """
        from .keyword_search import BM25Scorer
        from .text_utils import extract_keywords, contains_exact_phrase

        # Extract keywords for BM25
        keywords = extract_keywords(query)

        # Build BM25 index
        documents = [v['content'] for v in self.vectors['vectors'].values()]
        bm25 = BM25Scorer(documents)

        # Generate query embedding for semantic search
        query_embedding = await self._generate_embedding(query)

        # Score each memory
        results = []
        for chunk_id, vector_data in self.vectors['vectors'].items():
            content = vector_data['content']
            stored_embedding = vector_data['embedding']

            # Semantic score (cosine similarity)
            semantic_score = self._cosine_similarity(query_embedding, stored_embedding)

            # Keyword score (BM25)
            keyword_raw_score = bm25.score(keywords, content)
            keyword_score = bm25.normalize_score(keyword_raw_score)

            # Base hybrid score (weighted combination)
            base_score = (0.6 * semantic_score) + (0.4 * keyword_score)

            # Bonus: Exact phrase match
            phrase_bonus = 0.15 if contains_exact_phrase(query, content) else 0.0

            # Bonus: Consensus (both methods agree)
            consensus_bonus = 0.10 if semantic_score > 0.5 and keyword_score > 0.5 else 0.0

            raw_score = min(base_score + phrase_bonus + consensus_bonus, 1.0)

            metadata_dict = vector_data.get('metadata', {})
            if not metadata_dict:
                metadata_dict = {
                    'content': content,
                    'importance': 0.5,
                    'source': 'agent_inferred',
                    'category': vector_data.get('section'),
                    'created_at': vector_data.get('created_at'),
                    'pinned': False,
                    'access_count': 0
                }

            adjusted = self._apply_score_adjustments(raw_score, metadata_dict)

            if adjusted >= score_threshold:
                results.append({
                    'content': content,
                    'relevance_score': float(adjusted),
                    'score': float(adjusted),
                    'chunk_id': chunk_id,
                    'agent_id': vector_data.get('agent_id'),
                    'metadata': metadata_dict,
                    'score_breakdown': {
                        'semantic': float(semantic_score),
                        'keyword': float(keyword_score),
                        'phrase_bonus': float(phrase_bonus),
                        'consensus_bonus': float(consensus_bonus)
                    }
                })

        # Sort by score and return top results
        results.sort(key=lambda x: x['relevance_score'], reverse=True)
        return results[:limit]

    async def read_memory(self, file_path: str) -> str:
        """
        Read a memory file from S3.

        Args:
            file_path: S3 URI or relative path

        Returns:
            File contents
        """
        return await self.storage.read_file(file_path)

    async def list_memories(self) -> List[Dict[str, Any]]:
        """
        List all memory files in S3.

        Returns:
            List of file metadata dicts
        """
        return await self.storage.list_files()

    async def _generate_embedding(self, text: str) -> List[float]:
        """
        Generate embedding for text using configured provider.

        Args:
            text: Text to embed

        Returns:
            Embedding vector as list of floats
        """
        response = await self.embedder.embeddings.create(
            model=self.embedding_model,
            input=text
        )
        return response.data[0].embedding

    @staticmethod
    def _cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
        """
        Compute cosine similarity between two vectors.

        Args:
            vec1: First vector
            vec2: Second vector

        Returns:
            Similarity score (0-1)
        """
        v1 = np.array(vec1)
        v2 = np.array(vec2)
        return float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))

    @staticmethod
    def _filter_by_importance(
        results: List[Dict[str, Any]],
        min_importance: Optional[float],
        max_importance: Optional[float]
    ) -> List[Dict[str, Any]]:
        """
        Filter results by importance range.

        Args:
            results: Search results to filter
            min_importance: Minimum importance (inclusive)
            max_importance: Maximum importance (inclusive)

        Returns:
            Filtered results
        """
        filtered = []
        for result in results:
            metadata = result.get('metadata', {})
            importance = metadata.get('importance', 0.5)

            # Apply filters
            if min_importance is not None and importance < min_importance:
                continue
            if max_importance is not None and importance > max_importance:
                continue

            filtered.append(result)

        return filtered
