from __future__ import annotations

import base64
import logging
from typing import Any
from typing import Dict
from typing import Literal
from typing import Type
from typing import Union

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.dataset as ds
import pyarrow.flight as flight
from typing_extensions import TypeAlias

from .exceptions import parse_arrow_exception
from .server import SessionParams


log = logging.getLogger(__name__)

ArrowSourceData = Union[ds.Dataset, pa.Table, ds.Scanner]
ArrowSourceDataType = Union[Type[ds.Dataset], Type[pa.Table], Type[ds.Scanner]]
ArrowColumnMapping = Dict[str, pc.Expression]
Serializable: TypeAlias = Any

AsofDirection = Literal['forward', 'backward']
"""ASOF JOIN direction - see [FSDataset.join()][tempora.datasets.FSDataset.join]."""


def arrow_json_encoder(obj: Any) -> Serializable:
    """Custom JSON encoder for Arrow Schema objects.

    Args:
        obj: the object to be serialized to JSON.

    Returns:
        A dictionary containing the base64 encoded Arrow IPC message bytes if the object
        is a Pyarrow schema, otherwise the string representation of the object.
    """
    if isinstance(obj, pa.Schema):
        schema_bytes = obj.serialize().to_pybytes()
        # Convert Arrow IPC msg bytes to base64 bytes, then decode to ascii for JSON
        # See: https://stackoverflow.com/a/40000564
        b64_ascii = base64.b64encode(schema_bytes).decode('ascii')
        return {'class': 'pa.Schema', 'b64_ascii': b64_ascii}
    else:
        # If it's not a Schema then default to returning the string representation
        return str(obj)


def arrow_json_decoder(dct: dict[str, Any]) -> pa.Schema | dict[str, Any]:
    """Custom JSON decoder for Arrow Schema objects.

    Args:
        dct: the Python dictionary to be deserialized.

    Returns:
        A Pyarrow schema object if the dictionary conforms to the output of
        arrow_json_encoder(schema), otherwise the original dictionary object.
    """
    if 'class' in dct and 'b64_ascii' in dct and dct['class'] == 'pa.Schema':
        schema_bytes = base64.b64decode(dct['b64_ascii'])
        return pa.ipc.read_schema(pa.py_buffer(schema_bytes))
    return dct


def cleanup(json_key: str, session_params: SessionParams) -> None:
    """Cleanup function for Dataset objects."""
    client = session_params['client']
    assert client
    auth_token = session_params['auth_token']
    action = flight.Action('cleanup', json_key.encode())
    try:
        call_options = flight.FlightCallOptions(headers=[auth_token])
        next(client.do_action(action, call_options))
    except flight.FlightUnavailableError:
        log.warning('Unable to connect to the server for cleanup...')
    except flight.FlightError as e:
        msg = parse_arrow_exception(e)
        log.error(msg)
    except StopIteration:
        # Empty response from server - all good
        ...
    return None


def check_null_column_names(columns: list[str]) -> ArrowColumnMapping | None:
    """Replace empty column names in the schema that would otherwise error.

    Replace empty column names in the schema, that would otherwise throw a parser error
    when the Arrow data is queried with DuckDB. See:
    https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
    'Quoted identifiers can contain any character except the character with code zero'

    Args:
        columns: list of column names to check

    Returns:
        A dictionary that can be passed to the Arrow dataset scanner() method to rename
        any empty columns.
    """
    for column in columns:
        if column == '':
            break
    else:
        # No null columns found
        return None

    n, replacement_name_prefix = 1, '__missing_column_name'
    col_mapping: ArrowColumnMapping = {}

    for column in columns:
        if column == '':
            new_name = f'{replacement_name_prefix}_{n}'
            col_mapping[new_name] = pc.field(column)
        else:
            col_mapping[column] = pc.field(column)

    return col_mapping
