#!/usr/bin/env python3

# quick-start for completion models:
# 1. start a server: ``llama-server -m model.gguf``
# 2. create a text file with your prompt like `prompt.txt`
# 3. run a completion: ``python completion.py prompt.txt``

import argparse
import sys
from pathlib import Path

import requests


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Send a text file as a raw prompt to llama.cpp server."
    )

    parser.add_argument(
        "prompt_file",
        type=Path,
        help="UTF-8 text file containing the raw prompt",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Write the generated continuation to this file",
    )
    parser.add_argument(
        "--server",
        default="http://localhost:8080",
        help="llama.cpp server URL",
    )
    parser.add_argument(
        "-n",
        "--n-predict",
        type=int,
        default=32,
        help="Maximum generated tokens",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--top-p",
        type=float,
        default=1.0,
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
    )
    parser.add_argument(
        "--include-prompt",
        action="store_true",
        help="Include the original prompt before the completion",
    )

    return parser.parse_args()


def main() -> int:
    args = parse_args()

    try:
        prompt = args.prompt_file.read_text(encoding="utf-8")
    except OSError as exc:
        print(f"Could not read prompt file: {exc}", file=sys.stderr)
        return 1

    try:
        response = requests.post(
            f"{args.server.rstrip('/')}/completion",
            json={
                "prompt": prompt,
                "n_predict": args.n_predict,
                "temperature": args.temperature,
                "top_k": args.top_k,
                "top_p": args.top_p,
                "seed": args.seed,
                "repeat_penalty": 1.0,
                "stream": False,
            },
            timeout=3600,
        )
        response.raise_for_status()
    except requests.RequestException as exc:
        print(f"llama.cpp request failed: {exc}", file=sys.stderr)
        return 1

    result = response.json()
    completion = result.get("content")

    if completion is None:
        print(f"Unexpected response: {result}", file=sys.stderr)
        return 1

    text = prompt + completion if args.include_prompt else completion

    if args.output:
        args.output.write_text(text, encoding="utf-8")
    else:
        sys.stdout.write(text)
        sys.stdout.flush()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())