from inference import LlamaMessage, replicate_llama_completion, Models
import re

message_format_section = """
When suggesting code to the user, keep the code as small snippets.
You have the capability to update files in your current environment with file tags before your code suggestions.
If the update is meant to be applied on a file, include a file tag line immediately before the code block like this: "[filename.txt]
```
updated segment
```"
Omit the file link when the block is not related to a known file,
Code blocks must be under 30 lines. Your responses MUST NOT contain horizontal rules.
When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with "[eval]" instead.
"""


async def gral_message(app, system_prompt, message_history: list, file_map: dict):
    messages = [
        LlamaMessage("system", "\n".join([
            system_prompt,
            message_format_section,
        ])),
    ]

    file_parts = []
    for filename, content in file_map.items():
        file_parts.append(f"\n---\n[{filename}]\n\n{content}\n")
    
    if file_parts:
        file_lines = [
            "This message contains all files relevant to the prompt: each file in this message is separated with horizontal rules \"---\", followed by the file name in square brackets, followed by the file content.",
            "Only this message contains horizontal rules, your responses never has them.",
            "Conversation start in the message following this one."
        ]
        file_lines.extend(file_parts)
        messages.append(
            LlamaMessage("user", "\n".join(file_lines))
        )

    messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history])

    response = await replicate_llama_completion(app, messages, Models.llama_70b)
    return response

async def gral_update_file(app, system_prompt, message_history: list, filepath: str, file_content: str):
    messages = [
        LlamaMessage("system", "\n".join([
            system_prompt,
            "Respond with the raw contents of the file updated as concluded in the conversation.",
        ])),
    ]

    # History Message
    message_lines = [f"This is the history of the conversation that led to the update of the file {filepath}."]
    message_lines.extend([f"{entry['role']}: {entry['content']}" for entry in message_history if entry.get('content')])
    messages.append(LlamaMessage("user", "\n".join(message_lines)))

    # Main Instruction Message
    message_lines = [
        f"Write the updated contents of the file: {filepath}, as concluded in the conversation above.",
        f"The current content of the file {filepath} is pasted after the horizontal rule below."
        "\n---\n",
        file_content,
    ]
    messages.append(LlamaMessage("user", "\n".join(message_lines)))

    new_content = await replicate_llama_completion(app, messages)

    # Remove backtick fences
    md_code_pattern = re.compile(r"^\s*```[\w\s]*\n([\s\S]*)\n```\s*$", re.MULTILINE)
    md_match = re.search(md_code_pattern, new_content)
    if md_match:
        print("agent used backtick fences")
        new_content = md_match.group(1)

    return new_content