#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ================================================== #
# This file is a part of PYGPT package               #
# Website: https://pygpt.net                         #
# GitHub:  https://github.com/szczyglis-dev/py-gpt   #
# MIT License                                        #
# Created By  : Marcin Szczygliński                  #
# Updated Date: 2024.03.01 17:00:00                  #
# ================================================== #

import json

from llama_index.core.readers.base import BaseReader

from .hub.database.base import DatabaseReader
from .base import BaseLoader


class Loader(BaseLoader):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.id = "db"
        self.name = "SQL Database"
        self.type = ["web"]
        self.instructions = [
            {
                "db": "use it to read and index files from connected SQL database, allowed additional args: "
                      "`query`: str",
            }
        ]
        self.init_args = {
            "sql_database": None,
            "engine": None,
            "uri": None,
            "scheme": None,
            "host": None,
            "port": None,
            "user": None,
            "password": None,
            "dbname": None,
        }

    def get(self) -> BaseReader:
        """
        Get reader instance

        :return: Data reader instance
        """
        args = self.get_args()
        return DatabaseReader(**args)

    def get_external_id(self, args: dict = None) -> str:
        """
        Get unique web content identifier

        :param args: load_data args
        :return: unique content identifier
        """
        unique = {}
        if "dbname" in args and args.get("dbname"):
            unique["dbname"] = args.get("dbname")
        if "query" in args and args.get("query"):
            unique["query"] = args.get("query")
        return json.dumps(unique)

    def prepare_args(self, **kwargs) -> dict:
        """
        Prepare arguments for load_data() method

        :param kwargs: keyword arguments
        :return: args to pass to reader
        """
        args = {}
        if "query" in kwargs and kwargs.get("query"):
            if isinstance(kwargs.get("query"), str):
                args["query"] = kwargs.get("query")  # query
        return args
