#!/usr/bin/env python3
import ollama

# read file to process
with open('sample.py', 'r') as file:
    file_contents = file.read()

# prepare role description and prompt 
role = """You are a programming assistant. You analyze code 
  and return an improved version of it. Organize code into 
  functions where possible. Include code comments where 
  necessary.
  You return code ONLY! 
  DO NOT ADD EXPLANATIONS. 
  DO NOT USE MARKDOWN to format your response."""
prompt = f"Optimize the following code:\n\n{file_contents}"

response = ollama.chat(
    model="gpt-oss:20b", 
    messages = [ {"role": "system", "content": role},
                 {"role": "user",   "content": prompt} ] )


# print(response)

if response.done_reason != "stop":
    print("Code optimization failed.")
else:
    new_code = response.message.content.strip()
    with open('sample-improved.py', 'w') as file:
        file.write(new_code)
    print("\nImproved code\n\n")
    print(new_code)
