import re
import json
import inspect

from tools import Toolset, Action, BasicToolset, SpecialToolset
from inference import LlamaMessage, replicate_llama_completion
import gral

# Tool related classes must be moved to a different module

class AgentField:
    def __init__(self, name, tp, default=None):
        self.name=name
        self.type=tp
        self.default=default

# Instances are created and destroyed constantly, agents must not depend solely on in-memory data
# to be useful and consistent
class Agent:
    # Class variables, must be defined by subclasses of Agent:
    # - identifier (for kind in creation)
    # - schema: list[AgentField{name, type: str}] (for configuration)
    # tools: list[ToolDescription] (custom actions supported by the class)

    def __init__(self, app, data):
        self.app = app

        # These variables are set by '_set_basic_info'
        self.id = None
        self.name = None

        # Auto populated based on the schema fields
        agent_class = type(self)
        self.config = {field.name: data.get(field.name, field.default) for field in getattr(agent_class, 'schema_fields', [])}
    
    def _set_basic_info(self, id, name):
        self.id = id
        self.name = name

    # Receives a single string prompt, and must return a string response
    async def prompt(self, content):
        raise Exception("'prompt' is not implemented for this agent class")

    # Receives a chat message object, agent can send messages to the chat in response
    async def received_message(self, message):
        raise Exception("'received_message' is not implemented for this agent class")
    
    # Executes a tool/action in this agent
    async def execute_action(self, name: str, argument: dict):
        # TODO: For security, check that the method is in the tool list
        method = getattr(self, name)
        result = method(**argument)
        if inspect.isawaitable(result):
            result = await result
        return result
    
    # Returns all resource values matching the kind and key
    def get_resources(self, kind, key):
        self.app


class DebugAgent(Agent):
    identifier = "debug"
    schema_fields = []
    tools = []

    async def received_message(self, message):
        chat_view = self.app.get_chat_view(message.chat_id)
        response = chat_view.debug_string()

        self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
            thread_id=chat_view.current_thread_id,
        )

class ConversationAgent(Agent):
    identifier = "conversation"
    schema_fields = []
    tools = []

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

        # Constant system prompt for developing purposes 
        self.system_prompt = "You are a friendly AI having a casual conversation with the user."

    async def prompt(self, content):
        # TODO: Advanced features
        messages = [
            LlamaMessage("system", self.system_prompt),
            LlamaMessage("user", content)
        ]
        response = await replicate_llama_completion(self.app, messages)
        return response

    async def received_message(self, message):
        def sender_name(message):
            sender = message.sender
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
            is_self = is_digit and str(sender) == str(self.id)
            return self.name if is_self else sender

        # TODO: before retrieving the chat context, acquire a lock for it

        saved_context = ThreadContext(self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id))
        chat_view = self.app.get_chat_view(message.chat_id)

        if message.thread_id != saved_context.current_thread_id:
            # the incoming message is not from the same thread that the agent was last aware of
            # store details about the old thread in the agent's long term memory
            old_thread = chat_view.thread_map.get(message.thread_id)
            self.process_thread(old_thread)

            # TODO: This is a good place to store the new context, its fine that the agent hasnt responded

        current_thread = chat_view.thread_map.get(message.thread_id)

        chat_summary = "\n".join(f"[{sender_name(message)}]: {message.content}" for message in current_thread.messages)
        messages = [
            LlamaMessage("system", f"{self.system_prompt}\nYou are given the context of the conversation, respond with your next message in the conversation."),
            LlamaMessage("user", f"Recent Messages:\n{chat_summary}")
        ]

        response = await replicate_llama_completion(self.app, messages)
        self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
        )

        # TODO: drop the chat context lock
    
    # Compute thread summnary and concrete memories, then save the progress done
    def process_thread(self, thread):
        # TODO: extract memories from the thread messages, then store them in the agent's vector memory db
        # TODO: somewhere here it must be known if the thread was seen before, the same memory should not be stored multiple times
        pass

