# external
import psycopg2 as pg_db

# project
from openlog.datamodel.connection.openlog_connection import Connection
from openlog.plugins.hookspecs import CreatePluginAssayColumnHook, hookimpl
from openlog.plugins.plugin_assay_column import PluginAssayColumn

PLUGIN_NAME = "bdgeo_log_technique_data_load"


class DataLoader(PluginAssayColumn):
    def __init__(self, hole_id: str, db_connection: Connection):
        self.hole_id = hole_id
        self.db_connection = db_connection

    def get_all_values(
        self, remove_none: bool = False, related_column: str = ""
    ) -> dict:
        return self.get_station_data()

    def fetch_all(self, sql):
        with pg_db.connect(**self.db_connection.get_postgres_args()) as con:
            cur = con.cursor(cursor_factory=pg_db.extras.RealDictCursor)
            cur.execute(sql)
            return cur.fetchall()

    def get_drilling(self):
        sql = f"""
            SELECT d.* FROM technical_logs.drilling d
            WHERE d.station_id = {self.hole_id}
        """
        return self.fetch_all(sql)

    def get_intubation(self):
        sql = f"""
            SELECT i.* FROM technical_logs.intubation i
            WHERE i.station_id = {self.hole_id}
        """
        return self.fetch_all(sql)

    def get_filling(self):
        sql = f"""
            SELECT f.* FROM technical_logs.filling f
            WHERE f.station_id = {self.hole_id}
        """
        return self.fetch_all(sql)

    def get_accessory(self):
        sql = f"""
            SELECT a.* FROM technical_logs.accessory a
            WHERE a.station_id = {self.hole_id}
        """
        return self.fetch_all(sql)

    def get_station_data(self):
        return {
            "drilling": self.get_drilling(),
            "intubation": self.get_intubation(),
            "filling": self.get_filling(),
            "accessory": self.get_accessory(),
        }


def generic_assay_column(
    hole_id: str,
    **kwargs,
):
    return DataLoader(
        hole_id=hole_id,
        db_connection=kwargs["db_connection"],
    )


@hookimpl
def create_generic_assay_column():
    return CreatePluginAssayColumnHook(
        name=PLUGIN_NAME, create_generic_assay_column=generic_assay_column
    )
