from io import BytesIO
import logging
import pandas as pd
import re
import requests
from tqdm import tqdm
import urllib
import warnings
from zipfile import ZipFile

from .data_loader import Data_Loader, str2json, download_zip_and_extract, _url_error_msg, get_legacy_session, filter_dataframe
from ..datetime_parser import to_datetime
from ..exceptions import OPD_DataUnavailableError
from .. import httpio, log

logger = log.get_logger()

# https://stackoverflow.com/questions/73093656/progress-in-bytes-when-reading-csv-from-url-with-pandas
class TqdmReader:
    # Older versions of pandas check if reader has these properties even though they are not used
    write = []
    __iter__ = []
    def __init__(self, resp, pbar=True, nrows=None):
        total_size = int(resp.headers.get("Content-Length", 0))

        self.rows_read = 0
        if nrows != None:
            self.nrows = nrows
        else:
            self.nrows = float("inf")
        self.resp = resp
        self.pbar = pbar
        if self.pbar:
            self.bar = tqdm(
                desc=resp.url,
                total=total_size,
                unit="iB",
                unit_scale=True,
                unit_divisor=1024,
                leave=False
            )

        self.reader = self.read_from_stream()

    def read_from_stream(self):
        for line in self.resp.iter_lines():
            line += b"\n"
            if self.pbar:
                self.bar.update(len(line))
            yield line

    def read(self, n=0):
        try:
            if self.rows_read >= self.nrows:
                # Number of rows read is greater than user-requested limit
                return ""
            self.rows_read += 1
            return next(self.reader)
        except StopIteration:
            if self.pbar:
                self.bar.update(self.bar.total - self.bar.n)
            return ""

def read_zipped_csv(url, pbar=True, block_size=2**20, data_set=None):

    if data_set:
        logging.debug('Load CSV from zip using httpio method')
        # Load only requested dataset to minimize download size
        with httpio.open(url, block_size=block_size) as fp:
            with ZipFile(fp, 'r') as z:
                return pd.read_csv(BytesIO(z.read(data_set['file'])), encoding_errors='surrogateescape')
    else:
        logging.debug('Load CSV from zip by downloading and converting to pandas DataFrame')

        zip_data = download_zip_and_extract(url, block_size, pbar)
        zip_bytes_io = BytesIO(zip_data)
        logger.debug('Converting BytesIO to DataFrame')
        return pd.read_csv(zip_bytes_io, encoding_errors='surrogateescape')
  

def count_csv_rows(chunk_iter):
    if isinstance(chunk_iter, bytes):
        # Convert to iterator
        chunk_iter = [chunk_iter]
        
    count = 0
    open_quote = False
    no_quotes = True

    for chunk in chunk_iter:
        if no_quotes and chunk.count(b"\"")==0: # No need to worry about quotes
            count += chunk.count(b"\n")
        else: # Handle possible newlines in quotes
            no_quotes = False
            prev_end = chunk.find(b'\"') if open_quote else -1
            any_quotes = False
            while (next_start:=chunk.find(b'\"', prev_end+1))!=-1:
                any_quotes = True
                count += chunk[prev_end+1:next_start].count(b"\n")
                prev_end = chunk.find(b'\"', next_start+1)
                if prev_end==-1:
                    open_quote = True
                    break
            else:
                open_quote = False if any_quotes or prev_end!=-1 else open_quote # If not quotes found, leave unchanged
            
            if not open_quote:
                count+=chunk[prev_end+1:].count(b"\n")

    # Subtract off trailing newlines in last row
    newline = int.from_bytes(b"\n", "big")
    for c in reversed(chunk):
        if c==newline:
            count-=1
        else:
            break

    return count


