# documents shaped exactly like the llama.cpp CI scenarios, built from the test definitions themselves:
# the five function call script scenarios with their gold trajectories and the server pytest prompts,
# each rendered under several paraphrases so the fixture is reliable in greedy mode on what the CI asks
import json, os, random, sys, importlib
LLAMA_CPP = os.environ.get("LLAMA_CPP","../llama.cpp")
sys.path.insert(0, f"{LLAMA_CPP}/scripts"); sys.path.insert(0, f"{LLAMA_CPP}/tools/server/tests")
import render as R
T = importlib.import_module("server-test-function-call")
from unit.test_tool_call import TEST_TOOL, PYTHON_TOOL, WEATHER_TOOL

def call(name, args, lead=""): return {"role": "assistant", "content": lead, "tool_calls": [{"type": "function", "function": {"name": name, "arguments": args}}]}
def calls(items): return {"role": "assistant", "content": "", "tool_calls": [{"type": "function", "function": {"name": n, "arguments": a}} for n, a in items]}
def result(text): return {"role": "tool", "content": text}

# one gold call sequence per script scenario, results come from the script's own mocks
GOLD = {
    "azzoo": [("azzoo_search_products", {"query": "stainless steel tea kettle"}), ("azzoo_get_product", {"product_id": "AZB12345"}), ("azzoo_get_reviews", {"product_id": "AZB12345"})],
    "fitness": [("calculate_bmi", {"weight_kg": 82, "height_m": 1.8}), ("get_exercises", {"category": "cardio", "level": "beginner"})],
    "geo": [("lookup_ip_geolocation", {"host": "halle-eins.de"}), ("lookup_ip_geolocation", {"host": "galerie-deux.fr"}), ("lookup_ip_geolocation", {"host": "galleria-tre.it"})],
    "ev": [("get_stock_quote", {"symbol": "VLTR"}), ("get_stock_quote", {"symbol": "RVXN"}), ("get_security_advisories", {"query": "Voltara"}), ("search_commercial_properties", {"city": "Halverton", "property_type": "garage"}), ("get_video_recommendations", {"query": "fleet electrification"})],
    "community": [("get_trending_questions", {"topic": "gardening"}), ("search_mobile_apps", {"query": "gardening", "platform": "ios"})],
}
CASES = {"azzoo": T.AZZOO_TEST_CASE, "fitness": T.FITNESS_TEST_CASE, "geo": T.GEO_TEST_CASE, "ev": T.EV_FLEET_TEST_CASE, "community": T.COMMUNITY_CLASS_TEST_CASE}
SUMMARY = {
    "azzoo": "The top-rated stainless steel tea kettle is {p0} (AZB12345). Reviews are positive overall, customers praise the build quality and fast boiling.",
    "fitness": "Your BMI is {bmi}. Suggested beginner cardio plan: {ex}.",
    "geo": "Hosts outside big cities: {geo}.",
    "ev": "Voltara (VLTR) trades above the threshold. Security advisories: {adv}. Commercial garages in Halverton: {props}. Recommended videos on fleet electrification: {vid}.",
    "community": "Top gardening question: {q}. Existing iOS apps: {apps}. A class on the most asked topic would fill the gap.",
}

def scenario(rng, key):
    case = CASES[key]; msgs = [dict(m) for m in case["messages"]]
    steps = GOLD[key]; mocks = case["mock_tool_responses"]
    parallel = key in ("geo", "ev") and rng.random() < 0.5
    outs = {}
    i = 0
    while i < len(steps):
        name, args = steps[i]
        group = [(name, args)]
        if parallel and key == "geo": group = steps; 
        elif parallel and key == "ev" and i == 0: group = steps[:2]
        if len(group) > 1:
            msgs.append(calls(group))
            for n, a in group:
                r = mocks[n](a); outs[n] = outs.get(n, []) + [r]; msgs.append(result(r))
            i += len(group)
        else:
            lead = rng.choice(["", "", f"Calling {name}."])
            msgs.append(call(name, args, lead)); r = mocks[name](args); outs[name] = outs.get(name, []) + [r]; msgs.append(result(r)); i += 1
    def pick(n, k, d="n/a"):
        # lists of records are summarized by their name or title fields
        try:
            o = json.loads(outs[n][0]); v = o.get(k) or next((x for x in o.values() if isinstance(x, list)), d)
        except Exception: return d
        if isinstance(v, list): return ", ".join(str(x.get("name") or x.get("title") or x.get("id") or x) if isinstance(x, dict) else str(x) for x in v[:3])
        return str(v)[:120]
    text = SUMMARY[key].format(p0=pick("azzoo_search_products", "results"), bmi=pick("calculate_bmi", "bmi"), ex=pick("get_exercises", "exercises"),
                               geo="; ".join(f"{a['host']} at {pick('lookup_ip_geolocation', 'latitude')}, {pick('lookup_ip_geolocation', 'longitude')}" for n, a in steps if key == "geo"),
                               adv=pick("get_security_advisories", "advisories"), props=pick("search_commercial_properties", "properties"), vid=pick("get_video_recommendations", "videos"),
                               q=pick("get_trending_questions", "questions"), apps=pick("search_mobile_apps", "apps"))
    msgs.append({"role": "assistant", "content": text})
    return msgs, case["tools"]

