from __future__ import annotations

import logging
import re
from pathlib import Path
from typing import Any
from typing import cast
from typing import Literal

import boto3
import numpy as np  # noqa: F401
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.dataset as ds
import pyarrow.fs as fs
from dateutil.parser import ParserError

from .dataset import ArrowColumnMapping
from .dataset import ArrowSourceData
from .dataset import check_null_column_names
from .fs_options import ArrowConvertOptions
from .fs_options import ArrowFileFormat
from .fs_options import ArrowParseOptions
from .fs_options import ArrowPartitioning
from .fs_options import ArrowReadOptions
from .fs_options import ConvertOptions
from .fs_options import FileFormat
from .fs_options import FSType
from .fs_options import ParseOptions
from .fs_options import Partitioning
from .fs_options import ReadOptions
from .misc import to_list


log = logging.getLogger(__name__)


def load_fs_dataset(
    source: str | Path | list[str | Path],
    file_format: FileFormat | None,
    file_schema: pa.Schema | None,
    read_options: ReadOptions | None,
    parse_options: ParseOptions | None,
    convert_options: ConvertOptions | None,
    schema_inf_depth: float | None,  # in MiB
    partitioning: Literal['hive'] | Partitioning | None,
    fs_config: dict[str, Any] | None,
    columns: list[str] | None,
    time_column: str | None,
) -> ArrowSourceData:
    if isinstance(source, list):
        path = Path(source[0])
    else:
        path = Path(source)

    if file_format is None:
        # Infer the file format from the file extension
        if path.is_dir():
            raise ValueError('Must pass "format" arg for directory input')
        file_ext = path.suffix.lower()
        if file_ext == '.csv':
            file_format = 'csv'
        elif file_ext == '.json':
            file_format = 'json'
        elif file_ext in ('.parquet', '.pq'):
            file_format = 'parquet'
        elif file_ext in ('.feather', '.fea'):
            file_format = 'feather'
        elif file_ext == '.orc':
            file_format = 'orc'
        else:
            raise ValueError(
                f'Unknown file extension: "{file_ext}". '
                f'Please pass the "file_format" argument.',
            )

    # Create the Pyarrow read/parse/convert options
    pa_read_options: ArrowReadOptions | None = None
    if read_options:
        pa_read_options = read_options.as_pyarrow

    pa_parse_options: ArrowParseOptions | None = None
    if parse_options:
        pa_parse_options = parse_options.as_pyarrow

    pa_convert_options: ArrowConvertOptions | None = None
    if convert_options:
        pa_convert_options = convert_options.as_pyarrow

    # Configure the scheme inference logic
    if file_schema is None:
        # Set the block size in ReadOptions to control the schema inference depth:
        # https://github.com/apache/arrow/issues/34010
        if schema_inf_depth:
            if schema_inf_depth >= 100:
                # Warn if schema inf depth > 100 since the user may have mistakenly
                # specified it in bytes. A large block size could impact performance
                log.warning(f'schema_inf_depth {schema_inf_depth} > 100 MB')
            block_size_MB = schema_inf_depth
        else:
            # Default to 8 MB to minimize likelihood of schema inference errors
            # NB: Arrow default block size is 1 MB, see github link above
            block_size_MB = 8.0

        block_size_bytes = block_size_MB * 1024 * 1024
        if file_format == 'csv':
            if not pa_read_options:
                pa_read_options = pa.csv.ReadOptions()
            assert isinstance(pa_read_options, pa.csv.ReadOptions)
            pa_read_options.block_size = block_size_bytes
        elif file_format == 'json':
            if not pa_read_options:
                pa_read_options = pa.json.ReadOptions()
            assert isinstance(pa_read_options, pa.json.ReadOptions)
            pa_read_options.block_size = block_size_bytes
        else:
            if schema_inf_depth:
                raise ValueError('"schema_inf_depth" only valid for csv/json files')
    else:
        if schema_inf_depth:
            log.warning('Ignoring "schema_inf_depth" as "file_schema" argument passed')

    # Create the complete PyArrow file format with the passed-in options
    pa_file_format: ArrowFileFormat
    if file_format == 'parquet':
        pa_file_format = ds.ParquetFileFormat(read_options=pa_read_options)

        # Warn about options that don't apply to parquet
        if parse_options:
            log.warning('Ignoring "parse_options" argument - not valid for parquet')
        if convert_options:
            log.warning('Ignoring "convert_options" argument - not valid for parquet')
    elif file_format == 'csv':
        pa_file_format = ds.CsvFileFormat(
            read_options=pa_read_options,
            parse_options=pa_parse_options,
            convert_options=pa_convert_options,
        )
    elif file_format == 'json':
        pa_file_format = ds.JsonFileFormat(
            read_options=pa_read_options,
            parse_options=pa_parse_options,
        )

        # Warn about options that don't apply to JSON
        if convert_options:
            log.warning('Ignoring "convert_options" argument - not valid for JSON')
    else:
        raise ValueError(f'Invalid file format: "{file_format}"')

    # Create the Pyarrow dataset object
    pa_filesystem = create_pyarrow_filesystem(source=source, fs_config=fs_config)
    pa_partitioning: ArrowPartitioning | None = None
    if partitioning:
        if isinstance(partitioning, str):
            if partitioning == 'hive':
                pa_partitioning = ds.partitioning(flavor='hive')
            else:
                raise ValueError(f'Invalid value for partitioning: "{partitioning}"')
        else:
            pa_partitioning = partitioning.as_pyarrow

    # Strip any URI prefixes from the source paths (e.g. "s3://")
    processed_source: str | list[str]
    if isinstance(source, list):
        processed_source = [remove_uri_prefix(str(path)) for path in source]
    else:
        processed_source = remove_uri_prefix(str(source))

    ds_args = {
        'source': processed_source,
        'schema': file_schema,
        'format': pa_file_format,
        'filesystem': pa_filesystem,
        'partitioning': pa_partitioning,
    }
    fs_dataset = ds.dataset(**ds_args)  # type: ignore[arg-type]
    dataset_columns = fs_dataset.schema.names

    source_data: ArrowSourceData
    if columns:
        col_names = columns if isinstance(columns, dict) else to_list(columns)
        for col_name in col_names:
            if col_name not in dataset_columns:
                raise KeyError(
                    f'Column name "{col_name}" is not present in the dataset schema',
                )

        if isinstance(columns, dict):
            # Create the PyArrow mapping dict (str -> Expression)
            col_mapping: ArrowColumnMapping = {}
            for col_name, new_col_name in columns.items():
                col_mapping[new_col_name] = pc.field(col_name)
            source_data = fs_dataset.scanner(columns=col_mapping)
        else:
            source_data = fs_dataset.scanner(columns=to_list(columns))
    else:
        # Check if any of the column names are missing (null) and rename if so
        null_col_mapping = check_null_column_names(columns=dataset_columns)
        if null_col_mapping:
            source_data = fs_dataset.scanner(columns=null_col_mapping)
        else:
            source_data = fs_dataset

    # If the time column type is string try to convert the values to true timestamps
    if time_column:
        if isinstance(source_data, ds.Scanner):
            data_schema = source_data.projected_schema
        else:
            data_schema = source_data.schema

        if data_schema.field(time_column).type == pa.string():
            log.info(
                f'Time column "{time_column}" is of type string: '
                f'attempting to convert to type timestamp',
            )
            source_data = convert_string_timestamps(
                source_data=source_data,
                schema=data_schema,
                time_column=time_column,
            )

    return source_data


