"""This is the tool used to read/write parquet/csv files from/to s3/local.
Before using this tool, make sure you have .aws/credentials and /.aws/config file setup correctly.
"""


import os
import re

import pandas as pd
import s3fs
from pyarrow import parquet as pq, Table
import uuid

from kasheesh_tools.tool_logger import get_my_logger

def _is_s3_arn_path(s3_path):
    """Check if path is an S3 Access Point ARN format.
    
    Args:
        s3_path (str): S3 path to check
        
    Returns:
        bool: True if path is an ARN format
    """
    return bool(re.match(r's3://arn:aws:s3:', s3_path))

def _get_path_for_s3fs(s3_path, for_parquet=False):
    """Get the path format that s3fs expects.
    
    For ARN paths, we need to map the access point to the actual bucket name.
    The access point 'kasheesh-data-staging-accesspoint' maps to bucket 'kasheesh-data-staging'.
    
    Args:
        s3_path (str): S3 path, possibly with ARN format
        for_parquet (bool): If True, extract bucket name for ParquetDataset
        
    Returns:
        str: Path format for s3fs (without s3:// prefix)
    """
    if _is_s3_arn_path(s3_path):
        # Pattern: s3://arn:aws:s3:region:account-id:accesspoint/accesspoint-name/path
        arn_pattern = r's3://arn:aws:s3:[^:]+:\d+:accesspoint/([^/]+)(.*)'
        match = re.match(arn_pattern, s3_path)
        if match:
            accesspoint_name = match.group(1)
            path_suffix = match.group(2) if match.group(2) else ''
            
            # Map access point name to bucket name
            # kasheesh-data-staging-accesspoint -> kasheesh-data-staging
            if accesspoint_name.endswith('-accesspoint'):
                bucket_name = accesspoint_name[:-len('-accesspoint')]
            else:
                # Fallback: try to extract bucket name from access point name
                # Common pattern: bucket-name-accesspoint -> bucket-name
                bucket_name = accesspoint_name.replace('-accesspoint', '')
            
            # Return bucket name with path suffix (without s3:// prefix)
            return f'{bucket_name}{path_suffix}'
        # If ARN pattern doesn't match, try to remove s3:// prefix
        return s3_path.replace("s3://", "", 1)
    else:
        # For regular paths, remove s3:// prefix
        return s3_path.replace("s3://", "")

