# coding: utf-8
# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
from datetime import datetime, timezone
from typing import Any, List, Tuple

from openjiuwen.core.foundation.llm import Model
from openjiuwen.core.memory.manage.mem_model.semantic_store import SemanticStore
from openjiuwen.core.memory.common.base import generate_idx_name, parse_memory_hit_infos
from openjiuwen.core.memory.manage.index.base_memory_manager import BaseMemoryManager
from openjiuwen.core.memory.manage.mem_model.memory_unit import SummaryUnit, BaseMemoryUnit, MemoryType
from openjiuwen.core.memory.manage.mem_model.user_mem_store import UserMemStore
from openjiuwen.core.common.exception.codes import StatusCode
from openjiuwen.core.common.exception.errors import build_error
from openjiuwen.core.common.logging import memory_logger
from openjiuwen.core.common.logging.events import LogEventType


class SummaryManager(BaseMemoryManager):
    """Manages summary memory CRUD with encryption and vector storage"""
    def __init__(self,
                 semantic_recall_instance: SemanticStore,
                 user_mem_store: UserMemStore,
                 crypto_key: bytes):
        self.mem_store = user_mem_store
        self.semantic_recall = semantic_recall_instance
        self.crypto_key = crypto_key

    async def add(self, memory: BaseMemoryUnit, llm: Tuple[str, Model] | None = None):
        """add memory."""
        if not isinstance(memory, SummaryUnit):
            raise build_error(
                StatusCode.MEMORY_ADD_MEMORY_EXECUTION_ERROR,
                memory_type="summary",
                error_msg="summary add Must pass SummaryUnit class",
            )

        vector_success = await self._add_summary_memory_to_vector(summary_unit=memory, scope_id=memory.scope_id)
        if not vector_success:
            raise build_error(
                StatusCode.MEMORY_ADD_MEMORY_EXECUTION_ERROR,
                memory_type="summary",
                error_msg="summary add to vector store failed",
            )
        await self._add_summary_memory_to_mem_store(memory)

    async def update(self, user_id: str, scope_id: str, mem_id: str, new_memory: str, **kwargs):
        """update memory by its id."""
        time = datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')
        encrypt_new_memory = BaseMemoryManager.encrypt_memory_if_needed(key=self.crypto_key, plaintext=new_memory)
        new_data = {'mem': encrypt_new_memory, 'time': time}
        await self.mem_store.update(mem_id=mem_id, user_id=user_id, scope_id=scope_id, data=new_data)
        table_name = generate_idx_name(usr_id=user_id, scope_id=scope_id, mem_type=MemoryType.SUMMARY.value)
        await self.semantic_recall.delete_docs([mem_id], table_name)
        # semantic memory embedding must not encrypt
        await self.semantic_recall.add_docs([(mem_id, new_memory)], table_name)
        return True

    async def delete(self, user_id: str, scope_id: str, mem_id: str, **kwargs):
        """delete memory by its id."""
        data = await self.mem_store.get(user_id=user_id, scope_id=scope_id, mem_id=mem_id)
        if data is None:
            memory_logger.error(
                "Delete summary in store failed, the mem of mem_id is not exist.",
                event_type=LogEventType.MEMORY_STORE,
                memory_id=mem_id,
                user_id=user_id,
                scope_id=scope_id
            )
            return False
        await self.mem_store.delete(mem_id=mem_id, user_id=user_id, scope_id=scope_id)
        await self._delete_vector_summary_memory(memory_id=[mem_id],
                                                 user_id=user_id,
                                                 scope_id=scope_id)
        return True

    async def delete_by_user_id(self, user_id: str, scope_id: str):
        """delete memory by user id and app id."""
        data = await self.mem_store.get_all(user_id=user_id, scope_id=scope_id, mem_type=MemoryType.SUMMARY.value)
        if data is None:
            memory_logger.error(
                "Delete summary in store failed, the mem of user_id is not exist.",
                event_type=LogEventType.MEMORY_STORE,
                user_id=user_id,
                scope_id=scope_id
            )
            return False
        mem_ids = [item['id'] for item in data]
        await self.mem_store.batch_delete(user_id=user_id, scope_id=scope_id, mem_ids=mem_ids)
        await self._delete_vector_store_table(user_id=user_id,
                                              scope_id=scope_id)
        return True

    async def get(self, user_id: str, scope_id: str, mem_id: str) -> dict[str, Any] | None:
        """get memory by its id."""
        retrieve_res = await self.mem_store.get(user_id=user_id, scope_id=scope_id, mem_id=mem_id)
        retrieve_res["mem"] = BaseMemoryManager.decrypt_memory_if_needed(key=self.crypto_key,
                                                                         ciphertext=retrieve_res["mem"])
        return retrieve_res

    async def search(self, user_id: str, scope_id: str, query: str, top_k: int, **kwargs):
        """query memory, return top k results"""
        mem_ids, scores = await self._recall_by_vector(query, user_id, scope_id, top_k)
        retrieve_res = await self.mem_store.batch_get(user_id=user_id, scope_id=scope_id, mem_ids=mem_ids)
        if retrieve_res is None:
            return None
        for item in retrieve_res:
            item["score"] = scores.get(item['id'], 0)
            item["mem"] = BaseMemoryManager.decrypt_memory_if_needed(key=self.crypto_key, ciphertext=item["mem"])

        retrieve_res.sort(key=lambda x: x["score"], reverse=True)
        return retrieve_res

    async def _add_summary_memory_to_mem_store(self, summary_unit: SummaryUnit):
        """Encrypt and write summary memory to persistent storage."""
        mem = BaseMemoryManager.encrypt_memory_if_needed(key=self.crypto_key,
                                                         plaintext=summary_unit.summary)
        data = {
            'id': summary_unit.mem_id,
            'user_id': summary_unit.user_id or '',
            'scope_id': summary_unit.scope_id or '',
            'mem': mem,
            'source_id': summary_unit.message_mem_id,
            'mem_type': MemoryType.SUMMARY.value,
            'timestamp': summary_unit.timestamp,
            'context_summary': ''
        }
        await self.mem_store.write(user_id=summary_unit.user_id,
                                   scope_id=summary_unit.scope_id,
                                   mem_id=summary_unit.mem_id,
                                   data=data)

    async def _add_summary_memory_to_vector(self, summary_unit: SummaryUnit, scope_id: str) -> bool:
        """Add plaintext summary to vector store for semantic recall."""
        if not self.semantic_recall:
            raise build_error(
                StatusCode.MEMORY_ADD_MEMORY_EXECUTION_ERROR,
                memory_type="summary",
                error_msg="vector store must not be None",
            )
        table_name = generate_idx_name(usr_id=summary_unit.user_id,
                                       scope_id=summary_unit.scope_id,
                                       mem_type=MemoryType.SUMMARY.value)
        return await self.semantic_recall.add_docs(
            docs=[(summary_unit.mem_id, summary_unit.summary)],
            table_name=table_name,
            scope_id=scope_id
        )

    async def _delete_vector_summary_memory(
            self,
            user_id: str,
            scope_id: str,
            memory_id: List[str]
    ):
        """Delete summary memory from vector store by IDs."""
        if not self.semantic_recall:
            raise build_error(
                StatusCode.MEMORY_DELETE_MEMORY_EXECUTION_ERROR,
                memory_type="summary",
                error_msg="vector store must not be None",
            )
        table_name = generate_idx_name(usr_id=user_id, scope_id=scope_id, mem_type=MemoryType.SUMMARY.value)
        await self.semantic_recall.delete_docs(memory_id, table_name)

    async def _recall_by_vector(
            self,
            query: str,
            user_id: str,
            scope_id: str,
            top_k: int = 5
    ) -> tuple[List[str], dict[str, float]]:
        """Semantic recall summary memory IDs and similarity scores."""
        table_name = generate_idx_name(usr_id=user_id, scope_id=scope_id, mem_type=MemoryType.SUMMARY.value)
        memory_hit_info = await self.semantic_recall.search(query=query, table_name=table_name,
                                                            scope_id=scope_id, top_k=top_k)
        return parse_memory_hit_infos(memory_hit_info)

    async def _delete_vector_store_table(self, user_id: str, scope_id: str):
        """Delete entire vector table for user + scope summary memory."""
        if not self.semantic_recall:
            raise build_error(
                StatusCode.MEMORY_DELETE_MEMORY_EXECUTION_ERROR,
                memory_type="summary",
                error_msg="vector store must not be None",
            )
        table_name = generate_idx_name(usr_id=user_id, scope_id=scope_id, mem_type=MemoryType.SUMMARY.value)
        await self.semantic_recall.delete_table(table_name=table_name)
