#!/usr/bin/env python3
import os
import ollama

# recursively find all files in the given directory 
# with the specified extensions.
def find_files_with_extensions(directory, extensions):
    matching_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if any(file.endswith(ext) for ext in extensions):
                matching_files.append(os.path.join(root, file))
    return matching_files

# extract code from a markdown text
def extract_code(md):
    md = md.strip()
    lines = md.split("\n")
    first, last = 0, len(lines)
    for i in range(len(lines)):
        if lines[i].startswith("```"):
            first = i + 1
            break
    if first == 0:
        return md
    else:
        print("  Markdown found")
    for i in range(first + 1, len(lines)):
        if lines[i].startswith("```"):
            last = i
            break
    return "\n".join(lines[first:last])

# reads file, adds comments and writes the file back
def comment(filename):
    # read file to process
    with open(filename, "r") as file:
        content = file.read()
    with open(filename + ".bak", "w") as backup:
        backup.write(content)
    # prepare role description and prompt 
    role = """You are a coding assistant.
        Your goal is to add missing comments to existing code.
        At the very least, add a header to the file to indicate its purpose.
        Try to explain the purpose of each function/method.
        If necessary, add comments to explain blocks of code.
        Don't change/remove existing comments.
        Apart from comments, leave the code as it is. 
        The functionality of the code MUST NOT be changed.
        DO NOT ADD EXPLANATIONS in your response. 
        DO NOT USE MARKDOWN to format your response."""
    prompt = f"""Please add comments to the 
                 following code:\n\n{content}"""
    try:
        response = ollama.chat(
            model="llama3", 
            # model="codellama",
            # model="gemma2",
            # model="codestral",
            options = { "temperature": 0.5 },
            messages = [ {"role": "system", "content": role},
                         {"role": "user",   "content": prompt} ] )
    except Exception as e:
        print("  An error occurred:", str(e))
        return
    if response["done_reason"] != "stop":
        print("  Commenting failed.")
        print("  Output tokens:", response["eval_count"])  
    else:
        new_code = extract_code(response["message"]["content"])
        with open(filename, 'w') as file:
            file.write(new_code)
        print("  Commenting done.")
        print("  Output tokens:", response["eval_count"])  

# main code
dir = "/Users/kofler/no-sync/coding-sample-files/tst"
extensions = [".js", ".php"]
files = find_files_with_extensions(dir, extensions)
for filename in files:
    print(filename)
    comment(filename)
    exit
