# -*- coding: utf-8 -*-


import pandas as pd 
import pyarrow as pa
from pyarrow.fs import S3FileSystem
import pyarrow.dataset as ds
import io
import boto3
import json




def get_df_s3(s3_path_read, client, file_format="parquet", sheet_name=None, json_key=None):
    """
    Docstring for get_df_s3
    
    :param s3_path_read: The S3 path without including 's3://'
    :param client: aws credentials
    :param file_format: file format we want to read
    :param sheet_name: In case of reading an Excel file, specify the sheet name
    :param json_-key: In case of reading a json file and need a dataframe from a specific key
    """
    s3 = boto3.client(
        "s3",
        aws_access_key_id=client["key"],
        aws_secret_access_key=client["secret"],
        region_name="eu-west-2"
    )

    if file_format in (["tgz", "rar"]):
        # Extraer bucket y key del path estilo s3://bucket/key
        import rarfile
        
        bucket, key = s3_path_read.split("/", 1)

        obj = s3.get_object(Bucket=bucket, Key=key)
        rar_bytes = obj['Body'].read()

        # Trata el archivo como .rar aunque tenga extensión .tgz
        with rarfile.RarFile(io.BytesIO(rar_bytes)) as rf:
            csv_files = [f for f in rf.namelist() if f.endswith('.csv')]
            if not csv_files:
                raise ValueError("No se encontró ningún CSV dentro del archivo .rar")
            
            with rf.open(csv_files[0]) as csv_file:
                df = pd.read_csv(csv_file)

        return df
    
    elif file_format in (['xls', 'xlsx']):
        
        storage_options = {
            "key": client["key"],
            "secret": client["secret"],
            "client_kwargs": {
                "region_name": "eu-west-2"
            }
        }
        
        path = 's3://' + s3_path_read
        df = pd.read_excel(path, storage_options=storage_options, sheet_name=sheet_name)
        
        return df
    
    elif file_format == 'json':

        bucket, key = s3_path_read.split("/", 1)
        obj = s3.get_object(Bucket=bucket, Key=key) 
        data = json.loads(obj["Body"].read().decode("utf-8"))

        if json_key:
            return pd.DataFrame(data[json_key])
        else:
            return pd.DataFrame([data])
        
    
    else:
      
        s3_fs = S3FileSystem(
            access_key=client["key"],
            secret_key=client["secret"],
            region="eu-west-2"
        )
        dataset = ds.dataset(s3_path_read, format=file_format, filesystem=s3_fs)
        
        # Convertir a tabla y forzar tipo de 'date' a string
        table = dataset.to_table()
        
        original_schema = dataset.schema
        fields = []
        for field in original_schema:
            if field.name == "date":
                fields.append(pa.field("date", pa.string()))
            else:
                fields.append(field)
        new_schema = pa.schema(fields)
        
        # 5. Castear la tabla y convertir a pandas
        casted_table = table.cast(new_schema)
        df = casted_table.to_pandas()
        
        df = table.to_pandas()
        
    return df

