#%%
#import env
import pandas as pd
# from sqlalchemy import create_engine
# from snowflake.sqlalchemy import URL
# from snowflake.snowpark import Session
from snowflake.sqlalchemy import URL
import snowflake.connector
from snowflake.connector.pandas_tools import write_pandas
from datetime import datetime
import time
import math
from sqlalchemy import create_engine
#import utility_functions as uf
#%%
class Snowflake():
    
    def __init__(self,connection_params_dict,client_name) -> None:
        """Initialise the snowflake connection with a dictionary of connection parameters.

        Log File is created in the LOG_DIR environment variable with a subdirectory with the name of the client_name
        
        Args:
            connection_params_dict (dict): Dictionary of connection parameters
            client_name (str): Name of the client, this is used to create a folder in the LOG_DIR environment variable
        
        Returns:
            None"""
        
        self.connection_params_dict = connection_params_dict

        self.user = self.connection_params_dict['user']
        self.role = self.connection_params_dict['role']
        self.warehouse = self.connection_params_dict['warehouse']
        self.database = self.connection_params_dict['database']
        self.schema = self.connection_params_dict['schema']
        self.account = self.connection_params_dict['account']
        self.password = self.connection_params_dict['password']
        #self.session = Session().create(self.connection_params_dict)
        # self.session = Session.builder.configs(self.connection_params_dict).create()
        #self.logger = Logger(client_name,'snowflake')

    
    def read_snowflake_to_df(self, table_name, schema=None, chunk_size=200000):
        """Function to read Snowflake table using SQLAlchemy
        
        URL stands for Uniform Resource Locator. It is a reference (an address) to a resource on the Internet.
        
        Args:
            table_name (str): Name of the table to read
            schema (str, optional): Name of the schema to read from. Defaults to the schema specified on class initialisation.
            chunk_size (int, optional): Size of the data chunks to read in a single batch. Defaults to 200000.
        
        Returns:
            df: Pandas dataframe of the data read from Snowflake"""

        if schema == None:
            schema = self.schema

        with snowflake.connector.connect(
            user=self.user,
            password=self.password,
            account=self.account,
            warehouse=self.warehouse,
            database=self.database,
            schema=schema,
        ) as conn:
            with conn.cursor() as cur:
                cur.execute(f"SELECT COUNT(*) FROM {table_name}")
                num_rows = cur.fetchone()[0]
                print(f"Total rows in {table_name}: {num_rows}")

                num_chunks = math.ceil(num_rows / chunk_size)
                print(f"Fetching {num_rows} rows in {num_chunks} chunks of {chunk_size} rows each")
            
            with conn.cursor() as cur:
                cur.execute(f"SELECT * FROM {table_name}")
                
                # Initialize an empty DataFrame to hold the results
                df_total = pd.DataFrame()

                for chunk_index in range(num_chunks):
                    rows = cur.fetchmany(chunk_size)
                    if not rows:
                        break

                    df_chunk = pd.DataFrame(rows, columns=[x[0] for x in cur.description])
                    df_total = pd.concat([df_total, df_chunk], ignore_index=True)

                    # For now, we'll just print the size of each chunk
                    print(f"Fetched chunk {chunk_index + 1} of {num_chunks}, size: {len(df_chunk)} rows")

        return df_total

    
    
    def drop_table(self,table_name,database=None, schema=None):
        '''Function to drop Snowflake table
        
        Args:
            table_name (str): Name of the table to drop
            database (str, optional): Name of the database to drop the table from. Defaults to the database specified on class initialisation.
            schema (str, optional): Name of the schema to drop the table from. Defaults to the schema specified on class initialisation.
        
        Returns:
            None'''
        pass
    
    def write_df_to_snowflake(self, df, table_name, database=None, schema=None,
                          auto_create_table=False, overwrite=False, chunk_size=20000):
        '''Function to write Pandas dataframe to Snowflake table

        Truncates (if it exists) or creates new table and inserts the new data into the selcted table
        
        Args:
            df (dataframe): Pandas dataframe to write to Snowflake
            table_name (str): Name of the table to write to
            database (str, optional): Name of the database to write the table to. Defaults to the database specified on class initialisation.
            schema (str, optional): Name of the schema to write the table to. Defaults to the schema specified on class initialisation.
            auto_create_table (bool, optional): If True, creates the table if it does not exist. Defaults to False.
            overwrite (bool, optional): If True, overwrites the table if it exists. Defaults to False.
            chunk_size (int, optional): Size of the data chunks to write in a single batch. Defaults to 200000.
        
        Returns:
            None'''
        if schema == None:
            schema = self.schema
        if database == None:
            database = self.database
        
        with snowflake.connector.connect(
            user=self.user,
            password=self.password,
            account=self.account,
            warehouse=self.warehouse,
            database=self.database,
            schema=schema,
        ) as conn:
            print(f"Connected to Snowflake")
            try:
                now = time.time()
                success, nchunks, nrows, _ = write_pandas(conn, df, table_name, parallel=8,schema=schema,database=database,
                                        auto_create_table=auto_create_table,overwrite=overwrite,chuck_size=chunk_size)
                time_taken = round(time.time() - now,2)
                print(f"Sent Data to {table_name}, time taken: {time_taken} seconds")
                
            except Exception as error_message:
                print(f"Connection error {error_message}")
                time.sleep(10)  # wait for 10 seconds then try again
                try:
                    now = time.time()
                    success, nchunks, nrows, _ = write_pandas(conn,df, table_name, parallel=8,schema=schema,database=database,
                                            auto_create_table=auto_create_table,overwrite=overwrite,chuck_size=chunk_size)
                    time_taken = round(time.time() - now,2)
                    print(f"Sent Data to {table_name} time taken: {time_taken} seconds")
                    
                except Exception as error_message:
                    print("Connection failed again")
                    print(f'{table_name} error: ' + str(error_message))
                    return f'{table_name} error: ' + str(error_message) 
# %%
