"""Time extraction operation for personal memories.

This module provides functionality to extract time-related information from queries
using LLM-based extraction and pattern matching.
"""

import re
from typing import Dict

from flowllm.core.context import C
from flowllm.core.enumeration import Role
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import Message
from loguru import logger

from reme_ai.constants.common_constants import QUERY_WITH_TS, EXTRACT_TIME_DICT
from reme_ai.constants.language_constants import DATATIME_KEY_MAP
from reme_ai.utils.datetime_handler import DatetimeHandler


@C.register_op()
class ExtractTimeOp(BaseAsyncOp):
    """
    A specialized worker class designed to identify and extract time-related information
    from text generated by an LLM, translating date-time keywords based on the set language,
    and storing this extracted data within a shared context.
    """

    file_path: str = __file__
    EXTRACT_TIME_PATTERN = r"-\s*(\S+)[：:]\s*(\S+)"

    def get_language_value(self, value_dict: dict):
        """
        Get value from dictionary based on current language setting.

        Args:
            value_dict: Dictionary with language keys

        Returns:
            Value for current language or English fallback
        """
        return value_dict.get(self.language, value_dict.get("en"))

    async def async_execute(self):
        """
        Executes the primary logic of identifying and extracting time data from an LLM's response.

        This method first checks if the input query contains any datetime keywords. If not, it logs and returns.
        It then constructs a prompt with contextual information including formatted timestamps and calls the LLM.
        The response is parsed for time-related data using regex, translated via a language-specific key map,
        and the resulting time data is stored in the shared context.
        """
        query, query_timestamp = self.context[QUERY_WITH_TS]

        # Identify if the query contains datetime keywords
        contain_datetime = DatetimeHandler.has_time_word(query, self.language)
        if not contain_datetime:
            logger.info(f"Query contains no datetime keywords: {contain_datetime}")
            # Set empty time dict for downstream operations
            self.context[EXTRACT_TIME_DICT] = {}
            return

        # Prepare the prompt with necessary contextual details
        time_format = self.prompt_format(prompt_name="time_string_format")
        query_time_str = DatetimeHandler(dt=query_timestamp).string_format(time_format, self.language)

        # Create message with system and few-shot examples
        system_prompt = self.prompt_format(prompt_name="extract_time_system")
        few_shot = self.prompt_format(prompt_name="extract_time_few_shot")
        user_prompt = self.prompt_format(
            prompt_name="extract_time_user_query",
            query=query,
            query_time_str=query_time_str,
        )

        full_prompt = f"{system_prompt}\n\n{few_shot}\n\n{user_prompt}"
        logger.info(f"Extracting time from query: {query[:100]}...")

        # Invoke the LLM to generate a response
        response = await self.llm.achat([Message(role=Role.USER, content=full_prompt)])

        # Handle empty or unsuccessful responses
        if not response or not response.content:
            logger.warning("LLM returned empty response for time extraction")
            self.context[EXTRACT_TIME_DICT] = {}
            return

        response_text = response.content

        # Extract and parse time information from the LLM's response
        extract_time_dict = self._parse_time_from_response(response_text)

        logger.info(f"Extracted time information: {extract_time_dict}")
        self.context[EXTRACT_TIME_DICT] = extract_time_dict

    def _parse_time_from_response(self, response_text: str) -> Dict[str, str]:
        """
        Parse time information from LLM response using regex.

        Args:
            response_text: Raw LLM response content

        Returns:
            Dictionary of extracted time information
        """
        extract_time_dict: Dict[str, str] = {}
        matches = re.findall(self.EXTRACT_TIME_PATTERN, response_text)
        key_map: dict = DATATIME_KEY_MAP[DatetimeHandler.language_transform]

        for key, value in matches:
            if key in key_map.keys():
                extract_time_dict[key_map[key]] = value

        logger.debug(f"Time extraction - Response: {response_text[:200]}... Matches: {matches}")
        return extract_time_dict
