from typing import List, Optional, TYPE_CHECKING
from datetime import datetime
import uuid
from analogai.foundation.extensions.conversation.repositories.conversation_repository import ConversationRepository
from analogai.foundation.extensions.conversation.conversation import Conversation
from analogai.foundation.extensions.conversation.turn import Turn
from analogai.foundation.modules.user.user import User
from analogai.foundation.modules.agent.agent import Agent
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.infrastructure.storage.mongodb.mongodb_storage_adapter_impl import MongoDBStorageAdapterImpl

if TYPE_CHECKING:
    from analogai.foundation.modules.agent.agent_repository import AgentRepository


class MongoDBConversationRepositoryImpl(ConversationRepository):
    CONVERSATIONS_COLLECTION = "conversations"
    TURNS_COLLECTION = "turns"

    def __init__(self, storage_adapter: Optional[StorageAdapter] = None, agent_repository: Optional["AgentRepository"] = None):
        self.storage_adapter = storage_adapter if storage_adapter is not None else MongoDBStorageAdapterImpl()
        self._agent_repository = agent_repository

    def create_conversation(self, user_id: str, agent_id: str) -> Conversation:
        conversation_id = str(uuid.uuid4())
        now = datetime.now()

        user = User(id=user_id, name="", email="")
        agent = self._get_agent(agent_id)

        conversation = Conversation(
            id=conversation_id,
            user=user,
            agent=agent,
            created_at=now,
            updated_at=now
        )

        data = self._conversation_to_doc(conversation)
        stored_data = self.storage_adapter.store(self.CONVERSATIONS_COLLECTION, data)
        if not stored_data:
            raise ValueError(f"Failed to store conversation {conversation.id}")

        return conversation

    def find_by_id(self, conversation_id: str) -> Optional[Conversation]:
        doc = self.storage_adapter.retrieve_by_id(self.CONVERSATIONS_COLLECTION, conversation_id)
        if not doc:
            return None

        conversation = self._doc_to_conversation(doc)
        conversation.turns = self._get_turns_for_conversation(conversation_id)
        return conversation

    def find_by_agent(self, agent_id: str, page: int = 1, page_size: int = 10) -> List[Conversation]:
        docs = self.storage_adapter.retrieve_by_attributes(
            self.CONVERSATIONS_COLLECTION,
            {"agent_id": agent_id}
        )

        conversations = []
        for doc in docs:
            conversation = self._doc_to_conversation(doc)
            conversation.turns = self._get_turns_for_conversation(conversation.id)
            conversations.append(conversation)

        conversations.sort(key=lambda x: x.updated_at, reverse=True)

        start_idx = (page - 1) * page_size
        end_idx = start_idx + page_size

        return conversations[start_idx:end_idx]

    def find_latest_by_user_and_agent(self, user_id: str, agent_id: str) -> Optional[Conversation]:
        docs = self.storage_adapter.retrieve_by_attributes(
            self.CONVERSATIONS_COLLECTION,
            {"user_id": user_id, "agent_id": agent_id}
        )
        if not docs:
            return None
        try:
            docs.sort(key=lambda d: d.get("updated_at", ""), reverse=True)
        except Exception:
            pass
        doc = docs[0]
        conversation = self._doc_to_conversation(doc)
        conversation.turns = self._get_turns_for_conversation(conversation.id)
        return conversation

    def add_turn(self, conversation_id: str, turn: Turn) -> Turn:
        conversation = self.find_by_id(conversation_id)
        if not conversation:
            raise ValueError(f"Conversation {conversation_id} not found")

        conversation.add_turn(turn)

        turn_data = self._turn_to_doc(turn, conversation_id)
        stored_turn_data = self.storage_adapter.store(self.TURNS_COLLECTION, turn_data)
        if not stored_turn_data:
            raise ValueError(f"Failed to store turn {turn.id}")

        conversation_data = self._conversation_to_doc(conversation)
        stored_conversation_data = self.storage_adapter.store(self.CONVERSATIONS_COLLECTION, conversation_data)
        if not stored_conversation_data:
            raise ValueError(f"Failed to update conversation {conversation.id}")

        return turn

    def delete_conversation(self, conversation_id: str) -> bool:
        turns = self._get_turns_for_conversation(conversation_id)
        for turn in turns:
            self.storage_adapter.delete(self.TURNS_COLLECTION, turn.id)
        return self.storage_adapter.delete(self.CONVERSATIONS_COLLECTION, conversation_id)

    def _conversation_to_doc(self, conversation: Conversation) -> dict:
        return {
            "id": conversation.id,
            "user_id": conversation.user.id,
            "agent_id": conversation.agent.id,
            "created_at": conversation.created_at.isoformat(),
            "updated_at": conversation.updated_at.isoformat()
        }

    def _get_agent(self, agent_id: str) -> Agent:
        if self._agent_repository:
            agent = self._agent_repository.find_by_id(agent_id)
            if agent:
                return agent

        agent_doc = self.storage_adapter.retrieve_by_id("agents", agent_id)
        if agent_doc and agent_doc.get("creator_id"):
            agent_creator = User(id=agent_doc["creator_id"], name="", email="")
        else:
            agent_creator = User(id="system", name="", email="")
        return Agent(
            id=agent_id,
            name=agent_doc.get("name", "") if agent_doc else "",
            knowledge_base=agent_doc.get("knowledge_base", "") if agent_doc else "",
            creator=agent_creator
        )

    def _doc_to_conversation(self, doc: dict) -> Conversation:
        user = User(id=doc["user_id"], name="", email="")
        agent = self._get_agent(doc["agent_id"])

        conversation = Conversation(
            id=doc["id"],
            user=user,
            agent=agent,
            created_at=datetime.fromisoformat(doc["created_at"]),
            updated_at=datetime.fromisoformat(doc["updated_at"])
        )

        return conversation

    def _turn_to_doc(self, turn: Turn, conversation_id: str) -> dict:
        return {
            "id": turn.id,
            "conversation_id": conversation_id,
            "user_id": turn.user_id,
            "agent_id": turn.agent_id,
            "user_message": turn.user_message,
            "agent_response": turn.agent_response,
            "timestamp": turn.timestamp.isoformat(),
        }

    def _doc_to_turn(self, doc: dict) -> Turn:
        return Turn(
            id=doc["id"],
            user_id=doc["user_id"],
            agent_id=doc["agent_id"],
            user_message=doc["user_message"],
            agent_response=doc["agent_response"],
            timestamp=datetime.fromisoformat(doc["timestamp"]),
        )

    def _get_turns_for_conversation(self, conversation_id: str) -> List[Turn]:
        docs = self.storage_adapter.retrieve_by_attributes(
            self.TURNS_COLLECTION,
            {"conversation_id": conversation_id}
        )
        turns = [self._doc_to_turn(doc) for doc in docs]
        turns.sort(key=lambda x: x.timestamp)
        return turns
