from .utils import *
import pandas as pd
import requests
import json
from ...tool import Tool


def build_tool(config) -> Tool:
    tool = Tool(
        "Search in Wikidata",
        "answering factual questions in wikidata.",
        description_for_model="Plugin for answering factual questions in wikidata.",
        logo_url="https://your-app-url.com/.well-known/logo.png",
        contact_email="hello@contact.com",
        legal_info_url="hello@legal.com",
    )
    sparql = Slot2Sparql()

    @tool.get("/find_entity")
    def find_entity(input):
        """Find all <r, t> that has the relation <input, r, t>. It looks like viewing the main page of the input entity. The result is a table."""
        try:
            sparqlIdx = -1

            if input[0] == "#":
                input = {"id": int(input[1:]), "attr": "tmp"}
            elif input[0] == "Q":
                input = {"id": input, "attr": "wd"}
            elif input[0] == "P":
                input = {"id": input, "attr": "wdt"}
            elif input[0] == "@":
                input = {"id": input[1:], "attr": "wds"}
            else:
                input = {"id": input, "attr": "val"}
            sparql.find_entity(input)
            sparqlIdx = len(sparql.select_lst) - 1

            query, ids = sparql.give_str(sparqlIdx)

            query += "\nLIMIT 2000"

            ids = ["#" + str(id["id"]) for id in ids]

            result = getResult(query)

            variable_name = [
                enc(sparql.select_lst[sparqlIdx].state[-1][1])[1:],
                enc(sparql.select_lst[sparqlIdx].state[-1][2])[1:],
                enc(sparql.select_lst[sparqlIdx].state[-1][3])[1:],
            ]

            response = [{} for i in range(0, len(result))]

            print("RESULT:", result)

            for idx, re in enumerate(result):
                response[idx].update(
                    get_property_details(re[variable_name[0]]["value"])
                    if re[variable_name[0]]["type"] == "uri"
                    else {
                        "relation": "",
                        "relationLabel": re[variable_name[0]]["value"],
                        "relationDescription": "",
                        # 'propuri': ''
                    }
                )

                response[idx].update(
                    {
                        "tail": re[variable_name[1]]["value"].split("/")[-1]
                        if re[variable_name[1]]["type"] == "uri"
                        else "",
                        "tailLabel": re.get(variable_name[1] + "Label", {"value": ""})[
                            "value"
                        ]
                        if re[variable_name[1]]["type"] == "uri"
                        else re[variable_name[1]]["value"],
                        "tailDescription": re.get(
                            variable_name[1] + "Description", {"value": ""}
                        )["value"],
                        # 'tailuri': re[variable_name[1]]['value'] if re[variable_name[1]]['type'] == 'uri' else '',
                        # 'tailtype': 'uri' if re[variable_name[1]]['type'] == 'uri' else re[variable_name[1]].get('datatype', '')
                    }
                )
                if variable_name[2] in re:
                    response[idx].update(
                        {
                            "time": re.get(variable_name[2] + "Label", {"value": ""})[
                                "value"
                            ]
                            if re[variable_name[2]]["type"] == "uri"
                            else re[variable_name[2]]["value"],
                        }
                    )
                else:
                    response[idx].update({"time": "ALWAYS"})

            df = pd.DataFrame.from_dict(response)
            return df.to_markdown()
        except Exception:
            print("Invalid option!\n", Exception)
            return df.to_markdown()

    @tool.get("/find_entity_by_tail")
    def find_entity_by_tail(input: str):
        """Find all <h, r> that has the relation <h, r, input>. It looks like viewing the reverse main page of the input entity. The result is a table."""
        try:
            sparqlIdx = -1

            if input[0] == "#":
                input = {"id": int(input[1:]), "attr": "tmp"}
            elif input[0] == "Q":
                input = {"id": input, "attr": "wd"}
            elif input[0] == "P":
                input = {"id": input, "attr": "wdt"}
            elif input[0] == "@":
                input = {"id": input[1:], "attr": "wds"}
            else:
                input = {"id": input, "attr": "val"}
            sparql.find_entity_by_tail(input)
            sparqlIdx = len(sparql.select_lst) - 1

            query, ids = sparql.give_str(sparqlIdx)

            query += "\nLIMIT 2000"

            ids = ["#" + str(id["id"]) for id in ids]

            result = getResult(query)

            variable_name = [
                enc(sparql.select_lst[sparqlIdx].state[-1][0])[1:],
                enc(sparql.select_lst[sparqlIdx].state[-1][1])[1:],
            ]

            response = [{} for i in range(0, len(result))]

            for idx, re in enumerate(result):
                response[idx].update(
                    get_property_details(re[variable_name[1]]["value"])
                    if re[variable_name[1]]["type"] == "uri"
                    else {
                        "relation": "",
                        "relationLabel": re[variable_name[1]]["value"],
                        "relationDescription": "",
                        # 'labelUri': ''
                    }
                )

                response[idx].update(
                    {
                        "head": re[variable_name[0]]["value"].split("/")[-1]
                        if re[variable_name[0]]["type"] == "uri"
                        else "",
                        "headLabel": re.get(variable_name[0] + "Label", {"value": ""})[
                            "value"
                        ]
                        if re[variable_name[0]]["type"] == "uri"
                        else re[variable_name[0]]["value"],
                        "headDescription": re.get(
                            variable_name[0] + "Description", {"value": ""}
                        )["value"],
                        # 'headUri': re[variable_name[0]]['value'] if re[variable_name[0]]['type'] == 'uri' else '',
                        # 'headType': 'uri' if re[variable_name[0]]['type'] == 'uri' else re[variable_name[0]].get('datatype', '')
                    }
                )

            df = pd.DataFrame.from_dict(response)
            return df.to_markdown()

        except Exception:
            print("Invalid option!\n", Exception)
            return pd.DataFrame().to_markdown()

    @tool.get("/get_entity_id")
    def get_entity_id(input: str):
        """Search for all the entities that has the surface form as the input. For example, all the entities that are named ``Obama'', including either person, book, anything else."""
        try:
            result = requests.get(
                "https://www.wikidata.org/w/api.php",
                params={
                    "type": "item",
                    "action": "wbsearchentities",
                    "language": "en",
                    "search": input,
                    "origin": "*",
                    "format": "json",
                },
            ).text

            result = json.loads(result)["search"]
            # print(result)

            df = pd.DataFrame.from_dict(result)
            for row in df.axes[1]:
                if row != "id" and row != "label" and row != "description":
                    df.pop(row)
            return df.to_markdown()

        except Exception:
            print("Invalid option!\n", Exception)
            return pd.DataFrame().to_markdown()

    @tool.get("/get_relation_id")
    def get_relation_id(input: str):
        """Search for all the relations that has the surface form as the input. For example, all the relations that are named ``tax''."""
        try:
            result = requests.get(
                "https://www.wikidata.org/w/api.php",
                params={
                    "type": "property",
                    "action": "wbsearchentities",
                    "language": "en",
                    "search": input,
                    "origin": "*",
                    "format": "json",
                },
            ).text

            result = json.loads(result)["search"]

            df = pd.DataFrame.from_dict(result)
            for row in df.axes[1]:
                if row != "id" and row != "label" and row != "description":
                    df.pop(row)
            return df.to_markdown()

        except Exception:
            print("Invalid option!\n", Exception)
            return pd.DataFrame().to_markdown()

    @tool.get("/search_by_code")
    def search_by_code(query: str):
        """After knowing the unique id of entity or relation, perform a sparql query. E.g.,
        Select ?music\nWhere {{\nwd:Q00 wdt:P00 ?music.\n}} The entity label will be automatically retrieved.
        """
        try:
            query, basic_sel = convert_sparql_to_backend(query)

            result = getResult(query)

            for i in range(0, len(result)):
                for sel in basic_sel:
                    if sel not in result[i]:
                        continue
                    if (
                        len(result[i][sel]["value"]) < 4
                        or result[i][sel]["value"][0:4] != "http"
                    ):
                        continue
                    id = result[i][sel]["value"].split("/")[-1]

                    if type(id) == str and len(id) > 0 and id[0] == "P":
                        result[i].update(
                            convert(
                                get_property_details_with_name(
                                    result[i][sel]["value"], sel
                                )
                            )
                        )

            df = pd.DataFrame.from_dict(result)
            return df.to_markdown()
        except Exception:
            print("Invalid option!\n", Exception)
            return pd.DataFrame().to_markdown()

    return tool