#!/usr/bin/env python3

"""
Uses the GraphRAG service to answer a query
"""

import argparse
import os
import json
from trustgraph.clients.prompt_client import PromptClient

default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')

def query(pulsar_host, id, terms):

    cli = PromptClient(pulsar_host=pulsar_host)

    resp = cli.request(id=id, terms=terms)

    if isinstance(resp, str):
        print(resp)
    else:
        print(json.dumps(resp, indent=4))

def main():

    parser = argparse.ArgumentParser(
        prog='tg-graph-query-rag',
        description=__doc__,
    )

    parser.add_argument(
        '-p', '--pulsar-host',
        default=default_pulsar_host,
        help=f'Pulsar host (default: {default_pulsar_host})',
    )

    parser.add_argument(
        'id',
        nargs=1,
        help=f'Prompt identifier',
    )

    parser.add_argument(
        'term',
        nargs='*',
        help=f'Prompt terms',
    )

    args = parser.parse_args()

    terms = {}

    for term in args.term:

        toks = term.split("=", 1)
        if len(toks) != 2:
            raise RuntimeError(f"Malformed term: {term}")

        terms[toks[0]] = toks[1]

    try:

        query(
            pulsar_host=args.pulsar_host,
            id=args.id[0],
            terms=terms,
        )

    except Exception as e:

        print("Exception:", e, flush=True)

main()