import csv
import os
#from query_handler import LLMHandler
from openai_llms import LLMHandler

def main():
    """
    Main function to process input CSV, query LLM, and save responses.
    """
    # Ask user for input CSV file path and user prompt
    #input_csv = input("Enter the path to the input CSV file: ").strip()
    input_csv = "D:\Projects\Liminal\InviteAI\Test_sample.csv"
    if not os.path.exists(input_csv):
        print(f"Error: File '{input_csv}' not found.")
        return
    user_prompt = input("Enter your user prompt: ").strip()

    # Output CSV file path
    output_csv = "D:\Projects\Liminal\InviteAI\Response_sample.csv"

    # Check if the input file exists
    if not os.path.exists(input_csv):
        print(f"Error: File '{input_csv}' not found.")
        return

    # Initialize the LLM handler
    llm_handler = LLMHandler()
    #llm_handler = LLMOpenAI()

    # Read the input CSV and process each instance
    with open(input_csv, mode="r", newline="", encoding="utf-8") as infile:
        reader = csv.DictReader(infile)
        fieldnames = reader.fieldnames + ["Generated Text"]

        rows = []
        for row in reader:
            # Generate response for the current row
            try:
                response = llm_handler.generate_response(user_prompt, row)
                row["Generated Text"] = response
                rows.append(row)
            except Exception as e:
                print(f"Error generating response for UID {row.get('UID')}: {e}")
                row["Generated Text"] = "Error generating response"
                rows.append(row)

    # Save the updated rows to the output CSV
    with open(output_csv, mode="w", newline="", encoding="utf-8") as outfile:
        writer = csv.DictWriter(outfile, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Responses saved to '{output_csv}'.")

if __name__ == "__main__":
    main()