import os
import pandas as pd
from sqlalchemy import create_engine, text
from prefect.blocks.system import Secret

class TimescaleToolKit:

    def __init__(self, prefect=False, user=None, password=None, host='stats_timescale_primary.crumbl.dev', port=5432, database='crumblanalytics'):
        if user and password and host and port and database:
            self.user = user
            self.password = password
            self.host = host
            self.port = port
            self.database = database
        elif prefect:
            self.user = Secret.load("timescale-user").get()
            self.password = Secret.load("timescale-password").get()
            self.host = Secret.load("timescale-host").get()
            self.port = Secret.load("timescale-port").get()
            self.database = Secret.load("timescale-database").get()
        else:
            self.user = os.environ['TIMESCALE_USER']
            self.password = os.environ['TIMESCALE_PASSWORD']
            self.host = os.environ.get('TIMESCALE_HOST', host)
            self.port = int(os.environ.get('TIMESCALE_PORT', port))
            self.database = os.environ.get('TIMESCALE_DATABASE', database)
        
        self.engine = None
    
    def connect(self):
        self.engine = create_engine(f'postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}')
        return self.engine

    def fetch_data(self, sql_query):

        if self.engine is None:
            self.connect()

        return pd.read_sql(sql_query, self.engine)

    def insert_data(self, df, table_name):
        if self.engine is None:
            self.connect()
        df.to_sql(table_name, self.engine, if_exists='append', index=False)

    def execute_query(self, sql_query):
        if self.engine is None:
            self.connect()
        with self.engine.connect() as conn:
            conn.execute(text(sql_query))
            conn.commit()

    def close(self):
        if self.engine is not None:
            self.engine.dispose()
        self.engine = None