# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import re
import threading
from typing import TYPE_CHECKING

from typing_extensions import override

from . import _utils
from .base_memory_service import BaseMemoryService
from .base_memory_service import SearchMemoryResponse
from .memory_entry import MemoryEntry

if TYPE_CHECKING:
  from ..events.event import Event
  from ..sessions.session import Session


def _user_key(app_name: str, user_id: str):
  return f'{app_name}/{user_id}'


def _extract_words_lower(text: str) -> set[str]:
  """Extracts words from a string and converts them to lowercase."""
  return set([word.lower() for word in re.findall(r'[A-Za-z]+', text)])


class InMemoryMemoryService(BaseMemoryService):
  """An in-memory memory service for prototyping purpose only.

  Uses keyword matching instead of semantic search.

  This class is thread-safe, however, it should be used for testing and
  development only.
  """

  def __init__(self):
    self._lock = threading.Lock()

    self._session_events: dict[str, dict[str, list[Event]]] = {}
    """Keys are "{app_name}/{user_id}". Values are dicts of session_id to
    session event lists.
    """

  @override
  async def add_session_to_memory(self, session: Session):
    user_key = _user_key(session.app_name, session.user_id)

    with self._lock:
      self._session_events[user_key] = self._session_events.get(user_key, {})
      self._session_events[user_key][session.id] = [
          event
          for event in session.events
          if event.content and event.content.parts
      ]

  @override
  async def search_memory(
      self, *, app_name: str, user_id: str, query: str
  ) -> SearchMemoryResponse:
    user_key = _user_key(app_name, user_id)

    with self._lock:
      session_event_lists = self._session_events.get(user_key, {})

    words_in_query = _extract_words_lower(query)
    response = SearchMemoryResponse()

    for session_events in session_event_lists.values():
      for event in session_events:
        if not event.content or not event.content.parts:
          continue
        words_in_event = _extract_words_lower(
            ' '.join([part.text for part in event.content.parts if part.text])
        )
        if not words_in_event:
          continue

        if any(query_word in words_in_event for query_word in words_in_query):
          response.memories.append(
              MemoryEntry(
                  content=event.content,
                  author=event.author,
                  timestamp=_utils.format_timestamp(event.timestamp),
              )
          )

    return response
