#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import codecs
import csv
import logging
import time
import zipfile

import requests
from requests.exceptions import HTTPError
from six import string_types, BytesIO, PY2, PY3

from .models import WrappedJson
from .utils import conf_to_xml

logger = logging.getLogger(__name__)
long_type = int

if PY3:
    from collections.abc import Iterable as Iterable
elif PY2:
    from collections import Iterable as Iterable

    long_type = long


class LensQuery(WrappedJson):
    def __init__(self, client, *args, **kwargs):
        super(LensQuery, self).__init__(*args, **kwargs)
        self.client = client

    @property
    def finished(self):
        return self.status.status in ('SUCCESSFUL', 'FAILED', 'CANCELED', 'CLOSED')

    def get_result(self, *args, **kwargs):
        return self.client.get_result(self, *args, **kwargs)

    result = property(get_result)


type_mappings = {'BOOLEAN': bool,
                 'TINYINT': int,
                 'SMALLINT': int,
                 'INT': int,
                 'BIGINT': long_type,
                 'FLOAT': float,
                 'DOUBLE': float,
                 'TIMESTAMP': long_type,
                 'BINARY': bin,
                 'ARRAY': lambda x: list(y.strip() for y in x.split(",")),
                 'MAP': dict,
                 # 'STRUCT,': str,
                 # 'UNIONTYPE,': float,
                 # 3'USER_DEFINED,': float,
                 'DECIMAL,': float,
                 # 'NULL,': float,
                 # 'DATE,': float,
                 # 'VARCHAR,': float,
                 # 'CHAR': float
                 }
default_mapping = lambda x: x


class LensQueryResult(Iterable):
    def __init__(self, custom_mappings=None):
        if custom_mappings is None:
            custom_mappings = {}
        self.custom_mappings = custom_mappings

    def _mapping(self, type_name):
        if type_name in self.custom_mappings:
            return self.custom_mappings[type_name]
        if type_name in type_mappings:
            return type_mappings[type_name]
        return default_mapping


class LensInMemoryResult(LensQueryResult):
    def __init__(self, resp, custom_mappings=None):
        super(LensInMemoryResult, self).__init__(custom_mappings)
        self.rows = resp.in_memory_query_result.rows

    def __iter__(self):
        for row in self.rows:
            yield list(self._mapping(value.type)(value.value) if value else None for value in row['values'])


class LensPersistentResult(LensQueryResult):
    def __init__(self, header, response, encoding=None, is_header_present=True, delimiter=",",
                 custom_mappings=None):
        super(LensPersistentResult, self).__init__(custom_mappings)
        self.response = response
        self.is_zipped = 'zip' in self.response.headers['content-disposition']
        self.delimiter = str(delimiter)
        self.is_header_present = is_header_present
        self.encoding = encoding
        self.header = header

    def _parse_line(self, line):
        return list(self._mapping(self.header.columns[index].type)(line[index]) for index in range(len(line)))

    def get_csv_reader(self, file):
        if PY3:
            file = codecs.iterdecode(file, 'utf-8')
        return csv.reader(file, delimiter=self.delimiter)

    def __iter__(self):
        if self.is_zipped:
            byte_stream = BytesIO(self.response.content)
            with zipfile.ZipFile(byte_stream) as self.zipfile:
                for name in self.zipfile.namelist():
                    with self.zipfile.open(name) as single_file:
                        if name[-3:] == 'csv':
                            reader = self.get_csv_reader(single_file)
                        else:
                            reader = single_file
                        reader_iterator = iter(reader)
                        if self.is_header_present:
                            next(reader_iterator)
                        for line in reader_iterator:
                            yield self._parse_line(line)
            byte_stream.close()
        else:
            stream = codecs.iterdecode(self.response.iter_lines(),
                                       self.response.encoding or self.response.apparent_encoding)
            reader = csv.reader(stream, delimiter=self.delimiter)
            reader_iterator = iter(reader)
            if self.is_header_present:
                next(reader_iterator)
            for line in reader_iterator:
                yield self._parse_line(line)
            stream.close()


