from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generic, Iterator, TypeVar

from daft.recordbatch import MicroPartition
from daft.schema import Schema

T = TypeVar("T")


@dataclass
class WriteOutput(Generic[T]):
    """Wrapper for output of the DataSink's `.write()` method.

    Attributes:
        output: The actual output from the write operation
        bytes_written: Size of the written data in bytes
        rows_written: Number of rows written
    """

    output: T
    bytes_written: int
    rows_written: int


class DataSink(ABC, Generic[T]):
    """Interface for writing data to a sink that is not built-in.

    When a DataFrame is written using the `.write_sink()` method, the following sequence occurs:

    1. The sink's `.start()` method is called once at the beginning of the write process.
    2. The DataFrame is executed, and its output is split into micropartitions.
    3. The sink's `.write()` method is invoked on each micropartition, potentially in parallel
       and distributed across multiple tasks or workers.
    4. After all writes complete, the resulting `WriteOutput` objects are gathered on a single node.
    5. The `.finalize()` method is then called with all write outputs to produce a final `MicroPartition`.

    Warning:
        This API is early in its development and is subject to change.
    """

    def name(self) -> str:
        """Optional custom sink name."""
        return "User-defined Data Sink"

    @abstractmethod
    def schema(self) -> Schema:
        """The expected schema for the micropartition returned by the `.finalize()` method of this DataSink.

        If this given schema does not match the actual schema of the micropartition at runtime, we throw an error.
        """
        raise NotImplementedError

    def start(self) -> None:
        """Optional callback for when a write operation begins.

        For example, this can be used to initialize resources, open connections, start a transaction etc.
        The default implementation does nothing.
        """
        pass

    @abstractmethod
    def write(self, micropartitions: Iterator[MicroPartition]) -> Iterator[WriteOutput[T]]:
        """Writes a stream of micropartitions to the sink.

        This method should handle the ingestion of each micropartition and yield a result
        (e.g. metadata) for each successful write.

        Args:
            micropartitions (Iterator[MicroPartition]): An iterator of micropartitions to be written.

        Returns:
            Iterator[WriteOutput[T]]: An iterator of write results wrapped in a WriteOutput.
        """
        raise NotImplementedError

    @abstractmethod
    def finalize(self, write_outputs: list[WriteOutput[T]]) -> MicroPartition:
        """Finalizes the write process and returns a resulting micropartition.

        For example, this can be used to merge, summarize, or commit the results of individual writes
        into a single output micropartition.

        Args:
            write_outputs (list[WriteOutput[T]]): The list of results from the calls to `.write()`.

        Returns:
            MicroPartition: A final, single micropartition representing the result of all writes.
        """
        raise NotImplementedError
