from inference import LlamaMessage, replicate_llama_completion
import re

message_format_section = """
When suggesting code to the user, keep the code as small snippets.
If the update is meant to be applied to a file, include a line linking the file 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. Assistant 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 section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content.",
            "Only this message contains horizontal rules, assistant responses must not include them.",
            "The prompt and conversation start in the next message after this."
        ]
        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)
    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 raw file contents of the file {filepath} are included following 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