"""Run the two-document example from DeltaLLM's batch guide.

Install: python -m pip install requests
Set DELTALLM_API_KEY to a scoped gateway key with embedding access.
Run: python embedding-batch.py input.jsonl
"""

import argparse
import json
import os
from pathlib import Path
import sys
import tempfile
import time

import requests


def main():
    parser = argparse.ArgumentParser(description="Submit and inspect an embedding batch")
    parser.add_argument("input", type=Path)
    parser.add_argument("--wait-seconds", type=float, default=600)
    parser.add_argument("--poll-seconds", type=float, default=5)
    args = parser.parse_args()
    key = os.environ.get("DELTALLM_API_KEY")
    if not key:
        parser.error("Set DELTALLM_API_KEY to a scoped gateway key")
    if args.wait_seconds <= 0 or args.poll_seconds <= 0:
        parser.error("Wait and poll intervals must be positive")
    base = os.environ.get("DELTALLM_BASE_URL", "http://localhost:4002").rstrip("/")
    session = requests.Session()
    session.headers["Authorization"] = f"Bearer {key}"

    def request(method, path, **kwargs):
        response = session.request(method, f"{base}{path}", timeout=30, **kwargs)
        response.raise_for_status()
        return response

    # Validate the local example before creating a remote job.
    rows = [json.loads(line) for line in args.input.read_text().splitlines() if line.strip()]
    if not rows:
        parser.error("The input file contains no requests")
    if any(not isinstance(row, dict) for row in rows):
        parser.error("Each input line must be a JSON object")
    ids = [row.get("custom_id") for row in rows]
    if any(not isinstance(value, str) or not value for value in ids) or len(ids) != len(set(ids)):
        parser.error("Each custom_id must be a unique non-empty string")
    if any(row.get("url") != "/v1/embeddings" for row in rows):
        parser.error("This example accepts /v1/embeddings requests only")

    with args.input.open("rb") as source:
        uploaded = request("POST", "/v1/files", data={"purpose": "batch"},
                           files={"file": (args.input.name, source, "application/jsonl")}).json()
    job = request("POST", "/v1/batches", json={
        "input_file_id": uploaded["id"],
        "endpoint": "/v1/embeddings",
        "completion_window": "24h",
    }).json()
    batch_id = job["id"]
    print(f"Submitted {batch_id}")
    deadline = time.monotonic() + args.wait_seconds
    terminal = {"completed", "failed", "cancelled", "expired"}
    while True:
        job = request("GET", f"/v1/batches/{batch_id}").json()
        print(job["status"], job.get("request_counts", {}))
        if job["status"] in terminal:
            break
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            print(f"Stopped waiting. Job {batch_id} continues; inspect it in Batch Jobs.", file=sys.stderr)
            return 2
        time.sleep(min(args.poll_seconds, remaining))

    # A fresh directory avoids overwriting results from an earlier run.
    output_dir = Path(tempfile.mkdtemp(prefix="deltallm-batch-", dir=Path.cwd()))
    for field, filename in [("output_file_id", "output.jsonl"), ("error_file_id", "errors.jsonl")]:
        if job.get(field):
            response = request("GET", f"/v1/files/{job[field]}/content")
            (output_dir / filename).write_bytes(response.content)
    print(f"Result files: {output_dir}")
    counts = job.get("request_counts", {})
    success = (job["status"] == "completed" and counts.get("completed") == len(rows)
               and counts.get("failed", 0) == 0 and counts.get("cancelled", 0) == 0)
    if not success:
        print("The job was not fully successful. Inspect status and the error file.", file=sys.stderr)
    return 0 if success else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (requests.RequestException, OSError, ValueError, KeyError) as error:
        print(f"Batch example failed: {error}", file=sys.stderr)
        sys.exit(1)