"""
Firestore-backed implementation of the MemoryStore interface.
"""
from typing import List, Dict, Any, Optional
from datetime import datetime, timezone

from google.cloud import firestore
from google.cloud.firestore_v1.async_transaction import async_transactional

from contracts.common import EmmaConversation, MemoryRecord, AuditEvent
from .memory_store import MemoryStore

class FirestoreMemoryStore(MemoryStore):
    """
    Implements the MemoryStore interface using Google Cloud Firestore.
    """
    def __init__(self, project_id: str, client: Optional[firestore.AsyncClient] = None):
        self._project_id = project_id
        # Client can be injected for testing, otherwise it will be created on demand.
        self._client = client 

    def _get_client(self) -> firestore.AsyncClient:
        if not self._client:
            self._client = firestore.AsyncClient(project=self._project_id)
        return self._client

    async def create_conversation(self, owner_id: str, workspace_id: str, created_by: str) -> EmmaConversation:
        db = self._get_client()
        convo = EmmaConversation(owner_id=owner_id, workspace_id=workspace_id, created_by=created_by)
        doc_ref = db.collection("emma_conversations").document(convo.conversation_id)
        await doc_ref.set(convo.model_dump(mode='json'))
        return convo

    async def save_conversation(self, conversation: EmmaConversation) -> None:
        if not conversation.owner_id:
            raise ValueError("Conversation must have an owner_id.")

        db = self._get_client()
        doc_ref = db.collection("emma_conversations").document(
            conversation.conversation_id
        )
        conversation_data = conversation.model_dump(mode="json")

        @async_transactional
        async def _save(transaction) -> None:
            snapshot = await doc_ref.get(transaction=transaction)

            if not snapshot.exists:
                transaction.create(doc_ref, conversation_data)
                return

            existing_data = snapshot.to_dict() or {}

            if existing_data.get("owner_id") != conversation.owner_id:
                raise PermissionError(
                    "Cannot modify a conversation owned by another owner."
                )

            if existing_data.get("workspace_id") != conversation.workspace_id:
                raise PermissionError(
                    "Cannot move a conversation to another workspace."
                )

            conversation_data["owner_id"] = existing_data["owner_id"]
            conversation_data["workspace_id"] = existing_data["workspace_id"]

            if existing_data.get("created_at") is not None:
                conversation_data["created_at"] = existing_data["created_at"]

            transaction.update(doc_ref, conversation_data)

        await _save(db.transaction())

    async def get_conversation(self, conversation_id: str, owner_id: str) -> Optional[EmmaConversation]:
        db = self._get_client()
        doc_ref = db.collection("emma_conversations").document(conversation_id)
        doc = await doc_ref.get()
        if not doc.exists:
            return None
        
        convo = EmmaConversation(**doc.to_dict())
        # Security: Enforce ownership
        if convo.owner_id != owner_id:
            return None # Or raise an exception
        return convo

    async def append_message(self, conversation_id: str, message: Dict[str, Any]) -> None:
        db = self._get_client()
        # Note: message_id should be auto-generated by Firestore for ordering
        messages_ref = db.collection("emma_conversations").document(conversation_id).collection("messages")
        await messages_ref.add(message)
        # Update last_updated_at on parent conversation document
        convo_ref = db.collection("emma_conversations").document(conversation_id)
        await convo_ref.update({"last_updated_at": datetime.now(timezone.utc)})

    async def list_messages(self, conversation_id: str, limit: int = 50) -> List[Dict[str, Any]]:
        db = self._get_client()
        messages_ref = db.collection("emma_conversations").document(conversation_id).collection("messages")
        docs = messages_ref.order_by("created_at", direction=firestore.Query.DESCENDING).limit(limit).stream()
        return [doc.to_dict() async for doc in docs][::-1] # Reverse to get chronological order

    async def record_audit_event(self, event: AuditEvent) -> str:
        db = self._get_client()
        doc_ref = db.collection("emma_audit_events").document(event.event_id)
        # Use create() to ensure the document does not already exist, guaranteeing append-only.
        await doc_ref.create(event.model_dump(mode='json'))
        return event.event_id

    # --- Placeholder methods not fully implemented in Phase 2B ---

    async def save_conversation_summary(self, conversation_id: str, summary: str) -> None:
        # In a real implementation, this would update a summary field
        pass

    async def store_memory(self, record: MemoryRecord) -> None:
        # This would write to the emma_memories collection
        pass

    async def recall_memory(self, embedding: List[float], owner_id: str, scopes: List[str]) -> List[MemoryRecord]:
        # This would perform a query against a vector database, which is out of scope.
        return []

    async def get_profile(self, owner_id: str) -> Optional[Dict[str, Any]]:
        pass

    async def update_profile(self, owner_id: str, profile_data: Dict[str, Any]) -> None:
        pass