import logging
from typing import Any, Dict, List, Tuple
import json

import redis
from redis import Redis
from langchain.memory.chat_memory import BaseChatMemory
from langchain.schema.messages import BaseMessage, get_buffer_string, HumanMessage, AIMessage
from langchain.schema.messages import _message_to_dict, messages_from_dict
from langchain.memory.utils import get_prompt_input_key

from app.core.config import *


class UserConversationBufferWindowMemory(BaseChatMemory):
    """Buffer for storing conversation memory."""

    redis_client: Redis
    human_prefix: str = "Human"
    ai_prefix: str = "AI"
    memory_key: str = "history"  #: :meta private:
    k: int = 5

    @classmethod
    def from_redis(
            cls,
            host: str = REDIS_HOST,
            port: int = REDIS_PORT,
            db: int = 3,
            **kwargs
    ):
        redis_client = Redis(host=host, port=port, db=db)
        try:
            response = redis_client.ping()
            if response:
                print("Connect to redis server successfully.")
                logging.info("Connect to redis server successfully.")
            else:
                print("Fail to connect to redis server")
                logging.info("Fail to connect to redis server")
        except redis.RedisError as e:
            logging.info(f"Error occurs when connecting to redis server: {str(e)}")
        return cls(redis_client=redis_client, **kwargs)

    @property
    def memory_variables(self) -> List[str]:
        """Will always return list of memory variables.

        :meta private:
        """
        return [self.memory_key]

    def load_memory_variables(self, inputs: Dict[str, Any], key: str = "") -> Dict[str, str]:
        """Return history buffer."""
        _items: Any = self.redis_client.lrange(key, 0, self.k * 2) if self.k > 0 else []
        items = [json.loads(m.decode("utf-8")) for m in _items[::-1]]
        buffer = messages_from_dict(items)
        if not self.return_messages:
            buffer = get_buffer_string(
                buffer,
                human_prefix=self.human_prefix,
                ai_prefix=self.ai_prefix,
            )
        return {self.memory_key: buffer}

    def _get_input_output(
            self, inputs: Dict[str, Any], outputs: Dict[str, str]
    ) -> Tuple[str, str]:
        if self.input_key is None:
            prompt_input_key = get_prompt_input_key(inputs, self.memory_variables)
        else:
            prompt_input_key = self.input_key
        if self.output_key is None:
            if len(outputs) != 1:
                raise ValueError(f"One output key expected, got {outputs.keys()}")
            output_key = list(outputs.keys())[0]
        else:
            output_key = self.output_key
        return inputs[prompt_input_key], outputs[output_key]

    def add_message(self, key: str, message: BaseMessage) -> None:
        self.redis_client.lpush(key, json.dumps(_message_to_dict(message)))

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str], key: str = "") -> None:
        """Save context from this conversation to buffer."""
        input_str, output_str = self._get_input_output(inputs, outputs)
        self.add_message(key, HumanMessage(content=input_str))
        self.add_message(key, AIMessage(content=output_str))

    # def clear(self, key) -> None:
    #     """Clear memory contents."""
    #     self.redis_client.delete(key)
