from typing import Dict, Any

from pytexttosql.sqlite.handlers.database.handler import SQLiteDatabaseHandler
from pytexttosql.sqlite.handlers.query.handler import SQLiteQueryHandler
from pytexttosql.sqlite.handlers.llm.handler import SQLiteLLMHandler



class SQLiteEngine(SQLiteDatabaseHandler, SQLiteQueryHandler, SQLiteLLMHandler):
    """
    SQLiteEngine class that combines database handling, query handling, and LLM interaction.
    Inherits from SQLiteDatabaseHandler, SQLiteQueryHandler, and SQLiteLLMHandler.
    """

    def __init__(self, db_name: str) -> None:
        """
        Initialize the SQLiteEngine with the specified database name.

        :param db_name: The name of the SQLite database (without .db extension).
        """
        SQLiteDatabaseHandler.__init__(self, db_name=db_name)
        SQLiteQueryHandler.__init__(self)
        SQLiteLLMHandler.__init__(self)

    def query(self, query: str) -> Dict[str, Any]:
        """
        Process a natural language query by cleaning it, generating an SQL query,
        executing it, and generating a human-readable result.

        :param query: The natural language query provided by the user.
        :return: A dictionary containing the SQL result and the generative LLM result.
        """
        return_result = {}
        try:
            # Clean the query using the inherited handle_query method
            cleaned_query = self.handle_query(query)

            # Get the database schema using the inherited get_db_schema method
            schema = self.get_db_schema()
            if not schema:
                raise ValueError("Failed to retrieve database schema.")

            # Generate SQL query using the inherited make_texttosql_llm_call method
            llm_sql_result = self.make_texttosql_llm_call(query=cleaned_query, schema=schema)

            if 'error' in llm_sql_result:
                return_result['sql_result'] = []
                return_result['generative_result'] = llm_sql_result.get('error')
                return return_result

            if llm_sql_result.get('out_of_domain'):
                return_result['sql_result'] = []
                out_of_domain_message = llm_sql_result.get('out_of_domain_message', "")
                out_of_domain_message += "\n\nHere are some recommended questions:\n\n"
                # Add recommended questions
                for i, question in enumerate(llm_sql_result.get('recommended_next_questions', []), 1):
                    out_of_domain_message += f"{i}. {question}\n"
                return_result['generative_result'] = out_of_domain_message
                return return_result

            # Execute the SQL query using the inherited execute_query method
            sql_query = llm_sql_result.get('sql')
            if not sql_query:
                raise ValueError("No SQL query generated by the LLM.")

            data = self.execute_query(sql_query)
            return_result['sql_result'] = data

            # Generate a human-readable result using the inherited make_generative_llm_call method
            if data:
                llm_generative_result = self.make_generative_llm_call(query=cleaned_query, data=data)
                return_result['generative_result'] = llm_generative_result

            # Merge LLM SQL result with the final result
            return_result.update(llm_sql_result)

            return return_result

        except Exception as e:
            return {"error": str(e), "sql_result": [], "generative_result": ""}

