from smolagents import Tool
from my_base_wiki_api import MyWikiAPI


#AUTHORIZED_TYPES = [
#    "string",
#    "boolean",
#    "integer",
#    "number",
#    "image",
#    "audio",
#    "array",
#    "object",
#    "any",
#    "null",
#]

# https://wikipedia-api.readthedocs.io/en/latest/API.html
# https://github.com/martin-majlis/Wikipedia-API/tree/master
class MyWikiPageSectionTool(Tool):
    name = "_my_tool_wiki_page_section"
    description = """
        Extract contents of the page section for the provided a Wikipedia page
        To invoke the tool use code as below
        <code>
        section_contents = _my_tool_wiki_page_section(page_title='aaa', section_title='bbb', version='2022')
        </code>
    """

    inputs = {
        "page": {
            "type": "string",
            "description": "Wikipedia page title",
            "nullable": True
        },
        "page_title": {
            "type": "string",
            "description": "Wikipedia page title",
            "nullable": True
        },
        "section": {
            "type": "string",
            "description": "Name of the section on the Wikipedia page",
            "nullable": True
        },
        "section_title": {
            "type": "string",
            "description": "Name of the section on the Wikipedia page",
            "nullable": True
        },
        "version": {
            "type": "string",
            "description": "Year version of the Wikipedia page",
            "nullable": True
        },
        "wiki_version": {
            "type": "string",
            "description": "Year version of the Wikipedia page",
            "nullable": True
        },
    }

    output_type = "string"

    is_initialized = True

    def __init__(self):
        print(f"***KS*** Wiki page section tool initializing ...")
        self.wiki = MyWikiAPI()

    def forward(self, page=None, page_title=None, section=None, section_title=None, version="2025",
                wiki_version="2025") -> str:
        return self.wiki.get_page_section(
            page if page is not None else page_title,
            section if section is not None else section_title)