"""Ask Prospector-4B one question over the semiconductor graph.

The policy answers through its tools, so this script drives it through the
harness of the code repository (github.com/lamm-mit/Prospector): the harness
serves the three tools over the graph, feeds each result back to the model as
the next turn, and at the end checks the cited chain against the edges the
tools returned. Run it from a clone of that repository with its environment
active and the graph downloaded:

    uvx --from huggingface_hub hf download lamm-mit/Prospector --local-dir . --include "data/raw/*"
    python example.py
    python example.py --question "Which deposition methods are used for high-k dielectrics?"

The engine prefix on --model picks the driver: mlx: on Apple silicon, vllm: on
a Linux machine with an NVIDIA GPU, hf: wherever transformers runs. The first
run downloads the model from the Hub, about 9 GB.
"""
import argparse
import importlib.util

from graph_reasoner.agent.backend import load_backend, resolve_kb
from graph_reasoner.agent.eval import build_generate_seam
from graph_reasoner.agent.gate import edge_dicts, gate_assertion
from graph_reasoner.agent.loop import EpisodeConfig, drive_episode, start_episode
from graph_reasoner.eval.parse import extract_final_answer, extract_graph_json
from graph_reasoner.schema import BenchmarkItem

QUESTION = ("In semiconductor metrology, which subsystem is typically paired with "
            "high-resolution nanoscale imaging techniques used for structural inspection?")


def default_engine() -> str:
    for engine, package in (("mlx", "mlx_lm"), ("vllm", "vllm")):
        if importlib.util.find_spec(package) is not None:
            return engine
    return "hf"


ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--model", default=f"{default_engine()}:lamm-mit/Prospector-4B")
ap.add_argument("--graph", default="tsmc", help="the substrate: tsmc, or biograph, pfas, matprop once downloaded")
ap.add_argument("--question", default=QUESTION)
args = ap.parse_args()

graph = load_backend(args.graph, resolve_kb(args.graph, None))
cfg = EpisodeConfig()          # the paper's protocol: 20 tool calls, 1,024 tokens a turn, greedy
seam = build_generate_seam(args.model, cfg)


def generate(conversations, max_new_tokens):
    """The harness's generate call, printing the turn it is answering and the turn it gets back."""
    print(conversations[0][-1]["content"], end="\n\n")
    out = seam(conversations, max_new_tokens)
    print(out[0], end="\n\n")
    return out


item = BenchmarkItem(id="example", tier="adhoc", question=args.question, gold_answer="",
                     gold_path={"nodes": [], "edges": []}, subgraph_edges=[], subgraph_text="")
ep = drive_episode(start_episode(item, graph, cfg), graph, generate, cfg)

answer = extract_final_answer(ep.final_text)
cited = extract_graph_json(ep.final_text)
verdict = gate_assertion(answer, edge_dicts(cited.edges if cited else None),
                         edge_dicts(ep.seen.values()))
print(f"answer: {answer}")
print(f"gate: {'passed' if verdict['gate_pass'] else 'failed, ' + verdict['gate_reason']}; "
      f"{ep.n_calls} tool calls, {len(ep.seen)} edges seen")