from __future__ import annotations

import dataclasses as dc
import json
import logging
import signal
import threading
from collections import deque
from pathlib import Path
from types import TracebackType
from typing import Annotated
from typing import Any
from typing import Deque
from typing import Iterator
from typing import NoReturn
from weakref import finalize

import numpy as np
import pyarrow as pa
import pyarrow.flight as flight
import pyarrow.fs as fs
import pydantic.dataclasses as pydc
from pydantic import Field

from .target_spec import TargetSpec
from .transform_spec import EventEncoder
from .transform_spec import LabelEncoder
from .transform_spec import TransformSpec
from tempora.datasets.protocols import IsDataset
from tempora.utils.batch import Batch
from tempora.utils.batch import TSSegment
from tempora.utils.dataclass import get_args_dict
from tempora.utils.dataclass import pydc_config
from tempora.utils.dataset import arrow_json_encoder
from tempora.utils.dataset import cleanup
from tempora.utils.exceptions import parse_arrow_exception
from tempora.utils.exceptions import TemporaClientError
from tempora.utils.exceptions import TemporaServerError
from tempora.utils.output import OutputFormat
from tempora.utils.sampling import SamplerOptions
from tempora.utils.sampling import serialize_callable
from tempora.utils.server import get_connection
from tempora.utils.server import SessionParams
from tempora.utils.time import Length


log = logging.getLogger(__name__)


ClassWeight = Annotated[float, Field(gt=0, le=1)]


@pydc.dataclass(frozen=True, config=pydc_config)
class ClassSamplingSpec:
    """Class-level sampling specification for classification targets.

    Args:
        name: Class identifier used for logging/debugging.
        expr: SQL-compatible expression that matches dataset rows containing the
            desired target/label.
        weight: Normalized class sampling weight in the interval `(0, 1]`.
    """
    name: str
    expr: str
    weight: ClassWeight

    def __post_init__(self) -> None:
        if not self.name.strip():
            raise ValueError('"name" cannot be empty')

        if not self.expr.strip():
            raise ValueError('"expr" cannot be empty')


class SegmentReader(threading.Thread):
    def __init__(
        self, segments_buffer: Deque,
        json_sampler: str,
        targets_schema: pa.Schema | None,
        client: flight.FlightClient,
        auth_token: tuple[str, str],
    ) -> None:
        super().__init__(name='TemporaSegmentReader', daemon=True)
        self.segments_buffer = segments_buffer
        self.json_sampler = json_sampler
        self.targets_schema = targets_schema
        self.client = client
        self.auth_token = auth_token

        self.awake = threading.Event()
        self.terminate = threading.Event()

    def run(self) -> None:
        flight_reader: flight.FlightStreamReader | None = None

        # Set the awake event to True so that we always read segments on the first try
        self.awake.set()
        assert self.segments_buffer.maxlen is not None

        while not self.terminate.is_set():
            # Do nothing until the awake event is set by the _sample_batches method.
            # NB we wrap the call in a while loop in case wait timeouts before it's set.
            while not self.awake.is_set():
                self.awake.wait()

            if self.terminate.is_set():
                break

            if not flight_reader:
                try:
                    json_key = self.json_sampler.encode()
                    desc = flight.FlightDescriptor.for_command(json_key)
                    call_options = flight.FlightCallOptions(headers=[self.auth_token])
                    _, flight_reader = self.client.do_exchange(desc, call_options)
                except flight.FlightError as e:
                    msg = 'Unable to retrieve data from the server: '
                    msg += parse_arrow_exception(e)
                    log.error(msg)
                    raise TemporaServerError(msg) from None

            while len(self.segments_buffer) < self.segments_buffer.maxlen:
                if self.terminate.is_set():
                    break

                try:
                    chunk = flight_reader.read_chunk()
                    segment = TSSegment.from_chunk(chunk=chunk)
                    self.segments_buffer.appendleft(segment)
                except StopIteration:
                    # Received all segments written to the stream
                    flight_reader = None
                    break

            # Put the thread to sleep since we've filled the buffer or we're terminating
            self.awake.clear()

        return None


