import os
import json
import requests
import random

# Function to load presets from the "presets" folder
def load_presets():
    preset_files = [f for f in os.listdir("presets") if f.endswith(".json")]
    presets = {}
    for preset_file in preset_files:
        preset_name = os.path.splitext(preset_file)[0]
        with open(os.path.join("presets", preset_file), "r", encoding="utf-8") as preset_data:
            presets[preset_name] = json.load(preset_data)
    return presets

# Function to select a preset or use default
def select_preset(presets):
    while True:
        print("Select a preset:")
        for i, preset in enumerate(presets.keys()):
            print(f"{i + 1}. {preset}")
        choice = input("Enter the number of the preset (default: 'default'): ")
        if choice.isdigit() and 1 <= int(choice) <= len(presets):
            return presets[list(presets.keys())[int(choice) - 1]]
        elif choice.strip() == "":
            return presets["default"]
        else:
            print("Invalid choice. Please enter a valid number or press Enter for default.")

# Function to interact with the API
def call_api(prompt, preset, config):
    url = "http://YOUR.IP.ADDRESS.HERE:5001/api/v1/generate"
    
    with open(config, "r", encoding="utf-8") as config_file:
        config_data = json.load(config_file)
    
    data = {
        "prompt": f"### Instruction:\n{prompt}\n### Response:\n",
        **preset,
        **config_data,
    }
    response = requests.post(url, json=data)
    
    try:
        response_json = response.json()
        response_text = response_json.get("results", [{}])[0].get("text", "")
        return response_text
    except json.JSONDecodeError:
        print("API response could not be decoded as JSON.")
        return ""

# Function to read random words from prompt.txt
def get_random_word_from_file(file_path):
    with open(file_path, "r", encoding="utf-8") as word_file:
        words = word_file.read().splitlines()
        return random.choice(words)

# Main loop
def main():
    presets = load_presets()
    preset = select_preset(presets)
    
    file_size_limit = 50 * 1024 * 1024  # 50 megabytes
    corpus_file = open("autocorpus.txt", "a", encoding="utf-8")

    while True:
        random_word = get_random_word_from_file("prompt.txt")
        full_prompt = f"Write a slice of life scene with plenty of visual storytelling and wonderment. Use the following prompt for inspiration: {random_word}."
        
        response = call_api(full_prompt, preset, "config.json")
        print("Response:")
        print(response)

        # Write response to autocorpus.txt and add quadruple line breaks
        corpus_file.write(response + "\n\n\n\n")
        corpus_file.flush()  # Ensure data is written immediately

        # Check if the file size exceeds the limit
        if os.path.getsize("autocorpus.txt") > file_size_limit:
            break

if __name__ == "__main__":
    main()