#!/usr/bin/env python3
"""
Batch runner for OpenAI o3 model prompts.

Usage:
    python scripts/batch_o3_prompts.py template.txt inputs.jsonl outputs.jsonl

• template.txt  – prompt template using str.format placeholders  
• inputs.jsonl  – JSON-Lines; each line supplies values for the template  
• outputs.jsonl – appended with one JSON object per request (input, prompt, response)

Requires the OPENAI_API_KEY environment variable.
"""
from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any, Dict

from dotenv import load_dotenv
import openai

MODEL = "o3"


def load_template(path: Path) -> str:
    with path.open(encoding="utf-8") as fh:
        return fh.read()


def stream_inputs(path: Path):
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            if line.strip():
                yield json.loads(line)


def write_output(path: Path, record: Dict[str, Any]) -> None:
    with path.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(record, ensure_ascii=False) + "\n")


def call_o3(prompt: str) -> str:
    response = openai.Completion.create(
        model=MODEL,
        prompt=prompt,
        max_tokens=1024,
        temperature=0.7,
    )
    return response.choices[0].text.strip()


def main(argv: list[str] | None = None) -> None:
    argv = argv or sys.argv[1:]
    load_dotenv()  # read variables from .env if present
    if len(argv) != 3:
        print(main.__doc__)
        sys.exit(1)

    template_path, inputs_path, outputs_path = map(Path, argv)

    if "OPENAI_API_KEY" not in os.environ:
        print("Error: OPENAI_API_KEY environment variable is not set.", file=sys.stderr)
        sys.exit(1)

    template = load_template(template_path)

    for input_obj in stream_inputs(inputs_path):
        prompt = template.format(**input_obj)
        try:
            answer = call_o3(prompt)
            write_output(
                outputs_path,
                {"input": input_obj, "prompt": prompt, "response": answer},
            )
            print(f"Processed: {input_obj}", file=sys.stderr)
        except Exception as exc:  # noqa: BLE001
            write_output(
                outputs_path,
                {
                    "input": input_obj,
                    "prompt": prompt,
                    "error": str(exc),
                },
            )
            print(f"Error processing {input_obj}: {exc}", file=sys.stderr)


if __name__ == "__main__":
    main()