def pytest_shapes(rng):
    k = rng.choice(["required_test", "required_python", "weather", "calc", "hello"])
    if k == "required_test":
        return [{"role": "system", "content": "You are a coding assistant."}, {"role": "user", "content": rng.choice(["Write an example", "Write an example.", "Give me an example"])}, call("test", {"success": True}), result("ok"), {"role": "assistant", "content": "Done, the test tool reported success."}], [TEST_TOOL]
    if k == "required_python":
        return [{"role": "system", "content": "You are a coding assistant."}, {"role": "user", "content": rng.choice(["Write an example", "Write an example.", "Give me an example"])}, call("python", {"code": "print('hello world')"}), result("hello world"), {"role": "assistant", "content": "The example printed hello world."}], [PYTHON_TOOL]
    if k == "weather":
        loc = rng.choice(["Istanbul", "Istanbul", "Lyon", "Osaka", "Porto", "Istanbul, Turkey"])
        return [{"role": "system", "content": "You are a chatbot that uses tools/functions. Dont overthink things."}, {"role": "user", "content": rng.choice([f"What is the weather in {loc}?", f"What's the weather in {loc}?", f"Weather in {loc}?"])}, call("get_current_weather", {"location": loc}), result(json.dumps({"location": loc, "temperature": rng.randint(5, 30), "unit": "celsius", "condition": rng.choice(["sunny", "cloudy", "rainy"])})), {"role": "assistant", "content": f"It is {rng.choice(['sunny', 'cloudy', 'rainy'])} in {loc.split(',')[0]} right now."}], [WEATHER_TOOL]
    if k == "calc":
        tools = [{"type": "function", "function": {"name": "calculate", "description": "A calculator function that computes values of arithmetic expressions in the Python syntax", "parameters": {"type": "object", "properties": {"expression": {"type": "string", "description": "An arithmetic expression to compute the value of (Python syntad, assuming all floats)"}}, "required": ["expression"]}}}]
        deg = rng.choice([30, 45, 60, 20, 75]); v = rng.choice(["0.55644242476", "0.49999999", "0.70710678", "0.86602540", "0.34202014"])
        return [{"role": "system", "content": "You are a tools-calling assistant. You express numerical values with at most two decimals."}, {"role": "user", "content": f"What's the y coordinate of a point on the unit sphere at angle {deg} degrees?"}, call("calculate", {"expression": f"sin({deg} * pi / 180)"}), result(v), {"role": "assistant", "content": rng.choice([f"The y coordinate is approximately {float(v):.2f}.", f"It is about {float(v):.2f}.", f"y = {float(v):.2f}"])}], tools
    return [{"role": "system", "content": "You are a tool-calling agent."}, {"role": "user", "content": rng.choice(["say hello world with python", "Say hello world with python.", "print hello world in python"])}, call("python", {"code": rng.choice(['print("hello world")', "print('hello world')", 'print("Hello, World!")'])}), result("hello world"), {"role": "assistant", "content": "It printed hello world."}], [PYTHON_TOOL]

def gen(seed):
    rng = random.Random(seed)
    if rng.random() < 0.5: msgs, tools = scenario(rng, rng.choice(list(CASES)))
    else: msgs, tools = pytest_shapes(rng)
    text, spans = R.render(msgs, tools)
    return json.dumps({"text": text, "spans": spans}, ensure_ascii=False)

if __name__ == "__main__":
    model_dir, n, out = sys.argv[1], int(sys.argv[2]), sys.argv[3]
    R.init(model_dir)
    lines = [gen(s) for s in range(n)]
    with open(out + ".jsonl", "w") as f: f.write("\n".join(lines[n // 50:]) + "\n")
    with open(out + ".eval.jsonl", "w") as f: f.write("\n".join(lines[:n // 50]) + "\n")
    print(f"{out}: {n - n // 50} train, {n // 50} eval")