from carto.exceptions import CartoException, CartoRateLimitException

from .base_dataset import BaseDataset
from ...utils import is_sql_query


class QueryDataset(BaseDataset):
    def __init__(self, data, credentials, schema=None):
        super(QueryDataset, self).__init__(credentials)

        self._query = data

    @staticmethod
    def can_work_with(data):
        return is_sql_query(data)

    @classmethod
    def create(cls, data, credentials, schema=None):
        return cls(data, credentials)

    @property
    def query(self):
        return self._query

    @property
    def dataset_info(self):
        raise ValueError('The dataset_info method is not allowed on Datasets built on queries. '
                         'Use a table-based dataset instead: Dataset(my_table)')

    def update_dataset_info(self, privacy=None, table_name=None):
        raise ValueError('The dataset_info method is not allowed on Datasets built on queries. '
                         'Use a table-based dataset instead: Dataset(my_table)')

    def download(self, limit, decode_geom, retry_times):
        self._is_ready_for_dowload_validation()
        columns = self._get_query_columns()
        return self._copyto(columns, self._query, limit, decode_geom, retry_times)

    def upload(self, if_exists, with_lnglat):
        self._is_ready_for_upload_validation()

        if if_exists == BaseDataset.APPEND:
            raise CartoException('Error using append with a QueryDataset.'
                                 'It is not possible to append data to a query')
        elif if_exists == BaseDataset.REPLACE or not self.exists():
            self._create_table_from_query()
        elif if_exists == BaseDataset.FAIL:
            raise self._already_exists_error()

    def delete(self):
        raise ValueError('The delete method is not allowed on Datasets built on queries. '
                         'Use a table-based dataset instead: Dataset(my_table)')

    def get_query(self):
        return self._query

    def compute_geom_type(self):
        """Compute the geometry type from the data"""
        return self._get_geom_type(self._query)

    def get_table_names(self):
        query = "SELECT CDB_QueryTablesText('{}') as tables".format(self._query)
        result = self._context.execute_query(query)
        tables = []
        if result['total_rows'] > 0 and result['rows'][0]['tables']:
            # Dataset_info only works with tables without schema
            tables = [table.split('.')[1] if '.' in table else table for table in result['rows'][0]['tables']]

        return tables

    def _create_table_from_query(self):
        query = '''BEGIN; {drop}; {create}; {cartodbfy}; COMMIT;'''.format(
            drop=self._drop_table_query(),
            create=self._get_query_to_create_table_from_query(),
            cartodbfy=self._cartodbfy_query())

        try:
            self._context.execute_long_running_query(query)
        except CartoRateLimitException as err:
            raise err
        except CartoException as err:
            raise CartoException('Cannot create table: {}.'.format(err))

    def _get_query_to_create_table_from_query(self):
        return '''CREATE TABLE {table_name} AS ({query})'''.format(table_name=self._table_name, query=self._query)

@@ -1180,16 +1180,69 @@
 lumns()%0A
+        query = self._get_read_query(columns, limit)%0A
         
@@ -1270,22 +1270,16 @@
 olumns, 
-self._
 query, l
@@ -3328,20 +3328,635 @@
  query=self._query)%0A
+%0A    def _get_read_query(self, table_columns, limit=None):%0A        %22%22%22Create the read (COPY TO) query%22%22%22%0A        query_columns = %5Bcolumn.name for column in table_columns if column.name != 'the_geom_webmercator'%5D%0A%0A        query = 'SELECT %7Bcolumns%7D FROM (%7Bquery%7D) _q'.format(%0A            query=self._query,%0A            columns=', '.join(query_columns))%0A%0A        if limit is not None:%0A            if isinstance(limit, int) and (limit %3E= 0):%0A                query += ' LIMIT %7Blimit%7D'.format(limit=limit)%0A            else:%0A                raise ValueError(%22%60limit%60 parameter must an integer %3E= 0%22)%0A%0A        return query%0A

fix tables downloaded in query dataset