
import uuid
import json
import inspect
from typing import List, Union, Optional, Any, Dict, Iterator
from collections import defaultdict
from pydantic import BaseModel, model_validator, field_validator, Field, UUID4
from datetime import datetime

from gwenflow.llms import ChatOpenAI
from gwenflow.types import ChatCompletionMessage, ChatCompletionMessageToolCall, ChatCompletionChunk
from gwenflow.tools import BaseTool
from gwenflow.memory import ChatMemoryBuffer
from gwenflow.agents.types import AgentResponse
from gwenflow.agents.prompts import PROMPT_TOOLS, PROMPT_STEPS, PROMPT_TOOLS_REACT_GUIDELINES, PROMPT_TASK
from gwenflow.utils import logger
from gwenflow.utils.chunks import merge_chunk


MAX_TURNS = float('inf') #10


class Agent(BaseModel):

    # --- Agent Settings
    id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True)
    name: str = Field(description="Name of the agent")
    role: str = Field(description="Role of the agent")
    description: Optional[str] = Field(default=None, description="Description of the agent")

    # --- Settings for system message
    system_prompt: Optional[str] = None
    instructions: Optional[Union[str, List[str]]] = []
    add_datetime_to_instructions: bool = True
    markdown: bool = False
    response_model: Optional[dict] = None
    steps: Optional[List[str]] = []
    follow_steps: bool = False
    is_react: bool = False
    system_prompt_allowed: bool = True
 
    # --- Agent Model and Tools
    llm: Optional[Any] = Field(None, validate_default=True)
    tools: List[BaseTool] = []
    tool_choice: Optional[Union[str, Dict[str, Any]]] = None
    show_tool_calls: bool = False

    # --- Reasoning models
    reasoning_model: Optional[Any] = Field(None, validate_default=True)

    # --- Task, Context and Memory
    context_vars: Optional[List[str]] = []
    memory: Optional[ChatMemoryBuffer] = None
    metadata: Optional[Dict[str, Any]] = None
    # knowledge: Optional[Knowledge] = None

    # --- Team of agents
    team: Optional[List["Agent"]] = None


    @field_validator("id", mode="before")
    @classmethod
    def deny_user_set_id(cls, v: Optional[UUID4]) -> None:
        if v:
            raise ValueError("This field is not to be set by the user.")

    @field_validator("instructions", mode="before")
    @classmethod
    def set_instructions(cls, v: Optional[Union[List, str]]) -> str:
        if isinstance(v, str):
            instructions = [v]
            return instructions
        return v

    @field_validator("llm", mode="before")
    @classmethod
    def set_llm(cls, v: Optional[Any]) -> Any:
        llm = v or ChatOpenAI(model="gpt-4o-mini")
        return llm

    @model_validator(mode="after")
    def model_valid(self) -> Any:
        if self.memory is None and self.llm is not None:
             token_limit = self.llm.get_context_window_size()
             self.memory = ChatMemoryBuffer(token_limit=token_limit)
        return self
    
    def get_system_message(self, context: Optional[Any] = None):
        """Return the system message for the Agent."""

        if self.system_prompt is not None:
            return dict(role="system", content=self.system_prompt)

        system_message_lines = []

        # name, role and description
        txt = f"You are an AI agent named '{self.name}'."
        if self.role:
            txt += f" Your role is: {self.role.strip('.')}."
        if self.description:
            txt += f" {self.description.strip('.')}."
        system_message_lines.append(f"{txt}\n")

        if self.add_datetime_to_instructions:
            txt = f"The current date and time is:\n<current_date>\n{ datetime.now() }\n</current_date>\n"
            system_message_lines.append(txt)

        # tools
        if self.tools:
            if self.is_react:
                tools = self.get_tools_text_schema()
                system_message_lines.append(PROMPT_TOOLS.format(tools=tools).strip())
                system_message_lines.append(PROMPT_TOOLS_REACT_GUIDELINES)
                system_message_lines.append("")
            else:
                tools = ",".join(self.get_tool_names())
                system_message_lines.append(PROMPT_TOOLS.format(tools=tools).strip())
                system_message_lines.append("")

        # instructions
        instructions = []        
        if self.response_model:
            instructions.append("Use JSON to format your answers.")
        elif self.markdown:
            instructions.append("Use markdown to format your answers.")
        if self.tools is not None:
            instructions.append("Only use the tools you are provided.")
        if context is not None:
            instructions.append("Always prefer information from the provided context over your own knowledge.")

        if len(self.instructions) > 0:
            instructions += self.instructions

        if len(instructions) > 0:
            system_message_lines.append("## Guidelines:")
            system_message_lines.extend([f"- {instruction}" for instruction in instructions])
            system_message_lines.append("")

        if self.response_model:
            system_message_lines.append("## Provide your output using the following JSON schema:")
            system_message_lines.append("<json_schema>")
            system_message_lines.append(json.dumps(self.response_model, indent=4))
            system_message_lines.append("</json_schema>")
            system_message_lines.append("")

        if self.follow_steps:
            if len(self.steps) > 0:
                system_message_lines.append("## Follow these steps:\n")
                system_message_lines.extend([f"{i+1}. {step}" for i, step in enumerate(self.steps)])
                system_message_lines.append("")
            else:
                system_message_lines.append(PROMPT_STEPS.strip())
                system_message_lines.append("")

        # final system prompt
        if len(system_message_lines) > 0:
            return dict(role="system", content=("\n".join(system_message_lines)).strip())
        
        return None


    def format_context(self, context):
        prompt  = "Use the following information if it helps:\n\n"

        if isinstance(context, str):
            prompt += "<context>\n"
            prompt += context + "\n"
            prompt += "</context>\n\n"

        elif isinstance(context, dict):
            for key in context.keys():
                prompt += f"<{key}>\n"
                prompt += context.get(key) + "\n"
                prompt += f"</{key}>\n\n"

        return prompt

    def get_user_message(self, task: Optional[str] = None, context: Optional[Any] = None):
        """Return the user message for the Agent."""

        if not task and not context:
            raise ValueError("You need a task or a context (or both) to run the agent!")

        prompt = ""

        if context:
            prompt += self.format_context(context)

        if task:
            prompt += PROMPT_TASK.format(task=task)

        return { "role": "user", "content": prompt }

    
    def get_tools_openai_schema(self):
        return [tool.openai_schema for tool in self.tools]

    def get_tools_text_schema(self) -> str:
        descriptions = []
        for tool in self.tools:
            sig = inspect.signature(tool._run)
            description = f"{tool.name}{sig} - {tool.description}"
            descriptions.append(description)
        return "\n".join(descriptions)

    def get_tools_map(self):
        return {tool.name: tool for tool in self.tools}

    def get_tool_names(self):
        return [tool.name for tool in self.tools]

    def execute_tool_call(self, tool_name: str, arguments: Dict[str, str]) -> Any:
        available_tools = self.get_tools_map()
        if tool_name not in available_tools:
            logger.error(f"Unknown tool {tool_name}, should be instead one of { available_tools.keys() }.")
            return None

        logger.debug(f"Tool call: {tool_name}({arguments})")
        observation = available_tools[tool_name].run(**arguments)

        return observation

    def handle_tool_calls(
        self,
        tool_calls: List[ChatCompletionMessageToolCall],
    ) -> List:
        
        tool_map = self.get_tools_map()

        messages = []
        for tool_call in tool_calls:

            tool_name = tool_call.function.name

            # handle missing tool case, skip to next tool
            if tool_name not in tool_map.keys():
                logger.warning(f"Unknown tool {tool_name}, should be instead one of { tool_map.keys() }.")
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "tool_name": tool_name,
                        "content": f"Tool {tool_name} does not exist",
                    }
                )
                continue

            arguments = json.loads(tool_call.function.arguments)
            observation = self.execute_tool_call(tool_name, arguments)
            
            if observation:
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "tool_name": tool_name,
                        "content": f"Observation: {observation}",
                    }
                )

        return messages

    def invoke(self, messages: list, stream: bool = False) ->  Union[Any, Iterator[Any]]:

        tools = self.get_tools_openai_schema()

        params = {
            "messages": messages,
            "tools": tools or None,
            "tool_choice": self.tool_choice,
        }

        response_format = None
        if self.response_model:
            response_format = {"type": "json_object"}

        if stream:
            return self.llm.stream(**params, response_format=response_format)
        
        return self.llm.invoke(**params, response_format=response_format)

    def _run(
        self,
        task: Optional[str] = None,
        *,
        context: Optional[Any] = None,
        stream: Optional[bool] = False,
    ) ->  Iterator[AgentResponse]:

        # add reasoning
        if self.reasoning_model:
            reasoning_message = self.reason(task=task)
            if reasoning_message:
                if not context:
                    context = {}
                if isinstance(context, str):
                    text = context
                    context = { "context": text }
                context["thinking"] = reasoning_message

        messages_for_model = []

        # system messages
        system_message = self.get_system_message(context=context)
        if system_message:
            messages_for_model.append(system_message)

        # user messages
        user_message = self.get_user_message(task=task, context=context)
        if user_message:
            messages_for_model.append(user_message)
            self.memory.add_message(user_message)

        # global loop
        init_len = len(messages_for_model)
        while len(messages_for_model) - init_len < MAX_TURNS:

            if stream:
                message = {
                    "content": "",
                    "role": "assistant",
                    "function_call": None,
                    "tool_calls": defaultdict(
                        lambda: {
                            "function": {"arguments": "", "name": ""},
                            "id": "",
                            "type": "",
                        }
                    ),
                }

                completion = self.invoke(messages=messages_for_model, stream=True)

                for chunk in completion:
                    
                    if len(chunk.choices) > 0:
                        delta = json.loads(chunk.choices[0].delta.json())
                        if delta["content"]:
                            yield AgentResponse(
                                delta=delta["content"],
                                messages=None,
                                agent=self,
                                tools=self.tools,
                            )
                        elif delta["tool_calls"] and self.show_tool_calls:
                            if delta["tool_calls"][0]["function"]["name"] and not delta["tool_calls"][0]["function"]["arguments"]:
                                response = f"""**Calling:** {delta["tool_calls"][0]["function"]["name"]}"""
                                yield AgentResponse(
                                    delta=response,
                                    messages=None,
                                    agent=self,
                                    tools=self.tools,
                                )
                        delta.pop("role", None)
                        merge_chunk(message, delta)

                message["tool_calls"] = list(message.get("tool_calls", {}).values())
                message = ChatCompletionMessage(**message)
            
            else:
                completion = self.invoke(messages=messages_for_model)
                message = completion.choices[0].message

            # add messages to the current message stack
            message_dict = json.loads(message.model_dump_json())
            messages_for_model.append(message_dict)

            if not message.tool_calls:
                logger.debug("Task done.")
                self.memory.add_message(message_dict) # We only keep the answer in memory (not tool calls)
                break

            # handle tool calls and switching agents
            tool_response_messages = self.handle_tool_calls(message.tool_calls)
            messages_for_model.extend(tool_response_messages)

        content = messages_for_model[-1]["content"]
        if self.response_model:
            content = json.loads(content)

        yield AgentResponse(
            content=content,
            messages=messages_for_model[init_len:],
            agent=self,
            tools=self.tools,
        )


    def run(
        self,
        task: Optional[str] = None,
        *,
        context: Optional[Any] = None,
        stream: Optional[bool] = False,
        output_file: Optional[str] = None,
    ) ->  Union[AgentResponse, Iterator[AgentResponse]]:

        logger.debug("")
        logger.debug("------------------------------------------")
        logger.debug(f"Running Agent: { self.name }")
        logger.debug("------------------------------------------")
        logger.debug("")

        if stream:
            response = self._run(
                task=task,
                context=context,
                stream=True,
            )
            return response
    
        else:

            response = self._run(
                task=task,
                context=context,
                stream=False,
            )
            response = next(response)

            if output_file:
                with open(output_file, "a") as file:
                    file.write("\n")
                    file.write("---\n\n")
                    file.write(f"# Agent: { self.name }\n")
                    file.write(f"{ task }\n")
                    file.write("\n")
                    file.write(response.content)

            return response
