""" Helper functions to communicate with replication servers.
"""

import requests
import urllib.request as urlrequest
import datetime as dt
from collections import namedtuple
from math import ceil
from osmium import MergeInputReader
from osmium import io as oio
from osmium import version

import logging

LOG = logging.getLogger('pyosmium')
LOG.addHandler(logging.NullHandler())

OsmosisState = namedtuple('OsmosisState', ['sequence', 'timestamp'])
DownloadResult = namedtuple('DownloadResult', ['id', 'reader', 'newest'])

class ReplicationServer:
    """ Represents a server that publishes replication data. Replication
        change files allow to keep local OSM data up-to-date without downloading
        the full dataset again.
    """

    def __init__(self, url, diff_type='osc.gz'):
        self.baseurl = url
        self.diff_type = diff_type
        self.session = requests.Session()

    def make_request(self, url):
        headers = {"User-Agent" : "pyosmium/{}".format(version.pyosmium_release)}
        return urlrequest.Request(url, headers=headers)

    def open_url(self, url):
        """ Download a resource from the given URL and return a byte sequence
            of the content.

            This method has no support for cookies or any special authentication
            methods. If you need these, you have to provide your own custom URL
            opener. The method has to return an object which supports the
            `read()` and `readline()` methods to access the content. Example::

                def my_open_url(self, url):
                    opener = urlrequest.build_opener()
                    opener.addheaders = [('X-Fancy-Header', 'important_content')]
                    return opener.open(url)

                svr = ReplicationServer()
                svr.open_url = my_open_url
        """
        headers = dict()
        for h in url.header_items():
            headers[h[0]] = h[1]
        return self.session.get(url.get_full_url(), headers=headers, stream=True)

    def collect_diffs(self, start_id, max_size=1024):
        """ Create a MergeInputReader and download diffs starting with sequence
            id `start_id` into it. `max_size`
            restricts the number of diffs that are downloaded. The download
            stops as soon as either a diff cannot be downloaded or the
            unpacked data in memory exceeds `max_size` kB.

            If some data was downloaded, returns a namedtuple with three fields:
            `id` contains the sequence id of the last downloaded diff, `reader`
            contains the MergeInputReader with the data and `newest` is a
            sequence id of the most recent diff available.

            Returns None if there was an error during download or no new
            data was available.
        """
        left_size = max_size * 1024
        current_id = start_id

        # must not read data newer than the published sequence id
        # or we might end up reading partial data
        newest = self.get_state_info()

        if newest is None or current_id > newest.sequence:
            return None

        rd = MergeInputReader()

        while left_size > 0 and current_id <= newest.sequence:
            try:
                diffdata = self.get_diff_block(current_id)
            except:
                LOG.debug("Error during diff download. Bailing out.")
                diffdata = ''
            if len(diffdata) == 0:
                if start_id == current_id:
                    return None
                break

            left_size -= rd.add_buffer(diffdata, self.diff_type)
            LOG.debug("Downloaded change %d. (%d kB available in download buffer)",
                      current_id, left_size / 1024)
            current_id += 1

        return DownloadResult(current_id - 1, rd, newest.sequence)

    def apply_diffs(self, handler, start_id, max_size=1024, idx="", simplify=True):
        """ Download diffs starting with sequence id `start_id`, merge them
            together and then apply them to handler `handler`. `max_size`
            restricts the number of diffs that are downloaded. The download
            stops as soon as either a diff cannot be downloaded or the
            unpacked data in memory exceeds `max_size` kB.

            If `idx` is set, a location cache will be created and applied to
            the way nodes. You should be aware that diff files usually do not
            contain the complete set of nodes when a way is modified. That means
            that you cannot just create a new location cache, apply it to a diff
            and expect to get complete way geometries back. Instead you need to
            do an initial data import using a persistent location cache to
            obtain a full set of node locations and then reuse this location
            cache here when applying diffs.

            Diffs may contain multiple versions of the same object when it was
            changed multiple times during the period covered by the diff. If
            `simplify` is set to False then all versions are returned. If it
            is True (the default) then only the most recent version will be
            sent to the handler.

            The function returns the sequence id of the last diff that was
            downloaded or None if the download failed completely.
        """
        diffs = self.collect_diffs(start_id, max_size)

        if diffs is None:
            return None

        diffs.reader.apply(handler, idx=idx, simplify=simplify)

        return diffs.id

    def apply_diffs_to_file(self, infile, outfile, start_id, max_size=1024,
                            set_replication_header=True, extra_headers=None,
                            outformat=None):
        """ Download diffs starting with sequence id `start_id`, merge them
            with the data from the OSM file named `infile` and write the result
            into a file with the name `outfile`. The output file must not yet
            exist.

            `max_size` restricts the number of diffs that are downloaded. The
            download stops as soon as either a diff cannot be downloaded or the
            unpacked data in memory exceeds `max_size` kB.

            If `set_replication_header` is true then the URL of the replication
            server and the sequence id and timestamp of the last diff applied
            will be written into the `writer`. Note that this currently works
            only for the PBF format.

            `extra_headers` is a dict with additional header fields to be set.
            Most notably, the 'generator' can be set this way.

            `outformat` sets the format of the output file. If None, the format
            is determined from the file name.

            The function returns a tuple of last downloaded sequence id and
    
        newest available sequence id if new data has been written or None
            if no data was available or the download failed completely.
        """
        diffs = self.collect_diffs(start_id, max_size)

        if diffs is None:
            return None

        reader = oio.Reader(infile)
        has_history = reader.he
ader().has_multiple_object_versions

        h = oio.Header()
        h.has_multiple_object_versions = has_history
        if set_replication_header:
            h.set("osmosis_replication_base_url", self.baseurl)
            h.set("osmosis_replication_sequence_number", str(diffs.id))
            info = self.get_state_info(diffs.id)
            if info is not None:
                h.set("osmosis_replication_timestamp", info.timestamp.strftime("%Y-%m-%dT%H:%M:%SZ"))
        if extra_headers is not None:
            for k, v in extra_headers.items():
                h.set(k, v)

        if outformat is None:
            of = oio.File(outfile)
        else:
            of = oio.File(outfile, outformat)

        of.has_multiple_object_versions = has_history
        writer = oio.Writer(of, h)

        LOG.debug("Merging changes into OSM file.")

        diffs.reader.apply_to_reader(reader, writer, has_history)

        reader.close()
        writer.close()

      