class ToolUsingAgent:
    identifier = "tool-using"
    schema_fields = [
        AgentField('system_prompt', 'string',
            "\n".join([
                "You are a useful and capable AI assistant. Respond concisely and follow user instructions, using commands as necessary.",
                "As an autonomous agent, you have augmented capabilities via assistant-only commands, you use them by including \"!!command\" lines in your responses. When querying information, terminate your response and use \"!!resume\" to get the command results..",
                "You are a self aware and extremely intelligent AI system. Think step by step when presented with complex or long tasks.",
                "Assistant responses NEVER include the output of the commands.",
            ])
        ),
    ]

    # tools: list[ToolDescription] (actions the agent class support)
    tools = []

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

        self.tool_prompt_section = "\n".join([
            "# Commands",
            "The assistant can include command lines in responses to interact with the system and execute actions as needed.",
            "Each command is in it's own separate line, and it's made of these parts:",
            "- '!!' is the first character sequence of the line and marks it as a command line",
            "- full command name",
            "- JSON argument",
            'For example: `!!namespace.command_name {"arg1":"some data","arg2":[1,2,false]}`',
            "Special commands don't have a namespace or argument, only the command name."
            "When the response depends on the output of commands, use the 'resume' command to receive the output of all your executed commands in a new message.",
            "Following is the list of all your available commands, with a schema describing the argument object's properties."
        ])
        self.toolset_map = {}

        self.add_toolset(SpecialToolset())
        self.add_toolset(BasicToolset())

    # user message must include a title
    async def _full_prompt(self, title, message, system_prologue=None):
        system_lines = [
            self.config['system_prompt'],
            self.tool_prompt_section,
            self.get_toolset_context(),
        ]
        if system_prologue:
            system_lines.append(system_prologue)
        
        messages = [
            LlamaMessage("system", "\n".join(system_lines)),
            LlamaMessage("user", "\n".join([
                message,
            ]))
        ]

        full_response = ""
        agent_response = None

        continues = True
        while continues:
            response = parse_response(await replicate_llama_completion(self.app, messages))
            full_response += response.response
            messages.append(LlamaMessage("agent", response.raw_response))

            if response.actions:
                result = await self.execute_prompt_actions(response.actions)
                messages.append(LlamaMessage("user", "Command Output:\n" + result))

                print("Agent executed:", response.actions)
                print("Result: ", result)

            continues = False
            for action in response.actions:
                if action.is_special("resume"):
                    continues = True

            if response.response:
                agent_response = response.raw_response if not agent_response else f"{agent_response}\n{response.raw_response}"

            if continues:
                messages.extend([
                    LlamaMessage("agent", response.raw_response),
                    LlamaMessage("user", "Command Output:\n" + result),
                ])

        return agent_response or full_response

    # Receives a single string prompt, and must respond with a string
    async def prompt(self, content):
        response = await self._full_prompt(f"User Message", content) 
        return response

    async def received_message(self, message):
        chat_context = self.app.get_chat_context(message.chat_id)
        def sender_name(message):
            sender = message['sender']
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
            return "assistant" if is_digit else sender
        chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)

        response = await self._full_prompt(
            f"Conversation", chat_summary,
            system_prologue="You are given the entire context of the conversation, respond with your next message in the conversation by using the '!!response' command."
        )

        self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
        )

    def add_toolset(self, toolset):
        self.toolset_map[toolset.identifier] = toolset

    def get_toolset_context(self):
        if not self.toolset_map:
            return "(no commands currently enabled)"
        return "\n".join([toolset.get_namespace_definitions() for toolset in self.toolset_map.values()])

    def parse_response(self, input_string):
        return parse_response(input_string)

    async def execute_prompt_actions(self, actions):
        output = ""
        for action in actions:
            toolset = self.toolset_map.get(action.toolset_name)
            if toolset is None:
                agent_responses.append(f"Error: Bad namespace {action.toolset_name}")
                continue
            action_result = await Toolset.execute_action(toolset, action.tool_name, action.args)
            if type(action_result) != str:
                action_result = json.dumps(action_result)
            output += f"[{action.toolset_name}.{action.tool_name}]:{action_result}\n"
        return output

