from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Union, TYPE_CHECKING

import pyarrow as pa

from yggdrasil.data.cast import CastOptions
from yggdrasil.dataclasses.waiting import WaitingConfigArg
from yggdrasil.io.enums import SaveMode
from .statement_result import StatementResult

if TYPE_CHECKING:
    import polars
    import pandas
    import pyspark


__all__ = [
    "SQLEngine"
]


@dataclass
class SQLEngine(ABC):

    @abstractmethod
    def execute(
        self,
        statement: str,
        *,
        row_limit: Optional[int] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        wait: WaitingConfigArg = True,
        **kwargs
    ) -> StatementResult:
        raise NotImplementedError

    @abstractmethod
    def insert_into(
        self,
        data: Union[
            pa.Table, pa.RecordBatch, pa.RecordBatchReader,
            dict, list, str,
            "pandas.DataFrame", "polars.DataFrame",
            "pyspark.sql.DataFrame"
        ],
        *,
        mode: SaveMode | str | None = None,
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        cast_options: Optional[CastOptions] = None,
        overwrite_schema: bool | None = None,
        match_by: Optional[list[str]] = None,
        wait: WaitingConfigArg = True
    ):
        raise NotImplementedError