class FileHandler:
    def __init__(self, name="FileHdlr") -> None:
        self.logger = get_my_logger(name=name, level="INFO")

    def write_parquet(self, df, root_path, partition_cols=None, use_threads=6, overwrite=False,
                      date_of_job: str = "date", job_name: str = "job", keep_unique=True):
        """write parquet file from s3 or local

        Args:
            df (pd.DataFrame): dataframe to write
            root_path (str): root path to write
            partition_cols (list, optional): columns to partition, ORDERED AS SPECIFIED. Defaults to None.
            use_threads (int/bool, optional): threads count. Default cpu_count(). Forced Default to 6.
            overwrite (bool, optional): overwrite existing file. Defaults to False.
            date_of_job (str, optional): date of job (YYYYMMDD). Defaults to "date".
            job_name (str, optional): job name. Defaults to "job".
        """
        if 'datadate' not in df.columns:
            df['datadate'] = pd.Timestamp.now().date()
            self.logger.info(f"Added datadate column to dataframe.")
        self.logger.info(f"Writing {df.shape} dataframe to {root_path}")
        overwrite_behavior = None  # default: removing only the files with the same name.
        if overwrite:
            overwrite_behavior = 'delete_matching'  # removing the whole partition, all files in it.
        if keep_unique:
            uuid_str = uuid.uuid4().hex
        else:
            uuid_str = ""
        if "s3://" in root_path:
            # Get path format for s3fs (handles ARN paths correctly)
            s3fs_path = _get_path_for_s3fs(root_path)
            fs = s3fs.S3FileSystem()
            pq.write_to_dataset(
                Table.from_pandas(df),
                s3fs_path,
                partition_cols=partition_cols,
                filesystem=fs,
                use_dictionary=True,
                compression="snappy",
                existing_data_behavior=overwrite_behavior,
                basename_template = "{job_name}-{date_of_job}-{uuid}-{{i}}.parquet".format(date_of_job=date_of_job, 
                                                                                           job_name=job_name,
                                                                                           uuid=uuid_str)
            )
        else:
            # df.to_parquet(root_path, partition_cols=partition_cols, use_threads=use_threads)
            pq.write_to_dataset(
                Table.from_pandas(df),
                root_path,
                partition_cols=partition_cols,
                use_dictionary=True,
                compression="snappy",
                existing_data_behavior=overwrite_behavior,
                basename_template = "{job_name}-{date_of_job}-{{i}}-{uuid}.parquet".format(date_of_job=date_of_job, 
                                                                                           job_name=job_name,
                                                                                           uuid=uuid_str)
            )
            
        self.logger.info(f"Writing finished.")
        return None
    

    def read_parquet(self, root_path, filters=None, use_threads=6, use_env_creds=False):
        """read parquet file from s3 or local

        Args:
            root_path (str): root path to read from
            filters (list, optional): Specify AND/OR conditions here. Defaults to None.
                    #readout_partial = FH.read_parquet(root_path=root_path, 
                        filters=[[("a", "=", 1)], [("a", "=", 3)]])  # This is OR condition
                    # readout_partial = FH.read_parquet(root_path=root_path, 
                        filters=[[("a", "=", 1), ("a", "=", 3)]])  # This is AND condition
            use_threads (int/bool, optional): threads count, Default cpu_count(). Forced Default to 6.
            use_env_creds (bool, optional): use AWS credentials from environment variables. Defaults to False.

        Returns:
            df: pd.DataFrame
        """
        self.logger.info(f"Reading from {root_path}")
        if "s3://" in root_path:
            # Get path format for ParquetDataset (extracts access point name for ARN paths)
            s3fs_path = _get_path_for_s3fs(root_path, for_parquet=True)
            if use_env_creds and "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ:
                fs = s3fs.S3FileSystem(key=os.environ['AWS_ACCESS_KEY_ID'], 
                                       secret=os.environ['AWS_SECRET_ACCESS_KEY'])
            else:
                fs = s3fs.S3FileSystem()
            fs.invalidate_cache()
            df = pq.ParquetDataset(s3fs_path, filters=filters, filesystem=fs).read_pandas(use_threads=use_threads).to_pandas()
        else:
            df = pd.read_parquet(root_path, 
                                 filters=filters, use_threads=use_threads)
        self.logger.info(f"Reading finished. Data shape {df.shape}")
        return df
    
    def read_csv(self, file_path, use_env_creds=False):
        """read csv file from s3 or local

        Args:
            file_path (str): file path to read csv from

        Returns:
            df: pd.DataFrame
        """
        self.logger.info(f"Reading from {file_path}")
        if "s3://" in file_path:
            # Get path format for s3fs (handles ARN paths correctly)
            s3fs_path = _get_path_for_s3fs(file_path)
            if use_env_creds and "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ:
                fs = s3fs.S3FileSystem(key=os.environ['AWS_ACCESS_KEY_ID'], 
                                       secret=os.environ['AWS_SECRET_ACCESS_KEY'])
            else:
                fs = s3fs.S3FileSystem()
            df = pd.read_csv(fs.open(s3fs_path))
        else:
            df = pd.read_csv(file_path)
        self.logger.info(f"Reading finished. Data shape {df.shape}")
        return df
    
    def write_csv(self, df, file_path, use_env_creds=False, mode='w', index=False):
        """write csv file to s3 or local

        Args:
            df (pd.DataFrame): dataframe to write
            file_path (str): file path to write csv to
        """
        self.logger.info(f"Writing {df.shape} dataframe to {file_path}")
        if "s3://" in file_path:
            # Get path format for s3fs (handles ARN paths correctly)
            s3fs_path = _get_path_for_s3fs(file_path)
            if use_env_creds and "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ:
                fs = s3fs.S3FileSystem(key=os.environ['AWS_ACCESS_KEY_ID'], 
                                       secret=os.environ['AWS_SECRET_ACCESS_KEY'])
            else:
                fs = s3fs.S3FileSystem()
            with fs.open(s3fs_path, mode=mode) as f:
                df.to_csv(f, index=index)
        else:
            df.to_csv(file_path, mode=mode, index=index)
        self.logger.info(f"Writing finished.")
        return None

    def list_files(self, root_path):
        """list files in s3 or local

        Args:
            root_path (str): root path to list files from

        Returns:
            list: list of files
        """
        self.logger.info(f"Listing files in {root_path}")
        if "s3://" in root_path:
            # Get path format for s3fs (handles ARN paths correctly)
            s3fs_path = _get_path_for_s3fs(root_path)
            fs = s3fs.S3FileSystem()
            file_ls = fs.ls(s3fs_path)
        else:
            file_ls = os.listdir(root_path)
        self.logger.info(f"Listing finished. Total files {len(file_ls)}")
        return file_ls


