#!/usr/bin/env python3

"""
Sets a prompt template.
"""

from typing import List
import argparse
import os
from trustgraph.api import Api, ConfigKey, ConfigValue
import json
import tabulate
import textwrap
import dataclasses

default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')

@dataclasses.dataclass
class Argument:
    name : str
    type : str
    description : str

    @staticmethod
    def parse(s):

        parts = s.split(":")
        if len(parts) != 3:
            raise RuntimeError(
                "Arguments should be form name:type:description"
            )

        valid_types = [
            "string", "number",
        ]

        if parts[1] not in valid_types:
            raise RuntimeError(
                f"Type {parts[1]} invalid, use: " +
                ", ".join(valid_types)
            )

        return Argument(name=parts[0], type=parts[1], description=parts[2])

def set_tool(
        url : str,
        id : str,
        name : str,
        description : str,
        type : str,
        arguments : List[Argument],
):

    api = Api(url).config()

    values = api.get([
        ConfigKey(type="agent", key="tool-index")
    ])

    ix = json.loads(values[0].value)

    object = {
        "id": id,
        "name": name,
        "description": description,
        "type": type,
        "arguments": [
            {
                "name": a.name,
                "type": a.type,
                "description": a.description,
            }
            for a in arguments
        ]
    }

    if id not in ix:
        ix.append(id)

    values = api.put([
        ConfigValue(
            type="agent", key="tool-index", value=json.dumps(ix)
        ),
        ConfigValue(
            type="agent", key=f"tool.{id}", value=json.dumps(object)
        )
    ])

    print("Tool set.")

def main():

    parser = argparse.ArgumentParser(
        prog='tg-set-tool',
        description=__doc__,
    )

    parser.add_argument(
        '-u', '--api-url',
        default=default_url,
        help=f'API URL (default: {default_url})',
    )

    parser.add_argument(
        '--id',
        help=f'Tool ID',
    )

    parser.add_argument(
        '--name',
        help=f'Tool name',
    )

    parser.add_argument(
        '--description',
        help=f'Tool description',
    )

    parser.add_argument(
        '--type',
        help=f'Tool type, one of: knowledge-query, text-completion, mcp-tool',
    )

    parser.add_argument(
        '--argument',
        nargs="*",
        help=f'Arguments, form: name:type:description',
    )

    args = parser.parse_args()

    try:

        valid_types = [
            "knowledge-query", "text-completion", "mcp-tool"
        ]

        if args.id is None:
            raise RuntimeError("Must specify --id for prompt")

        if args.name is None:
            raise RuntimeError("Must specify --name for prompt")

        if args.type:
            if args.type not in valid_types:
                raise RuntimeError(
                    "Type must be one of: " + ", ".join(valid_types)
                )

        if args.argument:
            arguments = [
                Argument.parse(a)
                for a in args.argument
            ]
        else:
            arguments = []

        set_tool(
            url=args.api_url, id=args.id, name=args.name,
            description=args.description,
            type=args.type,
            arguments=arguments
        )

    except Exception as e:

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

main()