def create_pyarrow_filesystem(
    source: str | Path | list[str | Path],
    fs_config: dict[str, Any] | None = None,
) -> fs.FileSystem:
    fs_config = {} if fs_config is None else fs_config
    filesystem: fs.FileSystem

    paths = [str(p) for p in source] if isinstance(source, list) else [str(source)]
    fs_type = get_fs_type(paths[0])

    # If we have multiple source paths check that the fs_type is consistent across all
    if len(paths) > 1:
        if any(get_fs_type(path) != fs_type for path in paths[1:]):
            raise ValueError('Source paths list contains multiple file system types...')

    # Create the Arrow filesystem instance
    if fs_type == 's3':
        # Get the AWS credentials using boto3 and pass these directly to S3FileSystem.
        # This enables the user to specify the AWS credentials profile_name (which isn't
        # currently possible using pyarrow's S3FileSystem) and also gets round a bug
        # where S3FileSystem fails to read the credentials file when called with no args
        session = boto3.session.Session(**fs_config)
        credentials = session.get_credentials()
        assert credentials
        filesystem = fs.S3FileSystem(
            access_key=credentials.access_key,
            secret_key=credentials.secret_key,
            region=session.region_name,
        )
    elif fs_type == 'gcs':
        filesystem = fs.GcsFileSystem(**fs_config)
    elif fs_type == 'hdfs':
        filesystem = fs.HadoopFileSystem(**fs_config)
    else:
        filesystem = fs.LocalFileSystem(**fs_config)

    return filesystem


def get_fs_type(path: str) -> FSType:
    m = re.match(r'(server|s3|gcs|hdfs):', path, re.IGNORECASE)
    fs_type = m.group(1) if m else 'local'
    return cast(FSType, fs_type)


def remove_uri_prefix(path: str) -> str:
    # Remove the URI prefix from the source paths (req'd S3 format: bucket/key)
    m = re.match(r'server:|s3:/+|gcs:/+|hdfs:/+', path, re.IGNORECASE)
    if m:
        return path[len(m.group(0)):]
    else:
        return path


def convert_string_timestamps(
    source_data: ArrowSourceData, schema: pa.Schema, time_column: str,
) -> ds.InMemoryDataset:
    # The main use-case for this is parquet files where the timestamps have
    # been stored as strings. For csv/json files Pyarrow will automatically
    # attempt to infer and convert timestamps to ISO-8601, however for non-standard
    # formats this may not be possible. In this case, the user can pass a ConvertOptions
    # object specifying the timestamp format using the "timestamp_parsers" attribute:
    # https://arrow.apache.org/docs/python/generated/pyarrow.csv.ConvertOptions.html#pyarrow.csv.ConvertOptions.timestamp_parsers
    if isinstance(source_data, pa.Table):
        table = source_data
    else:
        log.info('Materializing data to convert time values from strings to timestamps')
        table = source_data.to_table()
    try:
        pd_dts = pd.to_datetime(
            table[time_column].to_pylist(), utc=False, errors='coerce',
        )
        time_column_idx = schema.get_field_index(time_column)
        new_table = table.drop([time_column])
        table = new_table.add_column(time_column_idx, time_column, pa.array(pd_dts))
        return ds.InMemoryDataset(table)
    except ParserError as e:
        raise ParserError(
            f'Problem converting time strings from column "{time_column}" to true '
            f'timestamps\nFor CSV data you can specify the timestamp format via '
            f'the "convert_options" argument (see: CSVConvertOptions)\n'
            f'{e}',
        )
