# type: ignore
import asyncio
import json
import logging
from abc import ABC, abstractmethod
from datetime import datetime
from json import JSONDecodeError
from typing import Any, AsyncGenerator, Optional, Type

from pydantic import BaseModel

from core.base.abstractions import (
    GenerationConfig,
    LLMChatCompletion,
    Message,
)
from core.base.providers import CompletionProvider, DatabaseProvider

from .base import Tool, ToolResult

logger = logging.getLogger()


class Conversation:
    def __init__(self):
        self.messages: list[Message] = []
        self._lock = asyncio.Lock()

    async def add_message(self, message):
        async with self._lock:
            self.messages.append(message)

    async def get_messages(self) -> list[dict[str, Any]]:
        async with self._lock:
            return [
                {**msg.model_dump(exclude_none=True), "role": str(msg.role)}
                for msg in self.messages
            ]


# TODO - Move agents to provider pattern
class AgentConfig(BaseModel):
    rag_rag_agent_static_prompt: str = "static_rag_agent"
    rag_agent_dynamic_prompt: str = "dynamic_reasoning_rag_agent_prompted"
    stream: bool = False
    include_tools: bool = True
    max_iterations: int = 10

    @classmethod
    def create(cls: Type["AgentConfig"], **kwargs: Any) -> "AgentConfig":
        base_args = cls.model_fields.keys()
        filtered_kwargs = {
            k: v if v != "None" else None
            for k, v in kwargs.items()
            if k in base_args
        }
        return cls(**filtered_kwargs)  # type: ignore