class LensQueryClient(object):
    def __init__(self, base_url, session):
        self._session = session
        self.base_url = base_url + "queryapi/"
        self.launched_queries = []
        self.finished_queries = {}
        self.query_confs = {}
        self.is_header_present_in_result = self._session['lens.query.output.write.header'].lower() \
                                           in ['true', '1', 't', 'y', 'yes', 'yeah', 'yup']

    def __call__(self, **filters):
        filters['sessionid'] = self._session._sessionid
        resp = requests.get(self.base_url + "queries/", params=filters, headers={'accept': 'application/json'})
        return self.sanitize_response(resp)

    def __getitem__(self, item):
        if isinstance(item, string_types):
            if item in self.finished_queries:
                return self.finished_queries[item]
            resp = requests.get(self.base_url + "queries/" + item, params={'sessionid': self._session._sessionid},
                                headers={'accept': 'application/json'})
            resp.raise_for_status()
            query = LensQuery(self, resp.json(object_hook=WrappedJson))
            if query.finished:
                query.client = self
                self.finished_queries[item] = query
            return query
        elif isinstance(item, LensQuery):
            return self[item.query_handle]
        elif isinstance(item, WrappedJson):
            if item._is_wrapper:
                return self[item._wrapped_value]
            if item.query_handle:
                return self[item.query_handle]
        raise Exception("Can't get query: " + str(item))

    def submit(self, query, operation=None, query_name=None, timeout=None, conf=None, wait=False, fetch_result=False,
               *args, **kwargs):
        payload = [('sessionid', self._session._sessionid), ('query', query)]
        if query_name:
            payload.append(('queryName', query_name))
        if timeout:
            payload.append(('timeoutmillis', str(int(timeout) * 1000)))
        if not operation:
            operation = "execute_with_timeout" if timeout else "execute"
        payload.append(('operation', operation))
        payload.append(('conf', conf_to_xml(conf)))
        resp = requests.post(self.base_url + "queries/", files=payload, headers={'accept': 'application/json'})
        query = self.sanitize_response(resp)
        logger.info("Submitted query %s", query)
        if conf:
            self.query_confs[str(query)] = conf
        if fetch_result:
            # get result and return
            return self.get_result(query, *args, **kwargs)  # query is handle here
        elif wait:
            # fetch details and return
            return self.wait_till_finish(query, *args, **kwargs)
        # just return handle. This would be the async case. Or execute with timeout, without wait
        return query

    def wait_till_finish(self, handle_or_query, poll_interval=5, *args, **kwargs):
        while not self[handle_or_query].finished:
            time.sleep(poll_interval)
        return self[handle_or_query]

    def get_result(self, handle_or_query, *args, **kwargs):
        query = self.wait_till_finish(handle_or_query, *args, **kwargs)
        handle = str(query.query_handle)
        if query.status.status == 'SUCCESSFUL' and query.status.is_result_set_available:
            resp = requests.get(self.base_url + "queries/" + handle + "/resultsetmetadata",
                                params={'sessionid': self._session._sessionid}, headers={'accept': 'application/json'})
            metadata = self.sanitize_response(resp)
            # Try getting the result through http result
            resp = requests.get(self.base_url + "queries/" + handle + "/httpresultset",
                                params={'sessionid': self._session._sessionid}, stream=True)
            if resp.ok:
                is_header_present = self.is_header_present_in_result
                if handle in self.query_confs and 'lens.query.output.write.header' in self.query_confs[handle]:
                    is_header_present = bool(self.query_confs[handle]['lens.query.output.write.header'])
                return LensPersistentResult(metadata, resp, is_header_present=is_header_present, *args, **kwargs)
            else:
                response = requests.get(self.base_url + "queries/" + handle + "/resultset",
                                        params={'sessionid': self._session._sessionid},
                                        headers={'accept': 'application/json'})
                resp = self.sanitize_response(response)
                # If it has in memory result, return inmemory result iterator
                if resp._is_wrapper and resp._wrapped_key == u'inMemoryQueryResult':
                    return LensInMemoryResult(resp)
                # Else, return whatever you got
                return resp

        else:
            raise Exception("Result set not available")

    def sanitize_response(self, resp):
        try:
            resp_json = resp.json(object_hook=WrappedJson)
        except:
            resp_json = resp.json()
        if resp_json is not None:
            if 'lensAPIResult' in resp_json:
                resp_json = resp_json.lens_a_p_i_result
                if 'error' in resp_json:
                    error = resp_json['error']
                    if "stackTrace" in error:
                        logger.error(error['stackTrace'])
                    raise HTTPError(error, request=resp.request, response=resp)
                if 'data' in resp_json:
                    data = resp_json.data
                    if len(data) == 2 and 'type' in data:
                        keys = list(data.keys())
                        keys.remove('type')
                        return WrappedJson({data['type']: data[keys[0]]})
                    return data
        if resp_json is not None:
            return resp_json
        else:
            resp.raise_for_status()
        logger.error(resp.text)
        raise Exception("Unknown error with response", resp)