from typing import *

from cognite.client._api_client import APIClient
from cognite.client.data_classes import Database, DatabaseList, Row, RowList, Table, TableList
from cognite.client.utils import _utils as utils


class RawAPI(APIClient):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.databases = RawDatabasesAPI(*args, **kwargs)
        self.tables = RawTablesAPI(*args, **kwargs)
        self.rows = RawRowsAPI(*args, **kwargs)


class RawDatabasesAPI(APIClient):
    _RESOURCE_PATH = "/raw/dbs"
    _LIST_CLASS = DatabaseList

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._LIST_LIMIT = 10000

    def __call__(self, chunk_size: int = None) -> Generator[Union[Database, DatabaseList], None, None]:
        """Iterate over databases

        Fetches dbs as they are iterated over, so you keep a limited number of dbs in memory.

        Args:
            chunk_size (int, optional): Number of dbs to return in each chunk. Defaults to yielding one db a time.
        """
        return self._list_generator(chunk_size=chunk_size, method="GET")

    def __iter__(self) -> Generator[Database, None, None]:
        return self.__call__()

    def create(self, name: Union[str, List[str]]) -> Union[Database, DatabaseList]:
        """Create one or more databases.

        Args:
            name (Union[str, List[str]]): A db name or list of db names to create.

        Returns:
            Union[Database, DatabaseList]: Database or list of databases that has been created.

        Examples:

            Create a new database::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> res = c.raw.databases.create("db1")
        """
        utils.assert_type(name, "name", [str, list])
        if isinstance(name, str):
            items = {"name": name}
        else:
            items = [{"name": n} for n in name]
        return self._create_multiple(items=items)

    def delete(self, name: Union[str, List[str]]) -> None:
        """Delete one or more databases.

        Args:
            name (Union[str, List[str]]): A db name or list of db names to delete.

        Returns:
            None

        Examples:

            Delete a list of databases::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> res = c.raw.databases.delete(["db1", "db2"])
        """
        utils.assert_type(name, "name", [str, list])
        if isinstance(name, str):
            name = [name]
        items = [{"name": n} for n in name]
        chunks = utils.split_into_chunks(items, self._DELETE_LIMIT)
        tasks = [{"url_path": self._RESOURCE_PATH + "/delete", "json": {"items": chunk}} for chunk in chunks]
        utils.execute_tasks_concurrently(self._post, tasks, max_workers=self._max_workers)

    def list(self, limit: int = None) -> DatabaseList:
        """List databases

        Args:
            limit (int, optional): Maximum number of databases to return. If not specified, all dbs will be returned.

        Returns:
            DatabaseList: List of requested databases.

        Examples:

            List the first 5 databases::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> db_list = c.raw.databases.list(limit=5)

            Iterate over databases::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for db in c.raw.databases:
                ...     db # do something with the db

            Iterate over chunks of databases to reduce memory load::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for db_list in c.raw.databases(chunk_size=2500):
                ...     db_list # do something with the dbs
        """
        return self._list(method="GET", limit=limit)


class RawTablesAPI(APIClient):
    _RESOURCE_PATH = "/raw/dbs/{}/tables"
    _LIST_CLASS = TableList

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._LIST_LIMIT = 10000

    def __call__(self, db_name: str, chunk_size: int = None) -> Generator[Union[Table, TableList], None, None]:
        """Iterate over tables

        Fetches tables as they are iterated over, so you keep a limited number of tables in memory.

        Args:
            db_name (str): Name of the database to iterate over tables for
            chunk_size (int, optional): Number of tables to return in each chunk. Defaults to yielding one table a time.
        """
        for tb in self._list_generator(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name),
            chunk_size=chunk_size,
            method="GET",
        ):
            yield self._set_db_name_on_tables(tb, db_name)

    def create(self, db_name: str, name: Union[str, List[str]]) -> Union[Table, TableList]:
        """Create one or more tables.

        Args:
            db_name (str): Database to create the tables in.
            name (Union[str, List[str]]): A table name or list of table names to create.

        Returns:
            Union[Table, TableList]: Table or list of tables that has been created.

        Examples:

            Create a new table in a database::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> res = c.raw.tables.create("db1", "table1")
        """
        utils.assert_type(name, "name", [str, list])
        if isinstance(name, str):
            items = {"name": name}
        else:
            items = [{"name": n} for n in name]
        tb = self._create_multiple(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name), items=items
        )
        return self._set_db_name_on_tables(tb, db_name)

    def delete(self, db_name: str, name: Union[str, List[str]]) -> None:
        """Delete one or more tables.

        Args:
            db_name (str): Database to delete tables from.
            name (Union[str, List[str]]): A table name or list of table names to delete.

        Returns:
            None

        Examples:

            Delete a list of tables::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> res = c.raw.tables.delete("db1", ["table1", "table2"])
        """
        utils.assert_type(name, "name", [str, list])
        if isinstance(name, str):
            name = [name]
        items = [{"name": n} for n in name]
        chunks = utils.split_into_chunks(items, self._DELETE_LIMIT)
        tasks = [
            {
                "url_path": utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name) + "/delete",
                "json": {"items": chunk},
            }
            for chunk in chunks
        ]
        utils.execute_tasks_concurrently(self._post, tasks, max_workers=self._max_workers)

    def list(self, db_name: str, limit: int = None) -> TableList:
        """List tables

        Args:
            db_name (str): The database to list tables from.
            limit (int, optional): Maximum number of tables to return. If not specified, all tables will be returned.

        Returns:
            TableList: List of requested tables.

        Examples:

            List the first 5 tables::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> table_list = c.raw.tables.list("db1", limit=5)

            Iterate over tables::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for table in c.raw.tables(db_name="db1"):
                ...     table # do something with the table

            Iterate over chunks of tables to reduce memory load::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for table_list in c.raw.tables(db_name="db1", chunk_size=2500):
                ...     table_list # do something with the tables
        """
        tb = self._list(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name), method="GET", limit=limit
        )
        return self._set_db_name_on_tables(tb, db_name)

    def _set_db_name_on_tables(self, tb: Union[Table, TableList], db_name: str) -> Union[Table, TableList]:
        if isinstance(tb, Table):
            tb._db_name = db_name
            return tb
        elif isinstance(tb, TableList):
            for t in tb:
                t._db_name = db_name
            return tb
        raise TypeError("tb must be Table or TableList")


