"""
Abstract base class for AI agents using LangChain.
This module provides a foundation for creating and managing AI agents
that interact with language models and retrieve information.
"""
import logging
from time import sleep
from abc import ABC, abstractmethod
from typing import List, Tuple, Callable, Iterable
import threading
from textwrap import dedent

import numpy as np
from langchain_core.language_models.llms import BaseLLM
from langchain_core.vectorstores.base import VectorStore


class BaseAgent(ABC):
    """
    Abstract base class for an AI agent that interacts with a language model
    and retrieves information to respond to queries.
    """

    base_llm: BaseLLM
    """
    The base language model used by the agent. This model is responsible for
    generating responses and performing language processing tasks.
    """

    vector_store: VectorStore
    """
    The vector store used by the agent for retrieving information.
    """

    q_and_a: List[Tuple[str, str | Callable]]
    """
    A list of tuples, where each tuple contains:
        - A question (str).
        - The expected response, which can either be:
            - A string (str), or
            - A callable function (Callable) that returns the expected
              response as a string.
    Used for evaluating the agent's performance.
    """

    _evaluation_queries_str: str
    _evaluation_queries: List[str]

    def __init__(self):
        self._is_running = False
        self.nap_time = 30
        self.received_response = None
        self._thread = None
        self.logger = logging.getLogger(self.__class__.__name__)
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        )
        handler = logging.StreamHandler()
        format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
        formatter = logging.Formatter(format)
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.setLevel(logging.WARNING)

    @property
    @abstractmethod
    def version(self) -> str:
        """
        Abstract property to return the version of the agent.

        Returns:
            str: The version of the agent.
        """
        return ""

    @abstractmethod
    def listen(self, context: str) -> bool:
        """
        Other agents can call this function to pass contextual information.

        Args:
            context (str): The contextual information in natural language.

        Returns:
            bool: Return True if the knowledge has been accepted
        """
        return True

    @abstractmethod
    def _retrieve(self, q: str) -> List[Tuple[str, dict]]:
        """
        Retrieve relevant information for a given query.

        Args:
            q (str): The query string for which relevant data needs to be
                     retrieved.

        Returns:
            List[Tuple[str, dict]]: A list of tuples where each tuple contains
            a relevant piece of information (e.g., a document or text snippet)
            as a string, and its associated metadata as a dictionary.
        """
        return []

    @abstractmethod
    def _prompt(self, q: str) -> str:
        """
        Generate a formatted prompt based on the given query.

        Args:
            q (str): The query string for which the prompt is generated.

        Returns:
            str: The formatted prompt that will be used for generating a
            response.
        """
        return ""

    @abstractmethod
    def respond(self, q: str) -> Iterable[str]:
        """
        Generate a response for the given query by retrieving information and
        interacting with the language model.

        Args:
            q (str): The query string for which the response is generated.

        Yields:
            Iterable[str]: An iterable of response strings generated by the
            agent.
        """
        while False:
            yield ""

    def ask(self, agent: "BaseAgent", q: str) -> Iterable[str]:
        """
        Asynchronously send a query to another agent and process its response.

        Args:
            agent (Agent): Another agent to which the query is sent. q (str):
            The query string to be sent.

        Yields:
            Iterable[str]: An iterable of response strings generated by the
            other agent.
        """
        self.logger.debug("Asking %s query: %s", agent, q)
        self.received_response = ""
        for chunk in agent.respond(q):
            self.received_response += chunk
            yield chunk
        self.logger.debug("Received %s from %s in response to query: %s",
                          self.received_response,
                          agent,
                          q)

    def receive(self, generator: Iterable) -> str:
        """
        Concatenates the output of an iterable generator into a single string.

        Args:
            generator (Iterable): An iterable object (e.g., a generator) that
            yields chunks of text or other string-like elements.

        Returns:
            str: The concatenated string formed by joining all elements from the
            generator.
        """
        return ''.join(list(generator))

    def tell(self, agent: "BaseAgent", context: str) -> bool:
        """
        Send information to another agent.

        Args:
            agent (Agent): The recipient agent to which the information is sent.
            context (str): The information or context to be shared, framed as a
            query.

        """
        self.logger.debug("Telling %s context: %s", agent, context)
        return agent.listen(context)

    @abstractmethod
    def run(self) -> None:
        """
        Run the agent. This method should contain the logic that needs to be
        executed in the loop.
        """

    def _run(self) -> None:
        """
        Internal method to run the agent in a loop.
        """
        while self._is_running:
            self.run()
            if self.nap_time is not None:
                sleep(self.nap_time)
            else:
                sleep(30)

    def start(self) -> None:
        """
        Start the agent by creating and starting the thread.
        """
        if not self._is_running:
            self.logger.info("Starting agent %s", self.__class__.__name__)
            self._is_running = True
            self._thread = threading.Thread(target=self._run, daemon=True)
            self._thread.start()

    def stop(self) -> None:
        """
        Stop the agent and wait for the thread to finish.
        """
        if self._is_running:
            self.logger.info("Stopping agent %s", self.__class__.__name__)
            self._is_running = False
            if self._thread is not None:
                self._thread.join()

    def evaluate(self) -> Tuple[float, np.ndarray]:
        """
        Evaluate the agent's performance by comparing its responses to a list of
        predefined question-answer pairs (`q_and_a`).

        This method iterates through the list of questions and their
        corresponding expected answers. For each question, the agent generates a
        response using the `respond` method. The response is compared to the
        expected answer (either a string or the output of a callable function)
        to determine correctness.

        Returns:
            Tuple[float, np.ndarray]: A tuple containing:
                - The accuracy as a float (correct answers / total questions),
                  formatted to three decimal places.
                - A NumPy array (`np.ndarray`) where each element is 1 if the
                  response matched the expected answer, or 0 otherwise.

        Example:
            If `q_and_a` contains 10 questions and the agent answers 7
            correctly, and their correctness is stored in a NumPy array, this
            method will return `(0.700, array([1, 1, 0, ...]))`.
        """
        matches = np.zeros(len(self.q_and_a))
        i = 0
        for q, expected in self.q_and_a:
            if isinstance(expected, str):
                if ''.join(list(self.respond(q))).strip() == expected.strip():
                    matches[i] = 1
            else:
                if ''.join(list(self.respond(q))).strip() == expected.strip():
                    matches[i] = 1
            i += 1

        return f'{matches.sum() / len(matches):.03f}', matches

    def _invoke(self, prompt: str) -> str:
        """
        Synchronously invoke the language model with the given prompt.

        Args:
            prompt (str): The input prompt to be processed by the language model.

        Returns:
            str: The response generated by the language model.
        """
        return self.base_llm.bind(skip_prompt=True).invoke(prompt)

    def _stream(self, prompt: str) -> str:
        """
        Stream responses from the language model for the given prompt.

        Args:
            prompt (str): The input prompt to be processed by the language model.

        Yields:
            str: A stream of responses generated by the language model.
        """
        return self.base_llm.bind(skip_prompt=True).stream(prompt)

    def _is_valid_response(response: str) -> bool:
        """
        Validate whether the response is 'yes' or 'no'.

        Args:
            response (str): The response string to validate.

        Returns:
            bool: True if the response is 'yes', False if 'no'.

        Raises:
            RuntimeError: If the response is not 'yes' or 'no'.
        """
        response = response.strip().lower()
        if response not in {'yes', 'no'}:
            raise RuntimeError("Incorrect response in evaluation: "
                               f"Expected: <yes|no> "
                               f"Actual: {response}.")
        return response == 'yes'

    def self_evaluate(self, n: int = 5) -> Tuple[float, np.ndarray]:
        """
        Evaluates the validity of the model's responses to generated
        evaluation queries.

        This method generates `n` evaluation queries based on the model's
        prompt template, evaluates the model's responses to those queries, and
        returns the results.

        Args:
            n (int, optional): The number of evaluation queries to generate.
            Defaults to 5.

        Returns:
            Tuple[float, np.ndarray]:
                - A float representing the fraction of valid responses
                  (validity score).
                - A numpy array containing a boolean value for each query,
                where `True` indicates a valid response and `False` indicates
                an invalid response.

        Raises:
            RuntimeError: If the number of generated evaluation queries does
            not match `n`. AssertionError: If the model's evaluation response
            is not "yes" or "no".

        Process:
            1. Generate `n` evaluation queries by replacing `<QUERY>` in the
               prompt template.
            2. Evaluate the model's responses to these queries.
            3. For each query, the model determines whether the response is
               valid ("yes") or not ("no").
            4. Return the fraction of valid responses and a boolean array
               representing the validity
            of responses for all queries.

        Notes:
            - The prompt template and separator used for evaluation queries
              are defined within the method.
            - Temperature and max tokens for the evaluation LLM are not
              currently configurable (marked as TODO).
            - This method relies on the `self._prompt` and `self.base_llm`
              components for generating
            and evaluating responses.
        """
        matches = np.zeros(n)

        sep = '### EVALUATION QUERY ###'

        prompt_template = self._prompt('<QUERY>')

        evaluation_query_prompt = dedent(f"""Consider the following prompt template.

        ******
        {prompt_template}
        ******

        Write {n} queries that could be used to evaluate the validity of the response.
        Separate them with the separator {sep}.
        These queries will replace <QUERY> in the prompt.

        Make sure to output exactly {n} queries!

        """)

        queries_str = self\
            .base_llm\
            .bind(skip_prompt=True)\
            .invoke(evaluation_query_prompt)
        self._evaluation_queries_str = queries_str
        queries = sep.split(queries_str)
        if len(queries) != n:
            raise RuntimeError(
                "Incorrect count for evaluation queries: "
                f"Expected: {n} "
                f"Actual: {len(queries)}"
            )

        self._evaluation_queries = queries

        for i, query in enumerate(queries):
            prompt = self._prompt(query)
            self_response = ''.join(list(self.respond(query)))
            evaluation_prompt = dedent(f"""Consider the following prompt:

            ******
            {prompt}
            ******

            Consider the following response to this prompt:

            ******
            {self_response}
            ******

            Is the response a valid response to the prompt.
            Respond only with "yes" if it is, otherwise, respond only with "no".

            One-word response:""")

            # TODO: Find a way to decrease the temperature and the max tokens.
            matches[i] = self._is_valid_response(
                self._invoke(evaluation_prompt)
            )

        return matches
