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

Uses append-only architecture for optimal performance:
- Individual memory files: memories/{timestamp}-{uuid}.json
- MEMORY.md: Reconstructed on-demand with caching
- Daily logs: Keep existing pattern (relatively small files)
"""

import boto3
import os
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional


class S3FileStorage:
    """S3-based memory storage for cloud deployments."""

    def __init__(self, s3_client, bucket: str, prefix: str, agent_id: Optional[str] = None):
        """
        Initialize S3 file storage.

        Args:
            s3_client: Boto3 S3 client instance
            bucket: S3 bucket name
            prefix: S3 key prefix (workspace path)
            agent_id: Optional agent identifier for attribution in shared workspaces
        """
        self.s3 = s3_client
        self.bucket = bucket
        self.prefix = prefix.rstrip('/') + '/'  # Ensure trailing slash
        self.agent_id = agent_id

    async def append_to_daily_log(self, content: str, category: Optional[str] = None) -> str:
        """
        Append to today's daily log in S3.

        Args:
            content: Content to append
            category: Optional category tag

        Returns:
            S3 URI of the log file
        """
        today = datetime.now().strftime("%Y-%m-%d")
        log_key = f"{self.prefix}logs/{today}.md"

        # Download existing log (if any)
        try:
            response = self.s3.get_object(Bucket=self.bucket, Key=log_key)
            existing_content = response['Body'].read().decode('utf-8')
        except self.s3.exceptions.NoSuchKey:
            existing_content = f"# Daily Log - {today}\n\n"

        # Format new entry with agent attribution
        timestamp = datetime.now().strftime("%H:%M:%S")
        entry = f"\n## {timestamp}"

        if self.agent_id:
            entry += f" [{self.agent_id}]"

        if category:
            entry += f" [{category}]"

        entry += f"\n\n{content}\n"

        # Upload updated log
        new_content = existing_content + entry
        self.s3.put_object(
            Bucket=self.bucket,
            Key=log_key,
            Body=new_content.encode('utf-8'),
            ContentType='text/markdown'
        )

        return f"s3://{self.bucket}/{log_key}"

    async def write_individual_memory(
        self,
        content: str,
        metadata: Dict[str, Any],
        section: Optional[str] = None
    ) -> str:
        """
        Write individual memory file (append-only pattern).

        Storage pattern:
            memories/{timestamp}-{uuid}.json

        Benefits:
            - No GET request needed (was causing 400ms latency)
            - Constant-time writes (doesn't degrade as memories grow)
            - 10x cost reduction at scale
            - Better concurrency (no read-modify-write)

        Args:
            content: Memory content
            metadata: Structured metadata dict
            section: Optional section for organization

        Returns:
            S3 URI of the memory file
        """
        # Generate unique filename
        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f")[:20]  # Include microseconds
        memory_id = str(uuid.uuid4())[:8]
        memory_key = f"{self.prefix}memories/{timestamp}-{memory_id}.json"

        # Create memory document
        memory_data = {
            "content": content,
            "metadata": metadata,
            "section": section,
            "agent_id": self.agent_id,
            "created_at": datetime.now().isoformat()
        }

        # Single PUT request (no GET needed - this is the optimization!)
        self.s3.put_object(
            Bucket=self.bucket,
            Key=memory_key,
            Body=json.dumps(memory_data, indent=2).encode('utf-8'),
            ContentType='application/json'
        )

        # Invalidate MEMORY.md cache
        await self._invalidate_memory_cache()

        return f"s3://{self.bucket}/{memory_key}"

    async def reconstruct_memory_md(self, use_cache: bool = True) -> str:
        """
        Reconstruct MEMORY.md from individual memory files.

        Uses caching to avoid expensive reconstruction on every read.
        Cache TTL: 5 minutes (good balance between freshness and performance)

        Args:
            use_cache: Whether to use cached version if available

        Returns:
            Markdown content of MEMORY.md
        """
        cache_key = f"{self.prefix}MEMORY.md.cache"

        # Check cache
        if use_cache:
            try:
                response = self.s3.get_object(Bucket=self.bucket, Key=cache_key)
                cache_data = json.loads(response['Body'].read().decode('utf-8'))

                # Check if cache is fresh (< 5 minutes old)
                cache_time = datetime.fromisoformat(cache_data['cached_at'])
                if (datetime.now() - cache_time) < timedelta(minutes=5):
                    return cache_data['content']
            except (self.s3.exceptions.NoSuchKey, json.JSONDecodeError, KeyError):
                pass  # Cache miss or corrupted, proceed to reconstruction

        # List all memory files
        memories = []
        try:
            paginator = self.s3.get_paginator('list_objects_v2')
            pages = paginator.paginate(
                Bucket=self.bucket,
                Prefix=f"{self.prefix}memories/"
            )

            for page in pages:
                for obj in page.get('Contents', []):
                    if obj['Key'].endswith('.json'):
                        try:
                            # Read memory file
                            mem_response = self.s3.get_object(
                                Bucket=self.bucket,
                                Key=obj['Key']
                            )
                            memory_data = json.loads(mem_response['Body'].read().decode('utf-8'))
                            memories.append(memory_data)
                        except (json.JSONDecodeError, KeyError):
                            # Skip corrupted files
                            continue

        except Exception as e:
            print(f"Error listing memory files: {e}")
            return "# Long-Term Memory\n\n(Error loading memories)"

        if not memories:
            return "# Long-Term Memory\n\n"

        # Sort by created_at
        memories.sort(key=lambda m: m.get('created_at', ''))

        # Generate markdown grouped by section
        md_content = "# Long-Term Memory\n\n"

        # Group by section
        sections = {}
        for mem in memories:
            section = mem.get('section') or 'General'
            if section not in sections:
                sections[section] = []
            sections[section].append(mem)

        # Build markdown
        for section, mems in sorted(sections.items()):
            md_content += f"\n## {section}\n\n"
            for mem in mems:
                # Add agent attribution
                if mem.get('agent_id'):
                    md_content += f"*[{mem['agent_id']}]* "

                # Add content
                md_content += f"{mem['content']}\n"

        # Cache the result
        cache_data = {
            "content": md_content,
            "cached_at": datetime.now().isoformat(),
            "memory_count": len(memories)
        }

        try:
            self.s3.put_object(
                Bucket=self.bucket,
                Key=cache_key,
                Body=json.dumps(cache_data).encode('utf-8'),
                ContentType='application/json'
            )
        except Exception as e:
            print(f"Warning: Failed to cache MEMORY.md: {e}")

        return md_content

    async def _invalidate_memory_cache(self):
        """Invalidate MEMORY.md cache by deleting cache file."""
        cache_key = f"{self.prefix}MEMORY.md.cache"
        try:
            self.s3.delete_object(Bucket=self.bucket, Key=cache_key)
        except Exception:
            pass  # Cache might not exist, that's fine

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

        Special handling for MEMORY.md: Reconstructs on-demand from individual
        memory files with caching.

        Args:
            file_path: S3 URI (s3://bucket/key) or relative path within workspace

        Returns:
            File contents
        """
        # Handle both full S3 paths and relative paths
        if file_path.startswith('s3://'):
            # Parse S3 URL
            parts = file_path.replace('s3://', '').split('/', 1)
            bucket = parts[0]
            key = parts[1] if len(parts) > 1 else ''
        else:
            # Relative path within workspace
            bucket = self.bucket
            key = f"{self.prefix}{file_path}"

        # Special case: MEMORY.md is reconstructed on-demand
        if key.endswith('MEMORY.md'):
            return await self.reconstruct_memory_md(use_cache=True)

        # Regular file read
        try:
            response = self.s3.get_object(Bucket=bucket, Key=key)
            return response['Body'].read().decode('utf-8')
        except self.s3.exceptions.NoSuchKey:
            return f"File not found: {file_path}"

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

        Returns:
            List of file metadata dicts
        """
        files = []

        # List objects with workspace prefix
        try:
            response = self.s3.list_objects_v2(
                Bucket=self.bucket,
                Prefix=self.prefix
            )

            for obj in response.get('Contents', []):
                key = obj['Key']

                # Skip vectors.json (internal) and metadata.json
                if key.endswith('vectors.json') or key.endswith('metadata.json'):
                    continue

                # Determine file type
                if key.endswith('MEMORY.md'):
                    file_type = 'long_term'
                elif '/logs/' in key and key.endswith('.md'):
                    file_type = 'daily_log'
                else:
                    continue  # Skip unknown files

                files.append({
                    'path': f"s3://{self.bucket}/{key}",
                    'type': file_type,
                    'size': obj['Size'],
                    'modified': obj['LastModified'].isoformat()
                })

        except Exception as e:
            print(f"Error listing S3 files: {e}")

        return files