class Csv(Data_Loader):
    """
    A class for accessing data from CSV download URLs

    Parameters
    ----------
    url : str
        URL
    date_field : str
        Name of the column that contains the date
    agency_field : str
        Name of column that contains agency names

    Methods
    -------
    load(year=None, nrows=None, offset=0, pbar=True, agency=None)
        Load data for query
    get_count(year=None, agency=None, force=False)
        Get number of records/rows generated by query
    get_years(force=False)
        Get years contained in data set
    """

    def __init__(self, url, date_field=None, agency_field=None, data_set=None, query=None):
        '''Create Csv object

        Parameters
        ----------
        url : str
            URL for CSV data
        date_field : str
            (Optional) Name of the column that contains the date
        agency_field : str
                (Optional) Name of the column that contains the agency name (i.e. name of the police departments)
        data_set : str
            (Optional) Name of CSV file in zip file if input is a zip file. Required if input is a zip file with more than one file.
        query : str, dict
            (Optional) Keys are data columns to filter. Values are values to filter for in columns.
        '''
        
        self.url = url
        self.date_field = date_field
        self.agency_field = agency_field
        self.data_set = data_set
        self.query = str2json(query)


    def isfile(self):
        '''Returns True to indicate that Csv data is file-based

        Returns
        -------
        True
        '''
        return True


    def get_count(self, year=None, *,  agency=None, force=False, **kwargs):
        '''Get number of records for a Csv 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.
        agency : str
            (Optional) Name of agency to filter for.
        force : bool
            (Optional) get_count for CSV file will only run if force=true. In many use cases, it will be more efficient to load the file and manually get the count.
            
        Returns
        -------
        int
            Record count or number of rows in data request
        '''

        logger.debug(f"Calculating row count for {self.url}")
        if self._last_count is not None and self._last_count[0] == (self.url, year, agency):
            logger.debug("Request matches previous count request. Returning saved count.")
            return self._last_count[1]
        if ".zip" not in self.url and year==None and agency==None and not self.query:
            logger.debug(f"Loading file to count rows from {self.url}")
            with requests.get(self.url, stream=True) as r:
                count = count_csv_rows(r.iter_content(chunk_size=2**16))
        elif force:
            count = len(self.load(year=year, agency=agency))
        else:
            raise ValueError("Extracting the number of records for a single year of a CSV file requires reading the whole file in. In most cases, "+
                "running load() with a year argument to load in the data and manually finding the record count will be more "
                "efficient. If running get_count with a year argument is still desired, set force=True")
        
        self._last_count = ((self.url, year, agency), count)
        return count


    def load(self, year=None, nrows=None, offset=0, *, pbar=True, agency=None, format_date=True, **kwargs):
        '''Download CSV file to pandas 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 CSV
        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
        agency : str
            (Optional) Name of the agency to filter for. None value returns data for all agencies.
        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 DataFrame
            DataFrame containing table imported from CSV
        '''

        if isinstance(nrows, float):
            nrows = int(nrows)
        
        logger.debug(f"Loading file from {self.url}")
        if ".zip" in self.url:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=pd.errors.DtypeWarning)
                try:
                    table = read_zipped_csv(self.url, pbar=pbar, data_set=self.data_set)
                    logger.debug("Completed reading CSV from zip file")
                except requests.exceptions.HTTPError as e:
                    if len(e.args) and 'Forbidden' in e.args[0]:
                        headers = {
                            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:106.0) Gecko/20100101 Firefox/106.0',
                            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8',
                            'Accept-Language': 'en-US,en;q=0.5',
                            # 'Accept-Encoding': 'gzip, deflate, br',
                            'DNT': '1',
                            'Connection': 'keep-alive',
                            'Upgrade-Insecure-Requests': '1',
                            'Sec-Fetch-Dest': 'document',
                            'Sec-Fetch-Mode': 'navigate',
                            'Sec-Fetch-Site': 'none',
                            'Sec-Fetch-User': '?1',
                        }
                        try:
                            table = pd.read_csv(self.url, encoding_errors='surrogateescape', storage_options=headers)
                        except urllib.error.HTTPError as e:
                            raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                        except:
                            raise
                    else:
                        raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                except Exception as e:
                    raise e
        else:
            use_legacy = False
            headers = None
            try:
                r = requests.head(self.url)
            except requests.exceptions.SSLError as e:
                if "[SSL: UNSAFE_LEGACY_RENEGOTIATION_DISABLED] unsafe legacy renegotiation disabled" in str(e.args[0]) or \
                    "[SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate" in str(e.args[0]):
                    use_legacy = True
                elif 'Max retries exceeded' in str(e):
                    raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                else:
                    raise e
            except requests.exceptions.ConnectionError as e:
                if 'Max retries exceeded' in str(e):
                    raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                else:
                    raise e
            except Exception as e:
                raise
                
            if not use_legacy:
                if r.status_code in [400,404]:
                    # Try get instead
                    r = requests.get(self.url)
                try:
                    r.raise_for_status()
                    r.close()
                except requests.exceptions.HTTPError as e:
                    try:
                        headers = {
                            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:106.0) Gecko/20100101 Firefox/106.0',
                            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8',
                            'Accept-Language': 'en-US,en;q=0.5',
                            # 'Accept-Encoding': 'gzip, deflate, br',
                            'DNT': '1',
                            'Connection': 'keep-alive',
                            'Upgrade-Insecure-Requests': '1',
                            'Sec-Fetch-Dest': 'document',
                            'Sec-Fetch-Mode': 'navigate',
                            'Sec-Fetch-Site': 'none',
                            'Sec-Fetch-User': '?1',
                        }
                        r = requests.get(self.url, headers=headers)
                        r.raise_for_status()
                        r.close()
                    except:
                        raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                except Exception as e:
                    raise e
            
            def get(url, use_legacy, headers=None):
                if use_legacy:
                    return get_legacy_session().get(url, params=None, stream=True, headers=headers)
                else:
                    return requests.get(url, params=None, stream=True, headers=headers)

            header = 'infer'   

            with get(self.url, use_legacy, headers) as resp:
                try:
                    with warnings.catch_warnings():
                        warnings.filterwarnings("ignore", message=r"Columns \(.+\) have mixed types", category=pd.errors.DtypeWarning)
                        nrows_read = offset+nrows if nrows is not None and not self.query else None
                        table = pd.read_csv(TqdmReader(resp, pbar=pbar), nrows=nrows_read, 
                            encoding_errors='surrogateescape', 
                            header=header)
                except (urllib.error.HTTPError, pd.errors.ParserError) as e:
                    raise OPD_DataUnavailableError(*e.args, _url_error_msg.format(self.url))
                except Exception as e:
                    raise e
                
        if len(table.columns)==1 and ('?xml' in table.columns[0] or re.search(r'^\<.+\>', table.columns[0])):
            # Read data was not a CSV file. It was an error code or HTML
            raise OPD_DataUnavailableError(table.iloc[0,0], _url_error_msg.format(self.url))
        
        table = filter_dataframe(table, date_field=self.date_field, year_filter=year, 
            agency_field=self.agency_field, agency=agency, format_date=format_date)
        
        if bool(self.query):
            for k,v in self.query.items():
                table = table[table[k]==v].reset_index(drop=True)

        if offset>0:
            rows_limit = offset+nrows if nrows is not None and offset+nrows<len(table) else len(table)
            logger.debug(f"Extracting {rows_limit} rows starting at {offset}")
            table = table.iloc[offset:rows_limit].reset_index(drop=True)
        if nrows is not None and len(table)>nrows:
            logger.debug(f"Extracting the first {nrows} rows")
            table = table.head(nrows)

        return table

    def get_years(self, *, force=False, **kwargs):
        '''Get years contained in data set
        
        Parameters
        ----------
        force : bool
            (Optional) If false, an exception will be thrown. It may be more efficient to load the table and extract years manually
            
        Returns
        -------
        list
            list containing years in data set
        '''

        if not force:
            raise ValueError("Extracting the years of a CSV file requires reading the whole file in. In most cases, "+
                "running load() with no arguments to load in the whole CSV file and manually finding the years will be more "
                "efficient. If running get_years is still desired, set force=True")
        else:
            if pd.isnull(self.date_field):
                raise ValueError("No date field provided to access year information")
            df = self.load()
            if self.date_field.lower()=="year":
                years = df[self.date_field].unique()
            else:
                date_col = to_datetime(df[self.date_field], ignore_errors=True)
                years = list(date_col.dt.year.dropna().unique())
            years.sort()
            return [int(x) for x in years]

