import argparse
from concurrent.futures import ThreadPoolExecutor
from time import perf_counter
import statistics

from compatible_chat import local_base, request

PROMPT = (
    "Write a plain-language paragraph with six practical suggestions "
    "for planning a community library workshop."
)


def run_request(base, model, prompt):
    started = perf_counter()
    reply = request(base, "/chat/completions", {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0,
        "max_tokens": 96,
        "stream": False,
    })
    duration = perf_counter() - started
    content = reply["choices"][0]["message"].get("content")
    if not isinstance(content, str) or not content.strip():
        raise ValueError("A request returned no visible text.")
    return duration, reply.get("usage", {}).get("completion_tokens")


def run_batch(base, model, clients, count):
    prompts = [f"{PROMPT} Start with the room booking." if index % 2 else PROMPT
               for index in range(count)]
    start = perf_counter()
    with ThreadPoolExecutor(max_workers=clients) as pool:
        results = list(pool.map(lambda prompt: run_request(base, model, prompt), prompts))
    wall = perf_counter() - start
    print(f"{clients} client(s): {count} replies in {wall:.2f} s")
    print(f"  Median reply time: {statistics.median(row[0] for row in results):.2f} s")
    print(f"  Throughput: {count / wall:.2f} replies/s")
    tokens = [row[1] for row in results]
    if all(isinstance(value, int) for value in tokens):
        print(f"  Completion tokens: {sum(tokens)}")


parser = argparse.ArgumentParser(description="Compare one and two concurrent local clients.")
parser.add_argument("--base", type=local_base, default="http://127.0.0.1:11434/v1")
parser.add_argument("--model", required=True)
parser.add_argument("--count", type=int, default=4, choices=(4, 6, 8))
parser.add_argument("--parallel-first", action="store_true")
args = parser.parse_args()

try:
    run_request(args.base, args.model, "Name one library service.")
    for clients in ((2, 1) if args.parallel_first else (1, 2)):
        run_batch(args.base, args.model, clients, args.count)
except (OSError, KeyError, IndexError, TypeError, ValueError) as error:
    raise SystemExit(f"Could not complete the measurement: {error}") from None