@dc.dataclass
class SamplerBase:
    """Abstract base class for samplers.

    Args:
        batch_size: Batch size.
        columns: Columns from the dataset to include in each batch. If None all columns
            are included.
        transform_spec: Optional transform specification for context-window data.
        target_spec: Optional target specification.
        output_format: Format of each sampled batch.
        as_tensor: Whether to generate batches in tensor form. If False batches will
            be generated in "packed" form.
        pad_value: Value to use to pad each batch.
        options: Options that influence the behavior of the sampler.
    """
    batch_size: int = 32
    columns: list[str] | None = None
    transform_spec: TransformSpec | None = None
    target_spec: TargetSpec | None = None
    output_format: OutputFormat = 'ndarray'
    as_tensor: bool = False
    pad_value: int | float = np.nan
    options: SamplerOptions = dc.field(default_factory=SamplerOptions)

    def __post_init__(self) -> None:
        self._dataset: IsDataset | None = None
        self._segment_reader: SegmentReader | None = None
        self._interrupt = threading.Event()

        # Create the segments buffer
        if self.options.segments_buffer_len is None:
            deque_maxlen = int(self.batch_size * 5)
        else:
            deque_maxlen = int(self.options.segments_buffer_len)
        self._segments_buffer: Deque = deque(maxlen=deque_maxlen)

        # Set up a signal handler for Ctrl-C interrupts during sampling
        signal.signal(signal.SIGINT, self._sigint_handler)

        self._finalize()

    def _finalize(self) -> None:
        # Set up the finalizer for cleanup. We must pass an actual json sampler string
        # and not a reference to "self._as_json", so that self can be garbage collected
        # https://docs.python.org/3/library/weakref.html#weakref.finalize
        # NB: we do this in a method (vs __post_init__) so we can mock it during testing
        json_sampler = self._as_json
        client, auth_token = get_connection()
        session_params = SessionParams(client=client, auth_token=auth_token)
        finalize(self, cleanup, json_key=json_sampler, session_params=session_params)

    def _init_sampler(self) -> None:
        self._validate_target_spec()

        # Initialize the sampler instance on the server
        action = flight.Action('init_sampler', self._as_json.encode())
        client, auth_token = get_connection()
        call_options = flight.FlightCallOptions(headers=[auth_token])
        result_iter = client.do_action(action, call_options)
        try:
            # Server responds with an empty result iter upon success
            next(result_iter)
        except flight.FlightError as e:
            msg = parse_arrow_exception(e)
            log.error(msg)
            raise TemporaServerError(msg) from None
        except StopIteration:
            # Empty response from server - all good
            ...
        return None

    def _validate_target_spec(self) -> None:
        if self._dataset is None or self.target_spec is None:
            return None

        transform_spec = self.target_spec.transform_spec
        if not isinstance(transform_spec, (LabelEncoder, EventEncoder)):
            return None

        assert self.target_spec.columns
        target_column = self.target_spec.columns[0]
        field_idx = self._dataset.schema.get_field_index(target_column)
        field = self._dataset.schema.field(field_idx)
        dtype = field.type
        if pa.types.is_dictionary(dtype):
            dtype = dtype.value_type

        if isinstance(transform_spec, EventEncoder):
            if not (pa.types.is_date(dtype) or pa.types.is_timestamp(dtype)):
                raise TemporaClientError(
                    'EventEncoder requires the target column to be a date or '
                    f'datetime type (including dictionary-encoded variants), but '
                    f'"{target_column}" has type "{field.type}"',
                )
            return None

        has_class_map = transform_spec.class_map is not None
        is_valid_label_dtype = (
            pa.types.is_boolean(dtype)
            or pa.types.is_integer(dtype)
            or pa.types.is_string(dtype)
            or pa.types.is_large_string(dtype)
            or (
                has_class_map and (
                    pa.types.is_floating(dtype)
                    or pa.types.is_decimal(dtype)
                )
            )
        )
        if not is_valid_label_dtype:
            expected_dtypes = 'bool, integer, or string'
            if not has_class_map:
                expected_dtypes += ' (float/decimal require "class_map")'
            raise TemporaClientError(
                'LabelEncoder requires the target column to be '
                f'{expected_dtypes} (including dictionary-encoded variants), but '
                f'"{target_column}" has type "{field.type}"',
            )

        return None

    @property
    def schema(self) -> pa.Schema | NoReturn:
        """Returns the schema of each batch."""
        return self._get_schema(action_type='get_sampler_data_schema')

    @property
    def targets_schema(self) -> pa.Schema | None | NoReturn:
        """Returns the schema of the batch targets."""
        if self.target_spec:
            return self._get_schema(action_type='get_sampler_targets_schema')
        else:
            return None

    def _get_schema(self, action_type: str) -> pa.Schema | NoReturn:
        if self._dataset is None:
            raise TemporaClientError(
                'Cannot return a schema until the sampler has been called on a dataset',
            )

        action = flight.Action(action_type, self._as_json.encode())
        try:
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            schema_bytes = next(client.do_action(action, call_options))
            return pa.ipc.read_schema(pa.py_buffer(schema_bytes.body))
        except flight.FlightError as e:
            msg = f'Problem encountered getting schema: {parse_arrow_exception(e)}'
            raise TemporaServerError(msg) from None
        except StopIteration:
            raise TemporaServerError('No response from server...') from None

    @staticmethod
    def _wire_converter(field_value: Any, field_type: str) -> Any:
        if field_value and 'TimeDelta | Length' in field_type:
            # Ensure we always send time lengths/deltas as Length objects to the server
            if not isinstance(field_value, Length):
                field_value = Length(field_value)
        if field_value and 'TransformFunction' in field_type:
            if not callable(field_value):
                raise TypeError(
                    f'Expected a callable transform, got {type(field_value).__name__}',
                )
            return serialize_callable(field_value)
        return field_value

    @property
    def _as_json(self) -> str:
        sampler_dict = get_args_dict(self, converter=SamplerBase._wire_converter)
        sampler_dict['id'] = id(self)
        if self._dataset:
            sampler_dict['args']['dataset'] = get_args_dict(self._dataset)
        else:
            sampler_dict['args']['dataset'] = None

        return json.dumps(sampler_dict, default=arrow_json_encoder)

    def _sigint_handler(self, signum: int, frame: Any) -> None:  # noqa
        self._interrupt.set()
        self._terminate_segment_reader()
        return None

    def _start_segment_reader(self) -> None:
        self._terminate_segment_reader()
        client, auth_token = get_connection()
        self._segment_reader = SegmentReader(
            segments_buffer=self._segments_buffer,
            json_sampler=self._as_json,
            targets_schema=self.targets_schema,
            client=client,
            auth_token=auth_token,
        )
        self._segment_reader.start()
        return None

    def _wake_segment_reader(self) -> None:
        if self._segment_reader:
            self._segment_reader.awake.set()
        return None

    def _terminate_segment_reader(self) -> None:
        if self._segment_reader:
            # NB: we have to wake up the segment reader thread in order to terminate it
            self._segment_reader.awake.set()
            self._segment_reader.terminate.set()
            self._segment_reader.join(5.0)
            if self._segment_reader.is_alive():
                log.warning('Unable to terminate existing segment reader thread...')
            self._segments_buffer.clear()

    def _sample_batches(
        self, num_batches: int | None = None, stop_on_sentinel: bool = False,
    ) -> Iterator[Batch]:
        # Ensure the SIGINT Event flag is cleared of any previous interrupts
        self._interrupt.clear()

        assert self._segments_buffer.maxlen is not None
        num_sampled = 0

        while num_batches is None or num_sampled < num_batches:
            segments: list[TSSegment] = []
            while len(segments) < self.batch_size and not self._interrupt.is_set():
                if len(self._segments_buffer) < int(self._segments_buffer.maxlen / 2):
                    self._wake_segment_reader()

                try:
                    segment = self._segments_buffer.pop()
                    if segment is None:
                        # Signal to start a new batch (cf. sequential sampling)
                        # Only yield the batch if we have at least one segment
                        # in the segments list (otherwise the sentinel signal
                        # has come at the transition between batches)
                        if stop_on_sentinel:
                            if segments:
                                # Yield the batch if we have any valid segments
                                batch = Batch._from_segments(
                                    segments=segments,
                                    as_tensor=self.as_tensor,
                                    pad_value=self.pad_value,
                                    output_format=self.output_format,
                                )
                                yield batch
                            return
                        else:
                            if segments:
                                # Break out of the inner sampling loop (and then yield)
                                break
                    else:
                        segments.append(segment)
                except IndexError:
                    # Deque is empty: wait for 1s for new segments to arrive
                    self._interrupt.wait(1)

            if self._interrupt.is_set():
                return
            else:
                batch = Batch._from_segments(
                    segments=segments,
                    as_tensor=self.as_tensor,
                    pad_value=self.pad_value,
                    output_format=self.output_format,
                )
                num_sampled += 1
                yield batch

    def write_batches(
        self, dataset: IsDataset, num_batches: int, path: str | Path,
        *, prefix: str = 'batch_', offset: int = 0,
        filesystem: fs.FileSystem | None = None,
        fs_config: dict[str, Any] | None = None,
        overwrite: bool = False, reset: bool = False,
    ) -> None:
        """Write sampled batches to disk (local directory or cloud storage).

        Args:
            dataset: Dataset to sample from.
            num_batches: Number of batches to write.
            path: Output directory path.
            prefix: Filename prefix for each batch.
            offset: Starting index to use for batch numbering.
            filesystem: Optional [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) instance to write to.
            fs_config: Custom arguments to pass to [PyArrow filesystem](https://arrow.apache.org/docs/python/api/filesystems.html#filesystem-implementations) if `filesystem` is not provided.
            overwrite: If True, overwrite existing files.
            reset: If True, reset the sampler's internal state before sampling.

        Returns:
            None.
        """  # noqa: E501
        batch_num = offset + 1
        for batch in self(dataset=dataset, num_batches=num_batches, reset=reset):
            filepath = f'{path}/{prefix}{batch_num}.pq'
            try:
                batch.to_parquet(
                    filepath=filepath, filesystem=filesystem, fs_config=fs_config,
                    overwrite=overwrite,
                )
            except FileExistsError as e:
                log.warning(e.args[0])

            batch_num += 1

    def __enter__(self) -> SamplerBase:
        return self

    def __exit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self._terminate_segment_reader()

    def __call__(
        self, dataset: IsDataset, num_batches: int | None = None,
        *, reset: bool = False,
    ) -> Iterator[Batch]:
        """Return an iterator over sampled batches from a dataset.

        Args:
            dataset: Dataset to sample from.
            num_batches: Maximum number of batches to yield. If None, yield until
                the dataset is exhausted.
            reset: If True, reset the sampler's internal state before sampling.

        Returns:
            Iterator over `Batch` objects.
        """
        if reset or dataset != self._dataset:
            self._dataset = dataset
            self._init_sampler()
            self._start_segment_reader()
        else:
            self._wake_segment_reader()

        yield from self._sample_batches(num_batches=num_batches)
