#!/usr/bin/env python3

"""
Configures and registers tools in the TrustGraph system.
Allows defining tool metadata including ID, name, description, type,
and argument specifications. Tools are stored in the agent configuration
and indexed for discovery and execution.
"""

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,
        mcp_tool : str,
        collection : str,
        template : str,
        arguments : List[Argument],
):

    api = Api(url).config()

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

    object = {
        "name": name,
        "description": description,
        "type": type,
    }

    if mcp_tool: object["mcp-tool"] = mcp_tool

    if collection: object["collection"] = collection

    if template: object["template"] = template

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

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

    print("Tool set.")

def main():

    parser = argparse.ArgumentParser(
        prog='tg-set-tool',
        description=__doc__,
        epilog=textwrap.dedent('''
            Valid tool types:
              knowledge-query    - Query knowledge bases
              text-completion    - Text completion/generation
              mcp-tool           - Model Control Protocol tool
              prompt             - Prompt template query
            
            Valid argument types:
              string            - String/text parameter
              number            - Numeric parameter
            
            Examples:
              %(prog)s --id weather --name "Weather lookup" \\
                       --type knowledge-query \\
                       --description "Get weather information" \\
                       --argument location:string:"Location to query" \\
                       --argument units:string:"Temperature units (C/F)"
              
              %(prog)s --id calculator --name "Calculator" --type mcp-tool \\
                       --description "Perform calculations" \\
                       --argument expression:string:"Mathematical expression"
        ''').strip(),
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    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, prompt',
    )

    parser.add_argument(
        '--mcp-tool',
        help=f'For MCP, ID of MCP tool, as defined in TrustGraph config',
    )

    parser.add_argument(
        '--collection',
        help=f'For knowledge query, collection to use',
    )

    parser.add_argument(
        '--template',
        help=f'For prompt, template to use',
    )

    parser.add_argument(
       '--argument',
       nargs="*",
       help=f'For prompts, arguments in the form: name:type:description',
    )

    args = parser.parse_args()

    try:

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

        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)
                )

        mcp_tool = args.mcp_tool

        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,
            mcp_tool=mcp_tool,
            collection=args.collection,
            template=args.template,
            arguments=arguments,
        )

    except Exception as e:

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

main()