import sys
import json
import random
from YOUR_API_CALLER_GOES_HERE import LLM_COMPLETION

def get_response(title, poet):
    """
    Creates poem-generation prompt from title and poet, and puts it into
    the expected penchant format
    """

    penchant_prompt = f'Write an artistic poem. The title should be "{title}". Use vivid imagery and creative metaphors throughout your work. Draw inspiration from the works of {poet}. Pay close attention to the sounds that words make when read aloud, and use these sounds to create rhythm and musicality within your piece.'

    penchant_template = f"<|im_start|>system\nYou are Dolphin, a helpful AI assistant.<|im_end|>\n<|im_start|>user\n{penchant_prompt}<|im_end|>\n<|im_start|>assistant\n"

    poem = LLM_COMPLETION(penchant_template)
    return poem


def process_files(titles_file, poets_file, output_file):
    try:
        # Read "titles" and "poets" from input files
        with open(titles_file, 'r') as f:
            titles_data = json.load(f)
            titles = titles_data.get("titles", [])

        with open(poets_file, 'r') as f:
            poets_data = json.load(f)
            poets = poets_data.get("poets", [])

        # Write responses to output JSONL file
        with open(output_file, 'w') as f:
            poems_to_write = 100
            count = 0
            while count < poems_to_write:
                title = random.choice(titles)
                poet = random.choice(poets)
                title = title.strip()
                poet = poet.strip()
                response = get_response(title, poet)
                output_data = {"poet": poet, "title": title, "poem": response}
                f.write(json.dumps(output_data) + "\n")
                count += 1
    except FileNotFoundError as e:
        print(f"Error: {e.filename} not found.")
    except Exception as e:
        print(f"An error occurred: {e}")

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python script.py titles_file.json poets_file.json output_file.jsonl")
    else:
        titles_file = sys.argv[1]
        poets_file = sys.argv[2]
        output_file = sys.argv[3]
        process_files(titles_file, poets_file, output_file)