from stream import Stream
import re
import json
import threading

from tools import Toolset, Action, BasicToolset, SpecialToolset

LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"

# 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:
    # - identifier (for kind in creation)
    # - schema: list[AgentField{name, type: str}] (for configuration)
    identifier = "basic"
    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.",
            ])
        ),
    ]

    # while this constructor accepts two parameters,
    # all subclasses must also accept a fourth 'data' parameter
    def __init__(self, app, agent_id, data):
        self.app = app
        self.agent_id = agent_id

        self.config = {field.name: data.get(field.name, field.default) for field in Agent.schema_fields}

        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 = {}

        print("tool prompt section", self.tool_prompt_section)

        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_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.agent_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_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
    
    def create_from_data(app, agent_id, kind, data):
        return Agent(app, agent_id, data)

class IntroAgent(Agent):
    def __init__(self, app, agent_id, _data):
        Agent.__init__(self, app, agent_id)
        self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."

class LlamaMessage:
    def __init__(self, role, content):
        self.role = role
        self.content = content

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

async def replicate_llama_completion(app, messages):
    # Llama 3 instruct
    full_prompt = "<|begin_of_text|>"
    for message in messages:
        header_name = message.role if message.role != "agent" else "assistant"
        full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
    full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
    
    content_stream = Stream(app)
    thread = threading.Thread(
        target=run_replicate_inference_loop,
        kwargs={
            'replicate_client': app.replicate_client,
            'model': LLAMA3_MODEL,
            'prompt': full_prompt,
            'content_stream': content_stream,
        }
    )
    thread.start()
    response = await content_stream.content_wait()
    return response.lstrip()


def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
    print("START INFERENCE")
    completion = replicate_client.run(
        model,
        input={
            "prompt": prompt,
            "temperature": 0.75,
            "max_new_tokens": 500,
            "system_prompt": "" # We handle system prompt in the formatted prompt
        }
    )
    for event in completion:
        if event:
            content_stream.feed(str(event))
    print("END INFERENCE")
    content_stream.close()