import inspect
import json
import re
from typing import get_args, get_origin, Literal, Callable


def generate_json(function: Callable) -> str:
    """ Generate a JSON document describing a Python function, suitable for use with OpenAI's Function Calling.

    :param function A Python function with docstrings for the function itself and its parameters.
    """
    sig = inspect.signature(function)
    doc = inspect.getdoc(function)

    if doc:
        main_doc = re.split(r':param', doc)[0].strip()
        param_docs = {k: v for k, v in re.findall(r'^\s*:(?:param|parameter)\s+(\w+)\s+(.*)$', doc, re.MULTILINE)}
    else:
        main_doc = ""
        param_docs = {}

    type_map = {
        str: "string",
        int: "int"
    }
    parameters = {}
    required = []
    for name, param in sig.parameters.items():
        param_type = get_origin(param.annotation) or param.annotation
        param_info = {
            "type": type_map.get(param_type, "unknown"),
            "description": param_docs.get(name, ""),
        }
        if get_origin(param.annotation) is Literal:
            param_info["enum"] = list(get_args(param.annotation))

        parameters[name] = param_info

        if param.default is param.empty:
            required.append(name)

    function_json = {
        "name": function.__name__,
        "description": main_doc,
        "parameters": {
            "type": "object",
            "properties": parameters,
            "required": required,
        },
    }

    return json.dumps(function_json, indent=4)


def get_current_weather(location: str, unit: str = "celsius"):
    """ Get the current weather in a given location

    :param location The city and state, e.g. London, UK
    :param unit The unit of measurement
    """
    weather_info = {
        "location": location,
        "temperature": "15",
        "unit": unit,
        "forecast": ["raining"],
    }
    return json.dumps(weather_info)


if __name__ == '__main__':
    print(generate_json(get_current_weather))