from math import ceil
import pandas as pd
import requests
from tqdm import tqdm

from .data_loader import Data_Loader, str2json, _url_error_msg, _process_date, _default_limit, _use_gpd_force, _has_gpd
from ..datetime_parser import to_datetime
from ..exceptions import OPD_DataUnavailableError, OPD_TooManyRequestsError
from .. import log

if _has_gpd:
    import geopandas as gpd
    from shapely.geometry import Point

logger = log.get_logger()

class Carto(Data_Loader):
    """
    A class for accessing data from Carto clients

    Parameters
    ----------
    url : str
        URL
    data_set : str
        Dataset ID
    date_field : str
        Name of the column that contains the date
    query : str
        Query that will be perfored for each request

    Methods
    -------
    load(year=None, nrows=None, offset=0, pbar=True)
        Load data for query
    get_count(year=None, where=None)
        Get number of records/rows generated by query
    get_years()
        Get years contained in data set
    """
    
    def __init__(self, url, data_set, date_field=None, query=None):
        '''Create Carto object

        Parameters
        ----------
        url : str
            URL or username for Carto data
        data_set : str
            Dataset ID
        date_field : str
            (Optional) Name of the column that contains the date
        query : str, dict
            (Optional) Additional query that will be added to each request
        '''

        # https://carto.com/developers/sql-api/guides/making-calls/
        # Format of URL is https://{username}.carto.com/api/v2/sql

        # Strip out username to ensure that URL is properly formatted
        username = url
        if username.startswith("https://"):
            username = username.replace("https://", "")

        if ".carto" in username:
            username = username[:username.find(".carto")]

        url_clean = "https://" + username + ".carto.com/api/v2/sql"

        self.url = url_clean
        self.data_set = data_set
        self.date_field = date_field
        self.query = str2json(query)

    
    def isfile(self):
        '''Returns False to indicate that Carto data is not file-based

        Returns
        -------
        False
        '''
        return False
    
    def get_api_url(self):
        return f'{self.url}?q=SELECT * FROM {self.data_set}'


    def get_count(self, year=None, **kwargs):
        '''Get number of records for a data request
        
        Parameters
        ----------
        year : int, list
            (Optional) Either the year or the year range [first_year, last_year] for the data that is being requested.  None value returns data for all years.
            
        Returns
        -------
        int
            Record count or number of rows in data request
        '''

        if self._last_count is not None and self._last_count[0]==year:
            logger.debug("Request matches previous count request. Returning saved count.")
            return self._last_count[1]
        else:
            where = self.__construct_where(year)
            json = self.__request(where=where, return_count=True)
            count = json["rows"][0]["count"]

        self._last_count = (year, count, where)

        return count


    def __request(self, where=None, return_count=False, out_fields="*", out_type="GeoJSON", offset=0, count=None):

        query = "SELECT "
        params = {}
        if return_count:
            query+="count(*)"
        else:
            query+=out_fields
            params["format"] = out_type

        query+=f" FROM {self.data_set}"

        if where != None:
            query+=" WHERE "+ where

        default_where = ''
        for k,v in self.query.items():
            if isinstance(v,str):
                v = f"'{v}'"
            default_where += f" AND {k}={v}"

        if len(default_where):
            if where != None:
                query+=default_where
            else:
                query+=" WHERE"+ default_where[4:]

        if not return_count and count!=0:
            # Order results to ensure data order remains constant if paging
            query+=" ORDER BY cartodb_id"

        query+=f" OFFSET {offset}"

        if count!=None:
            query+=f" LIMIT {count}"

        params["q"] = query

        logger.debug(f"Request data from {self.url}")
        for k,v in params.items():
            logger.debug(f"\t{k} = {v}")

        r = requests.get(self.url, params=params)

        try:
            r.raise_for_status()
        except requests.HTTPError as e:
            if len(e.args)>0:
                if "503 Server Error" in e.args[0]:
                    raise OPD_DataUnavailableError(self.get_api_url(), e.args, _url_error_msg.format(self.get_api_url()))
                else:
                    raise

            else: raise e
        except: raise
        
        return r.json()


    def __construct_where(self, year=None):
        if self.date_field!=None and year!=None:
            start_date, stop_date = _process_date(year, date_field=self.date_field)
            where_query = f"{self.date_field} >= '{start_date}' AND {self.date_field} <= '{stop_date}'"
        else:
            where_query = None

        return where_query

    
    def load(self, year=None, nrows=None, offset=0, *, pbar=True, format_date=True, **kwargs):
        '''Download table to pandas or geopandas DataFrame
        
        Parameters
        ----------
        year : int, list
            (Optional) Either the year or the year range [first_year, last_year] for the data that is being requested.  None value returns data for all years.
        nrows : int
            (Optional) Only returns the first nrows rows of the table
        offset - int
            (Optional) Number of records to offset from first record. Default is 0 to return records starting from the first.
        pbar : bool
            (Optional) If true (default), a progress bar will be displayed
        format_date : bool, optional
            If True, known date columns (based on presence of date_field in datasets table or data type information provided by dataset owner) will be automatically formatted
            to be pandas datetimes (or pandas Period in rare cases), by default True
            
        Returns
        -------
        pandas or geopandas DataFrame
            DataFrame containing downloaded
        '''
        
        if self._last_count is not None and self._last_count[0]==year:
            record_count = self._last_count[1]
            where_query = self._last_count[2]
        else:
            where_query = self.__construct_where(year)
            json = self.__request(where=where_query, return_count=True)
            record_count = json["rows"][0]["count"]
            self._last_count = (year, record_count, where_query)

        record_count-=offset
        if record_count<=0:
            return pd.DataFrame()

        batch_size = _default_limit
        nrows = nrows if nrows!=None and record_count>=nrows else record_count
        batch_size = nrows if nrows < batch_size else batch_size
        num_batches = ceil(nrows / batch_size)
            
        pbar = pbar and num_batches>1
        if pbar:
            bar = tqdm(desc=self.url, total=nrows, leave=False)

        # When requesting data as GeoJSON, no type information is returned so request it now
        type_info = self.__request(count=0, out_type="JSON")
            
        features = []
        for batch in range(num_batches):
            bs = batch_size if batch<num_batches-1 else nrows-batch*batch_size

            try:
                data = self.__request(where=where_query, offset=offset+batch*batch_size, count=bs)
                features.extend(data["features"])

                if batch==0 and len(features)>0:
                    date_cols = [key for key, x in type_info["fields"].items() if x["type"]=='date']
                    if len(data["features"]) not in [batch_size, nrows]:
                        num_rows = len(data["features"])
                        raise ValueError(f"Number of rows is {num_rows} but is expected to be max rows to read {batch_size} or total number of rows {nrows}")
            except Exception as e:
                if len(e.args)>0 and "Error Code: 429" in e.args[0]:
                    raise OPD_TooManyRequestsError(self.url, *e.args, _url_error_msg.format(self.get_api_url()))
                else:
                    raise
            except:
                raise

            if pbar:
                bar.update(len(data["features"]))

        if pbar:
            bar.close()

        df = pd.DataFrame.from_records([x["properties"] for x in features])
        if format_date:
            for col in date_cols:
                if col in df:
                    logger.debug(f"Column {col} had a data type of date. Converting values to datetime objects.")
                    df[col] = to_datetime(df[col])

        if len(df) > 0:
            has_point_geometry = any("geometry" in x and x["geometry"]!=None for x in features)
            if has_point_geometry:
                if _use_gpd_force is not None:
                    if not _has_gpd and _use_gpd_force:
                        raise ValueError("User cannot force GeoPandas usage when it is not installed")
                    use_gpd = _use_gpd_force
                else:
                    use_gpd = _has_gpd

                if use_gpd:
                    geometry = []
                    for feat in features:
                        if "geometry" not in feat or feat["geometry"]==None or len(feat["geometry"]["coordinates"])<2:
                            geometry.append(None)
                        else:
                            geometry.append(Point(feat["geometry"]["coordinates"][0], feat["geometry"]["coordinates"][1]))

                    logger.debug("Geometry found. Contructing geopandas GeoDataFrame")
                    df = gpd.GeoDataFrame(df, crs=4326, geometry=geometry)
                else:
                    geometry = [feat["geometry"] if "geometry" in feat else None for feat in features]

                    if "geolocation" not in df:
                        logger.debug("Adding geometry column generated from spatial data provided by request.")
                        df["geolocation"] = geometry
                    else:
                        raise KeyError("geolocation already exists in DataFrame")

            return df
        else:
            return pd.DataFrame()

