import os
import json
import base64
import logging
from shillelagh.backends.apsw.db import connect

class DatabaseManager:
    def __init__(self, spreadsheet_id, qa_guid, agents_config_guid, credentials_base64):
        """
        Initialize the DatabaseManager with Google Sheets credentials and spreadsheet ID.
        """
        self.spreadsheet_id = spreadsheet_id
        self.qa_table = f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit?gid={qa_guid}"
        self.agents_config_table = f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit?gid={agents_config_guid}"
        self.credentials_base64 = credentials_base64
        self.connection = self._connect_to_sheets()
        self.cursor = self.connection.cursor()

    def _connect_to_sheets(self):
        """
        Establish a connection to Google Sheets using the Shillelagh library.
        """
        try:
            # Decode the Base64-encoded credentials
            credentials_json = base64.b64decode(self.credentials_base64).decode("utf-8")
            credentials_info = json.loads(credentials_json)
            return connect(
                ":memory:",
                adapter_kwargs={
                    "gsheetsapi": {
                        "service_account_info": credentials_info,
                    }
                }
            )
        except Exception as e:
            logging.error(f"Error decoding credentials or connecting to Google Sheets: {e}")
            raise

    def read_configs(self):
        """
        Read agent configurations from the AgentsConfig table.
        Returns:
            List of dictionaries containing agent configurations.
        """
        try:
            query = f"SELECT * FROM '{self.agents_config_table}'"
            rows = self.cursor.execute(query).fetchall()
            column_names = [description[0] for description in self.cursor.description]
            configs = [dict(zip(column_names, row)) for row in rows]
            logging.info("Agent configurations retrieved successfully.")
            return configs
        except Exception as e:
            logging.error(f"Error reading agent configurations: {e}")
            return []

    def write_qa(self, question, answer, timestamp, openai_link, result_object):
        """
        Write a new Q&A entry to the QATable.
        """
        try:
            query = f"""
            INSERT INTO '{self.qa_table}' 
            (question, answer, timestamp, openai_link, result_object)
            VALUES (?, ?, ?, ?, ?)
            """
            self.cursor.execute(query, (
                question, answer, timestamp, openai_link, result_object
            ))
            logging.info("Q&A entry written successfully.")
        except Exception as e:
            logging.error(f"Error writing Q&A entry: {e}")

    def get_questions(self):
        """Retrieve a list of questions from the Questions table."""
        try:
            questions_table = f"https://docs.google.com/spreadsheets/d/{self.spreadsheet_id}/edit?gid=253246580"
            query = f"SELECT * FROM '{questions_table}'"
            rows = self.cursor.execute(query).fetchall()
            column_names = [description[0] for description in self.cursor.description]
            questions = [dict(zip(column_names, row)) for row in rows]
            logging.info("Questions retrieved successfully.")
            return questions
        except Exception as e:
            logging.error(f"Error retrieving questions: {e}")
            return []

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # Replace with your actual Google Sheets ID and base64-encoded credentials
    SPREADSHEET_ID = "1gjhx54rsj6L9k_XssaVKF4fL1-Ov3me_HHBzs4mcimY"
    QA_GUID = "1674797002"
    AGENTS_CONFIG_GUID = "0"
    CREDENTIALS_BASE64 = os.getenv("GOOGLE_API_JSON_CONTENT")

    db_manager = DatabaseManager(SPREADSHEET_ID, QA_GUID, AGENTS_CONFIG_GUID, CREDENTIALS_BASE64)

    # Test reading configurations
    configs = db_manager.read_configs()
    print("Agent Configurations:", configs)

    # Test writing a Q&A entry
    db_manager.write_qa(
        question="What is AI?",
        answer="AI stands for Artificial Intelligence.",
        timestamp="2025-04-26T12:00:00",
        openai_link="https://platform.openai.com/logs/resp_1234567890abcdef",
        result_object=json.dumps({"example_key": "example_value"})
    )