class Agent(ABC):
    def __init__(
        self,
        llm_provider: CompletionProvider,
        database_provider: DatabaseProvider,
        config: AgentConfig,
        rag_generation_config: GenerationConfig,
    ):
        self.llm_provider = llm_provider
        self.database_provider: DatabaseProvider = database_provider
        self.config = config
        self.conversation = Conversation()
        self._completed = False
        self._tools: list[Tool] = []
        self.tool_calls: list[dict] = []
        self.rag_generation_config = rag_generation_config
        self._register_tools()

    @abstractmethod
    def _register_tools(self):
        pass

    async def _setup(
        self, system_instruction: Optional[str] = None, *args, **kwargs
    ):
        await self.conversation.add_message(
            Message(
                role="system",
                content=system_instruction
                or (
                    await self.database_provider.prompts_handler.get_cached_prompt(
                        self.config.rag_rag_agent_static_prompt,
                        inputs={
                            "date": str(datetime.now().strftime("%m/%d/%Y"))
                        },
                    )
                    + f"\n Note,you only have {self.config.max_iterations} iterations or tool calls to reach a conclusion before your operation terminates."
                ),
            )
        )

    @property
    def tools(self) -> list[Tool]:
        return self._tools

    @tools.setter
    def tools(self, tools: list[Tool]):
        self._tools = tools

    @abstractmethod
    async def arun(
        self,
        system_instruction: Optional[str] = None,
        messages: Optional[list[Message]] = None,
        *args,
        **kwargs,
    ) -> list[LLMChatCompletion] | AsyncGenerator[LLMChatCompletion, None]:
        pass

    @abstractmethod
    async def process_llm_response(
        self,
        response: Any,
        *args,
        **kwargs,
    ) -> None | AsyncGenerator[str, None]:
        pass

    async def execute_tool(self, tool_name: str, *args, **kwargs) -> str:
        if tool := next((t for t in self.tools if t.name == tool_name), None):
            return await tool.results_function(*args, **kwargs)
        else:
            return f"Error: Tool {tool_name} not found."

    def get_generation_config(
        self, last_message: dict, stream: bool = False
    ) -> GenerationConfig:
        if (
            last_message["role"] in ["tool", "function"]
            and last_message["content"] != ""
            and "ollama" in self.rag_generation_config.model
            or not self.config.include_tools
        ):
            return GenerationConfig(
                **self.rag_generation_config.model_dump(
                    exclude={"functions", "tools", "stream"}
                ),
                stream=stream,
            )

        return GenerationConfig(
            **self.rag_generation_config.model_dump(
                exclude={"functions", "tools", "stream"}
            ),
            # FIXME: Use tools instead of functions
            # TODO - Investigate why `tools` fails with OpenAI+LiteLLM
            tools=(
                [
                    {
                        "function": {
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool.parameters,
                        },
                        "type": "function",
                        "name": tool.name,
                    }
                    for tool in self.tools
                ]
                if self.tools
                else None
            ),
            stream=stream,
        )

    async def handle_function_or_tool_call(
        self,
        function_name: str,
        function_arguments: str,
        tool_id: Optional[str] = None,
        save_messages: bool = True,
        *args,
        **kwargs,
    ) -> ToolResult:
        logger.debug(
            f"Calling function: {function_name}, args: {function_arguments}, tool_id: {tool_id}"
        )
        if tool := next(
            (t for t in self.tools if t.name == function_name), None
        ):
            try:
                function_args = json.loads(function_arguments)

            except JSONDecodeError as e:
                error_message = f"Calling the requested tool '{function_name}' with arguments {function_arguments} failed with `JSONDecodeError`."
                if save_messages:
                    await self.conversation.add_message(
                        Message(
                            role="tool" if tool_id else "function",
                            content=error_message,
                            name=function_name,
                            tool_call_id=tool_id,
                        )
                    )

                # raise R2RException(
                #     message=f"Error parsing function arguments: {e}, agent likely produced invalid tool inputs.",
                #     status_code=400,
                # )

            merged_kwargs = {**kwargs, **function_args}
            try:
                raw_result = await tool.results_function(
                    *args, **merged_kwargs
                )
                llm_formatted_result = tool.llm_format_function(raw_result)
            except Exception as e:
                raw_result = f"Calling the requested tool '{function_name}' with arguments {function_arguments} failed with an exception: {e}."
                logger.error(raw_result)
                llm_formatted_result = raw_result

            tool_result = ToolResult(
                raw_result=raw_result,
                llm_formatted_result=llm_formatted_result,
            )
            if tool.stream_function:
                tool_result.stream_result = tool.stream_function(raw_result)

            if save_messages:
                await self.conversation.add_message(
                    Message(
                        role="tool" if tool_id else "function",
                        content=str(tool_result.llm_formatted_result),
                        name=function_name,
                        tool_call_id=tool_id,
                    )
                )
                # HACK - to fix issues with claude thinking + tool use [https://github.com/anthropics/anthropic-cookbook/blob/main/extended_thinking/extended_thinking_with_tool_use.ipynb]
                if self.rag_generation_config.extended_thinking:
                    await self.conversation.add_message(
                        Message(
                            role="user",
                            content="Continue...",
                        )
                    )

            self.tool_calls.append(
                {
                    "name": function_name,
                    "args": function_arguments,
                }
            )
        return tool_result


# TODO - Move agents to provider pattern
class RAGAgentConfig(AgentConfig):
    rag_rag_agent_static_prompt: str = "static_rag_agent"
    rag_agent_dynamic_prompt: str = "dynamic_reasoning_rag_agent_prompted"
    stream: bool = False
    include_tools: bool = True
    max_iterations: int = 10
    # tools: list[str] = [] # HACK - unused variable.

    # Default RAG tools
    rag_tools: list[str] = [
        "search_file_descriptions",
        "search_file_knowledge",
        "get_file_content",
    ]

    # Default Research tools
    research_tools: list[str] = [
        "rag",
        "reasoning",
        # DISABLED by default
        "critique",
        "python_executor",
    ]

    @classmethod
    def create(cls: Type["AgentConfig"], **kwargs: Any) -> "AgentConfig":
        base_args = cls.model_fields.keys()
        filtered_kwargs = {
            k: v if v != "None" else None
            for k, v in kwargs.items()
            if k in base_args
        }
        filtered_kwargs["tools"] = kwargs.get("tools", None) or kwargs.get(
            "tool_names", None
        )
        return cls(**filtered_kwargs)  # type: ignore