from datetime import datetime
from deltalake import DeltaTable

import pyarrow as pa
from typing import List, Optional, Tuple, Any, Union
from bmsdna.lakeapi.core.types import FileTypes
from bmsdna.lakeapi.context.df_base import FLAVORS, ExecutionContext, ResultData, get_sql
import pyarrow.dataset
import pypika.queries
from pypika.terms import Term
import pypika.functions
import pypika.enums
import pypika
import os
from datetime import datetime, timezone
from bmsdna.lakeapi.core.config import SearchConfig
from uuid import uuid4
from adbc_driver_sqlite.dbapi import Connection, Cursor, connect


class SqliteResultData(ResultData):
    def __init__(
        self,
        original_sql: Union[pypika.queries.QueryBuilder, str],
        connection: Connection,
        chunk_size: int,
    ) -> None:
        super().__init__(chunk_size=chunk_size)
        self.original_sql = original_sql
        self.connection = connection
        self._arrow_schema = None
        self._df = None

    def columns(self):
        return self.arrow_schema().names

    def query_builder(self) -> pypika.queries.QueryBuilder:
        return pypika.Query.from_(self.original_sql)

    def arrow_schema(self) -> pa.Schema:
        if self._arrow_schema is not None:
            return self._arrow_schema
        query = get_sql(self.original_sql, limit=10)  # limit of at least 1 is needed for sqlite!
        with self.connection.cursor() as cur:
            cur.execute(query)
            self._arrow_schema = cur.fetch_arrow_table().schema
        return self._arrow_schema

    @property
    def df(self):
        if self._df is None:
            query = get_sql(self.original_sql)
            with self.connection.cursor() as cur:
                cur.execute(query)
                self._df = cur.fetch_arrow_table()
        return self._df

    def to_pandas(self):
        return self.df.to_pandas()

    def to_arrow_table(self):
        return self.df

    def to_arrow_recordbatch(self, chunk_size: int = 10000):
        query = get_sql(self.original_sql)
        return self.df.to_batches(chunk_size)


class SqliteExecutionContext(ExecutionContext):
    def __init__(self, chunk_size: int):
        super().__init__(chunk_size=chunk_size)
        self.res_con = None
        self.connections: dict[str, Connection] = dict()
        self.persistance_file_name = None
        self.len_func = "length"

    def register_arrow(self, name: str, ds: Union[pyarrow.dataset.Dataset, pyarrow.Table]):
        raise NotImplementedError("Cannot read arrow in remote sql")

    def close(self):
        pass

    def execute_sql(
        self,
        sql: Union[
            pypika.queries.QueryBuilder,
            str,
        ],
    ) -> SqliteResultData:
        # todo: get correct connection string somehow
        assert len(self.connections) == 1
        return SqliteResultData(
            sql, chunk_size=self.chunk_size, connection=self.connections[list(self.connections.keys())[0]]
        )

    def distance_m_function(self, lat1: Term, lon1: Term, lat2: Term, lon2: Term):
        raise NotImplementedError("Not implemented")

    def json_function(self, term: Term, assure_string=False):
        raise NotImplementedError(
            "Cannot convert to JSON in remote sql"
        )  # we could but sql does not support structured types anyway, so...

    def init_search(
        self,
        source_view: str,
        search_configs: list[SearchConfig],
    ):
        raise NotImplementedError("Not supported")

    def register_datasource(
        self, name: str, uri: str, file_type: FileTypes, partitions: List[Tuple[str, str, Any]] | None
    ):
        assert file_type == "sqlite"
        self.connections[name] = connect(uri)

    def list_tables(self) -> ResultData:
        return self.execute_sql("SELECT type as table_type, name from sqlite_schema where type='table'")

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        for _, v in self.connections.items():
            v.close()
        self.connections = dict()
