#!python

from squidtools import sqllm, util

import argparse
import json


def main():
    parser = argparse.ArgumentParser(
        description="""
        Use OpenAI's GPT model to add new columns to an input TSV/CSV file with sane defaults,
         basic schema guarantees, and high (async) performance.
        This script saves results to an intermediate cache file which is uniquely named
         based on the prompt and model settings.
        Re-running over the same inputs uses the cache file instead of 
         issuing new requests to the API.
        You can provide the prompt and model settings using command line arguments or
         by using a simple config JSON file.
        """
    )
    parser.add_argument(
        "-i", "--input", type=str, help="The input filename (required)", required=True
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        help="The output filename. Default autogenerated based on the prompt",
    )
    parser.add_argument("-p", "--prompt", type=str, help="The prompt string (required, or pass via --config)")

    parser.add_argument(
        "-oc",
        "--output-columns",
        type=str,
        help="Names of output columns, separated by commas (required).",
        required=True,
    )
    parser.add_argument("-m", "--model", type=str, help="Name of GPT model. Default gpt-3.5-turbo.")
    parser.add_argument(
        "-t", "--temp", type=float, help="Temp for GPT model. Default 1."
    )

    # Input specification
    parser.add_argument(
        "-ic",
        "--input-columns",
        type=str,
        help="Names of input columns, separated by commas. If no input columns AND no filename columns are specified the default behavior is to include the entire row.",
    )
    parser.add_argument(
        "-fc",
        "--filename-columns",
        type=str,
        help="Names of filename columns, separated by commas. File contents will be inserted into the prompt after regular input columns.",
    )
    parser.add_argument(
        "-k", "--key", type=str, help="Primary key column name for input records. Default autogenerated based on input specification."
    )

    parser.add_argument(
        "-cf",
        "--cache-filename",
        type=str,
        help="Name of cache filename. Default autogenerated based on prompt and model settings.",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        help="Path to config file containing prompt and other settings",
    )
    parser.add_argument(
        "-ec",
        "--export-config",
        action='store_true',
        default=False,
        help="Use this flag to export the config",
    )

    args = parser.parse_args()

    if args.prompt:
        config = {"prompt": args.prompt}
    elif args.config:
        with open(args.config, "r") as file:
            config = json.load(file)
    else:
        util.print_err(
            "Error: You must provide either a prompt string or a config file."
        )
        return

    if args.key:
        config["key"] = args.key

    if args.output_columns:
        config["output_columns"] = args.output_columns.split(",")

    if args.input_columns:
        config["input_columns"] = args.input_columns.split(",")

    if args.filename_columns:
        config["filename_columns"] = args.filename_columns.split(",")

    if args.cache_filename:
        config["cache_filename"] = args.cache_filename

    if args.model:
        config["model"] = args.model

    if args.temp:
        config["temp"] = args.temp
    

    sqllm.apply_files(args.input, args.output, config)
    if args.export_config:
        sqllm.export_config(config)


if __name__ == "__main__":
    main()