class RawRowsAPI(APIClient):
    _RESOURCE_PATH = "/raw/dbs/{}/tables/{}/rows"
    _LIST_CLASS = RowList

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._LIST_LIMIT = 10000

    def __call__(
        self, db_name: str, table_name: str, chunk_size: int = None
    ) -> Generator[Union[Row, RowList], None, None]:
        """Iterate over rows.

        Fetches rows as they are iterated over, so you keep a limited number of rows in memory.

        Args:
            db_name (str): Name of the database
            table_name (str): Name of the table to iterate over rows for
            chunk_size (int, optional): Number of rows to return in each chunk. Defaults to yielding one row a time.
        """
        return self._list_generator(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name, table_name),
            chunk_size=chunk_size,
            method="GET",
        )

    def insert(
        self, db_name: str, table_name: str, row: Union[List[Row], Row, Dict], ensure_parent: bool = False
    ) -> None:
        """Insert one or more rows into a table.

        Args:
            db_name (str): Name of the database.
            table_name (str): Name of the table.
            row (Union[List[Row], Row, Dict]): The row(s) to insert
            ensure_parent (bool): Create database/table if they don't already exist.

        Returns:
            None

        Examples:

            Insert new rows into a table::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> rows = {"r1": {"col1": "val1", "col2": "val1"}, "r2": {"col1": "val2", "col2": "val2"}}
                >>> res = c.raw.rows.insert("db1", "table1", rows)
        """
        chunks = self._process_row_input(row)

        tasks = [
            {
                "url_path": utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name, table_name),
                "json": {"items": chunk},
                "params": {"ensureParent": ensure_parent},
            }
            for chunk in chunks
        ]
        utils.execute_tasks_concurrently(self._post, tasks, max_workers=self._max_workers)

    def _process_row_input(self, row: List[Union[List, Dict, Row]]):
        utils.assert_type(row, "row", [list, dict, Row])
        rows = []
        if isinstance(row, dict):
            for key, columns in row.items():
                rows.append({"key": key, "columns": columns})
        elif isinstance(row, list):
            for elem in row:
                if isinstance(elem, Row):
                    rows.append(elem.dump(camel_case=True))
                else:
                    raise TypeError("list elements must be Row objects.")
        elif isinstance(row, Row):
            rows.append(row.dump(camel_case=True))
        return utils.split_into_chunks(rows, self._CREATE_LIMIT)

    def delete(self, db_name: str, table_name: str, key: Union[str, List[str]]) -> None:
        """Delete rows from a table.

        Args:
            db_name (str): Name of the database.
            table_name (str): Name of the table.
            key (Union[str, List[str]]): The key(s) of the row(s) to delete.

        Returns:
            None

        Examples:

            Delete rows from table::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> keys_to_delete = ["k1", "k2", "k3"]
                >>> res = c.raw.rows.delete("db1", "table1", keys_to_delete)
        """
        utils.assert_type(key, "key", [str, list])
        if isinstance(key, str):
            key = [key]
        items = [{"key": k} for k in key]
        chunks = utils.split_into_chunks(items, self._DELETE_LIMIT)
        tasks = [
            (
                {
                    "url_path": utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name, table_name) + "/delete",
                    "json": {"items": chunk},
                }
            )
            for chunk in chunks
        ]
        utils.execute_tasks_concurrently(self._post, tasks, max_workers=self._max_workers)

    def retrieve(self, db_name: str, table_name: str, key: str) -> Row:
        """Retrieve a single row by key.

        Args:
            db_name (str): Name of the database.
            table_name (str): Name of the table.
            key (str): The key of the row to retrieve.

        Returns:
            Row: The requested row.

        Examples:

            Retrieve a row with key 'k1' from tablew 't1' in database 'db1'::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> row = c.raw.rows.retrieve("db1", "t1", "k1")
        """
        return self._retrieve(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name, table_name), id=key
        )

    def list(self, db_name: str, table_name: str, limit: int = None) -> RowList:
        """List rows in a table.

        Args:
            db_name (str): Name of the database.
            table_name (str): Name of the table.
            limit (int): The number of rows to retrieve.

        Returns:
            RowList: The requested rows.

        Examples:

            List rows::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> row_list = c.raw.rows.list("db1", "t1", limit=5)

            Iterate over rows::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for row in c.raw.rows(db_name="db1", table_name="t1"):
                ...     row # do something with the row

            Iterate over chunks of rows to reduce memory load::

                >>> from cognite.client import CogniteClient
                >>> c = CogniteClient()
                >>> for row_list in c.raw.rows(db_name="db1", table_name="t1", chunk_size=2500):
                ...     row_list # do something with the rows
        """
        return self._list(
            resource_path=utils.interpolate_and_url_encode(self._RESOURCE_PATH, db_name, table_name),
            limit=limit,
            method="GET",
        )
