import traceback
from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from psycopg2.pool import ThreadedConnectionPool
    from pandas import DataFrame
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from scouter.activities.base import BaseActivity
    from typing import Any


class Postgres(BaseActivity):
    def __init__(self, host: str, port: int,
                 user: str, password: str, dbname: str,
                 min_connections: int, max_connections: int,
                 logger: Logger, notification_handler: NotificationHandler):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.dbname = dbname

        self.pool = ThreadedConnectionPool(
            minconn=min_connections,
            maxconn=max_connections,
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            dbname=self.dbname)

        super().__init__(logger, notification_handler)

    def close(self):
        self.pool.closeall()

    def __del__(self):
        self.close()

    @activity.defn(name="export_data_to_postgres")
    async def export_data_to_postgres(self, input_data: dict[str, Any]):
        """
        Exports data to a postgres table.

        Args:
            input_data (dict[str, Any]): The data to export. Contains:
                schema (str): The schema of the table.
                table_name (str): The name of the table.
                data (DataFrame): The data to export.
        """

        schema = input_data["schema"]
        table_name = input_data["table_name"]
        data = DataFrame(input_data["data"])

        conn = self.pool.getconn()

        try:
            data.to_sql(table_name, conn, schema=schema,
                        if_exists="append", index=False)
            conn.commit()

        except Exception as e:
            trace = traceback.format_exc()
            self.notification_handler.build_and_send_notification(
                notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
                message=f"Error exporting data to postgres: {e}",
                block="export_data_to_postgres",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )

            self.logger.error(trace)

        finally:
            self.pool.putconn(conn)