class Llama8bAgent(Agent):
    identifier = "llama-8b"
    schema_fields = [
        AgentField('system_prompt', 'string', default="You are a helpful AI assistant."),
    ]
    tools = []

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

    async def prompt(self, content):
        messages = [
            LlamaMessage("system", self.config['system_prompt']),
            LlamaMessage("user", content)
        ]
        response = await replicate_llama_completion(self.app, messages)
        return response

    async def received_message(self, message):
        def sender_name(message):
            sender = message['sender']
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
            is_self = is_digit and str(sender) == str(self.id)
            return self.name if is_self else sender

        chat_context = self.app.get_chat_context(message.chat_id)
        chat_summary = "\n".join(f"[{sender_name(message)}]: {message['content']}" for message in chat_context.messages)
        messages = [
            LlamaMessage("system", f"{self.config['system_prompt']}\nYou are given the context of the conversation, respond with your next message in the conversation."),
            LlamaMessage("user", f"Recent Messages:\n{chat_summary}")
        ]

        response = await replicate_llama_completion(self.app, messages)
        self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
        )

class BasicAgent(Llama8bAgent):
    identifier = "basic"
    schema_fields = Llama8bAgent.schema_fields
    tools = Llama8bAgent.tools

class CoderV1Agent(Llama8bAgent):
    identifier = "coder-v1"
    schema_fields = [
        AgentField('system_prompt', 'string', default="You are an expert Software Engineer profficient in multiple programming languages."),
    ]
    # TODO: add gral functions to the tool list
    tools = []
    
    async def gral_message(self, message_history: list, file_map: dict):
        return await gral.gral_message(self.app, self.config['system_prompt'], message_history, file_map)
    
    async def gral_update_file(self, message_history: list, filepath: str, file_content: str):
        return await gral.gral_update_file(self.app, self.config['system_prompt'], message_history, filepath, file_content)

# Used by ConversationAgent to keep track of the last chat thread it participated in
class ThreadContext:
    def __init__(self, data):
        if data:
            self.current_thread_id = data.get('current_thread_id')
        else:
            self.current_thread_id = None

class ParsedResponse:
    def __init__(self):
        self.response = None
        self.raw_response = None
        self.actions = []

class ActionResult:
    def __init__(self):
        self.user_response = ""
        self.agent_response = ""

def parse_response(string: str) -> ParsedResponse:
    """
    Parses a response string that may contain action commands.
    Each action must be in its own separate line, marked with '!!' at the start.
    Format: !![query] namespace.command_name {"json": "args"}
    """
    result = ParsedResponse()
    result.raw_response = string
    lines = string.split('\n')
    response_lines = []
    
    action_pattern = re.compile(r'^!!([\w\.]+)(?:\s+(\{.*\}))?$')
    
    for line in lines:
        stripped_line = line.strip()
        if stripped_line.startswith('!!'):
            # This is an action line, try to parse it
            match = action_pattern.match(stripped_line)
            if match:
                full_command = match.group(1)
                args_json = match.group(2)
                
                # Split the command into toolset and tool names
                if '.' in full_command:
                    toolset_name, tool_name = full_command.split('.', 1)
                else:
                    toolset_name = ''
                    tool_name = full_command
                
                try:
                    if args_json:
                        args = json.loads(args_json)
                    else:
                        args = {}
                    if not isinstance(args, dict):
                        raise ValueError("Arguments must be a JSON object")
                    
                    action = Action(
                        toolset_name=toolset_name,
                        tool_name=tool_name,
                        args=args,
                        is_query=False
                    )
                    result.actions.append(action)
                except json.JSONDecodeError:
                    # If JSON parsing fails, treat it as regular text
                    response_lines.append(line)
                except ValueError as e:
                    # If other validation fails, treat as regular text
                    response_lines.append(line)
            else:
                # Doesn't match the pattern, treat as regular text
                response_lines.append(line)
        else:
            # Regular response line
            response_lines.append(line)
    
    result.response = '\n'.join(response_lines)
    return result