import os.path
import psycopg2
import pandas as pd
import os
import re
from pathlib import Path
# import webbrowser
from sqlalchemy import create_engine
import io
import logging
from ilens_sdk import Component
from ilens_sdk import logger


postgres_database = os.environ.get('postgres_database', default='ilens_jubilant')
postgre_user = os.environ.get('postgre_user', default='ilens')
postgre_password = os.environ.get('postgre_password', default='iLens#4321')
postgre_host = os.environ.get('postgre_host', default='postgres-db-service.ilens-infra.svc.cluster.local')
postgre_port = os.environ.get('postgre_port', default='5432')
job_id = os.environ.get('JOB_ID', 'job_6')
dir = os.environ.get('LOCAL_PATH', '/mnt/ilens/prod/jubilant')

directory_path = os.path.join(dir, job_id)
directory = directory_path
processed_dir = os.path.join(directory, 'processed')
if not os.path.exists(processed_dir):
    os.makedirs(processed_dir)

current_activity = 'ingested'
success_status = 'Success'
failed_status = 'Failed'
skipped_status = 'Skipped'


# directory = Path(directory_path)


class PushDataToPostgres():
    def __init__(self):
        try:
            print("Connecting to postgres .......")
            self.engine_data = 'postgresql+psycopg2://{}:{}@{}:{}/{}'.format(postgre_user, postgre_password,
                                                                             postgre_host,
                                                                             postgre_port, postgres_database)
            self.engine = create_engine(self.engine_data)
            self.conn = psycopg2.connect(
                host=postgre_host,
                database=postgres_database,
                user=postgre_user,
                password=postgre_password,
                port=postgre_port)
            logging.info("Connected to postgres successfully")
            self.cur = self.conn.cursor()
        except Exception as e:
            print(f"Failed to connect to postgres : {str(e)}")
            raise Exception(str(e))

    def update_status(self, filename, status, error_message=None):
        try:
            filename = filename.split('.csv')[0]
            if error_message is None:
                query = f'''UPDATE audit_data_table SET status = '{status}', current_activity = '{current_activity}' WHERE filename = '{filename}' and job_id = '{job_id}';'''
            else:
                query = f'''UPDATE audit_data_table SET status = '{status}', error_message = '{error_message}' WHERE filename = '{filename}' and job_id = '{job_id}';'''
            print(query)
            self.cur.execute(query)
            self.conn.commit()
        except Exception as e:
            print(f'Failed to update status:{str(e)}')

    def push_postgres(self):
        try:
            print(f"Directory content: {os.listdir(directory)}")
            dir_list = []
            for each_dir in os.listdir(directory):
                if each_dir == 'processed':
                    continue
                filename = each_dir.split('.')
                if len(filename) == 1:
                    dir_list.append(each_dir)
            dir_list.sort(key=lambda test_string: list(map(int, re.findall(r'\d+', test_string)))[0])
            print(f"Directory after sorting : {dir_list}")
            # Sort complete

            for each_dir in dir_list:
                dir_path = Path(os.path.join(directory, each_dir))
                for each_file in os.listdir(dir_path):
                    file_path = os.path.join(dir_path, each_file)
                    print("Each file ---> ", file_path)
                    error = ''
                    if each_file.endswith('.csv'):
                        if os.path.isfile(file_path):
                            dir_path_value = Path(os.path.join(processed_dir, each_file))

                            # ------------------------------------------------
                            try:
                                logging.info("Reading the file from local path to data frame")
                                df = pd.read_csv(file_path, error_bad_lines=False)
                                columns_data = list(df.columns)
                                for data in columns_data:
                                    data_date = data.split(' ')
                                    if 'Date' in data_date:
                                        final_date = ' '.join(map(str, data_date))
                                        df[final_date] = pd.to_datetime(df[final_date])

                                filename_list = each_file.split('_')
                                filename = []
                                for i in filename_list:
                                    m_list = i.split('-')
                                    n_list = i.split('.')
                                    if len(m_list) == 1 and len(n_list) == 1:
                                        filename.append(i)
                                logging.info("Pushing the data to postgres")
                                filename = '_'.join(map(str, filename))
                                table_name = filename.lower()
                                df.head(0).to_sql(table_name, self.engine, if_exists='append', index=False)
                                conn = self.engine.raw_connection()
                                cur = conn.cursor()
                                output = io.StringIO()
                                df.to_csv(output, sep='\t', header=False, index=False)
                                output.seek(0)
                                cur.copy_from(output, table_name, null="")
                                conn.commit()
                                self.update_status(filename=each_file, status=success_status)
                            except Exception as file_read_error:
                                error = f"Failed to read the file"
                                print(f"Failed to read the file : {file_read_error}")
                                self.update_status(filename=each_file, status=failed_status, error_message=error)
                                pass
        except Exception as e:
            print(f'Failed to ingest Data:{str(e)}')

class PushPostgres(Component):
    def __init__(self):
        logger.info("starting component")
        super().__init__()

    def run(self):
        push_data_postgres = PushDataToPostgres()
        push_data_postgres.push_postgres()


if __name__ == "__main__":
   object = PushPostgres()
   object.start()
