#!/home/retoor/.venv/bin/python3
import os
import sys
import json
import requests
import time

def clean_content(content):
    if content.startswith("```"):
        lines = content.split("\n")[1:-1]
        return "\n".join(lines).strip()
    return content.strip()

def main():
    try:
        api_key = os.environ["OPENAI_API_KEY"]
        user_input = os.environ["AI_PROMPT"]
        doc_input = sys.stdin.read()

        system_prompt = (
            "1. You receive a selection of source code from user.\n"
            "2. Your task is to refactor the code according to the user's prompt.\n"
            "3. You repond exclusively with source code. Do not respond with anything else.\n"
            "4. Given selection will be replaced with your response.\n"
            f"5. {user_input}\n"
        )

        data = {
            "model": "gpt-4.1",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": doc_input}
            ],
            "temperature": 0.7
        }

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

        with open("/tmp/ai_output.txt", "w") as f:
            f.write(json.dumps(data, indent=2))

        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=data)

        output = response.json()["choices"][0]["message"]["content"]
        print(clean_content(output), flush=True, end="")

    except Exception as e:
        print(e)

if __name__ == "__main__":
    main()
