# -*- coding: utf-8 -*-
"""
Created on Wed Jul 16 13:23:51 2025

@author: adria
"""

import boto3
import time
import pandas as pd
import streamlit as st

output_location = 's3://test-bucket-767397873198-eu-west-2/tmp_adria/test_python_conversion/'

def create_table(
        client,
        database,
        table_name,
        cols_typed,
        file_type,
        location
        ):
    
    try:
        query_create_table = f"""
        CREATE EXTERNAL TABLE {table_name} (
          {cols_typed}
        )
        STORED AS {file_type}
        LOCATION '{location}';
        """
        
        response = client.start_query_execution(
            QueryString=query_create_table,
            QueryExecutionContext={'Database': database},
            ResultConfiguration={'OutputLocation': output_location}
        )
    
    except Exception as e:
        print("Error executing the creation of the table")
        print(e)
        

def drop_table(
        client,
        database,
        table_name,
        ):

    try:
        query_drop_table = f'drop table if exists {table_name}'
        
        response = client.start_query_execution(
            QueryString=query_drop_table,
            QueryExecutionContext={'Database': database},
            ResultConfiguration={'OutputLocation': output_location}
        )
        
        
    except Exception as e:
        print(f"Error droping the table {table_name}")
        print(e)
        

def get_df(
        client,
        query,
        database
        ):
    
    try: 
        
        response = client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={'OutputLocation': output_location}
        )

        execution_id = response['QueryExecutionId']
        attempt = 0
        max_attempts = 5

        while attempt < max_attempts:
            while True:
                status = client.get_query_execution(QueryExecutionId=execution_id)['QueryExecution']['Status']['State']
                if status in ['SUCCEEDED', 'CANCELLED', 'FAILED']:
                    break
                time.sleep(2)

            if status in ['SUCCEEDED', 'CANCELLED']:
                break
            elif status == 'FAILED':
                    attempt += 1
                    if attempt <= max_attempts:
                        execution_id = retry_query(client, query, database)
                    else:
                        break



        result = []
        next_token = None
        
        while True:
            if next_token:
                response = client.get_query_results(
                    QueryExecutionId=execution_id,
                    NextToken=next_token
                )
            else:
                response = client.get_query_results(QueryExecutionId=execution_id)
    
            result.extend(response['ResultSet']['Rows'])
    
            next_token = response.get('NextToken')
            if not next_token:
                break

    
    except Exception as e:
        print("Error executing the query:")
        print(query)
        print(e)  
        pass

    rows = result

    try:
        columns = [col['VarCharValue'] for col in rows[0]['Data']]
    except:
        pass
    data = [[col.get('VarCharValue', '') for col in row['Data']] for row in rows[1:]]

    return pd.DataFrame(data, columns=columns)
    

def retry_query(client, query, database, output_location=output_location):
    response = client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={'Database': database},
        ResultConfiguration={'OutputLocation': output_location}
    )
    return response['QueryExecutionId']