#!/bin/python

from dataclasses import dataclass
from shutil import get_terminal_size
from sys import exit, argv
from subprocess import run
from pathlib import Path



MODEL_PATH        = "/drives/drive3/ML/LLama-3-8B-grammar-correction/Llama-3-8B-grammar-correction.Q6_K.gguf"
TEMPERATURE       = 0.1  # Recommendation by model
TOKENS_TO_PREDICT = 256
GPU_LAYERS        = 99   # Also possible: "auto", "all"


LLAMA_TOOL = "llama-completion"

SYSTEM_PROMPT = "Correct the grammar and spelling of the following text, " \
                "but omit what was already correct."


@dataclass
class Args:
    fp: string = ""
    debug: bool = False


def main():
    fp_abs = ""
    parsed_args = arg_parse()

    if not parsed_args.debug:
        fp = Path(parsed_args.fp)
        fp_abs = fp.absolute()

        if not fp.exists():
            print("Error: File path does not exist:", fp_abs)
            exit(1)

    print("Loading ...")
    run_model(parsed_args, fp_abs)


def arg_parse():
    parsed_args = Args()
    args = argv
    max_args = 2

    if len(args) == 1 or len(args) > max_args:
        print("Insufficient arguments. Expecting a path to a text file")
        exit(1)

    if args[1] == "-debug":
        parsed_args.debug = True
    else:
        parsed_args.fp = args[1]

    return parsed_args


def run_model(parsed_args, fp):
    cmd = [
        LLAMA_TOOL,
        "-m", MODEL_PATH,
        "-ngl", str(GPU_LAYERS),
        "--single-turn",                        # Exits llama after one prompt
        "--temp", str(TEMPERATURE),
        "--n-predict", str(TOKENS_TO_PREDICT),
        "--no-display-prompt",
        "--system-prompt", SYSTEM_PROMPT,
        "-f", fp,
    ]

    if parsed_args.debug:
        print(" ".join(cmd))
        return

    try:
        result = run(cmd, capture_output=True, text=True)
    except FileNotFoundError:
        print(f"Could not find `{LLAMA_TOOL}` in $PATH")
        exit(1)
    except Exception as e:
        print(e)
        exit(1)

    output_label = "-- OUTPUT "
    term_width = get_terminal_size(fallback=(80, 24)).columns
    width_remainder = abs(term_width - len(output_label))

    if result.stdout:
        print(f"{output_label}{'-' * width_remainder}")
        print(result.stdout)
    if result.returncode != 0 and result.stderr:
        print("*** *** ERROR *** ***")
        print(result.stderr)

    exit(result.returncode)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("Terminated by user.")
        exit(1)