# This file is part of the MapProxy project.
# Copyright (C) 2016 Omniscale <http://omniscale.de>
#
# Licensed 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 calendar
import hashlib
import sys
import threading
from typing import Optional

from mapproxy.cache.tile import TileCollection
from mapproxy.cache.tile import Tile
from mapproxy.image import ImageResult
from mapproxy.cache import path
from mapproxy.cache.base import tile_buffer, TileCacheBase
from mapproxy.util import async_
from mapproxy.util.py import reraise_exception
from urllib import request as urllib2

from mapproxy.util.coverage import Coverage

try:
    import boto3
    import botocore
except ImportError:
    boto3 = None


import logging
log = logging.getLogger('mapproxy.cache.s3')


_s3_sessions_cache = threading.local()


def s3_session(profile_name=None):
    if not hasattr(_s3_sessions_cache, 'sessions'):
        _s3_sessions_cache.sessions = {}
    if profile_name not in _s3_sessions_cache.sessions:
        _s3_sessions_cache.sessions[profile_name] = boto3.session.Session(profile_name=profile_name)
    return _s3_sessions_cache.sessions[profile_name]


class S3ConnectionError(Exception):
    pass


class S3Cache(TileCacheBase):

    def __init__(self, base_path, file_ext, directory_layout='tms',
                 bucket_name='mapproxy', profile_name=None, region_name=None, endpoint_url=None,
                 _concurrent_writer=4, access_control_list=None, coverage: Optional[Coverage] = None,
                 use_http_get=False):
        super().__init__(coverage)
        md5 = hashlib.new('md5', base_path.encode('utf-8') + bucket_name.encode('utf-8'), usedforsecurity=False)
        self.lock_cache_id = md5.hexdigest()
        self.bucket_name = bucket_name
        self.profile_name = profile_name
        self.region_name = region_name
        self.endpoint_url = endpoint_url
        self.access_control_list = access_control_list
        self.use_http_get = use_http_get

        try:
            self.bucket = self.conn().head_bucket(Bucket=bucket_name)
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == '404':
                raise S3ConnectionError('No such bucket: %s' % bucket_name)
            elif e.response['Error']['Code'] == '403':
                raise S3ConnectionError('Access denied. Check your credentials')
            else:
                raise reraise_exception(
                    S3ConnectionError('Unknown error: %s' % e),
                    sys.exc_info(),
                )

        self.base_path = base_path
        self.file_ext = file_ext
        self._concurrent_writer = _concurrent_writer

        self._tile_location, _ = path.location_funcs(layout=directory_layout)

    def get_bucket_url(self, tile):
        return f"https://{self.bucket_name}.s3.{self.region_name}.amazonaws.com/{self.tile_key(tile)}"

    def tile_key(self, tile):
        return self._tile_location(tile, self.base_path, self.file_ext).lstrip('/')

    def conn(self):
        if boto3 is None:
            raise ImportError("S3 Cache requires 'boto3' package.")

        try:
            return s3_session(self.profile_name).client(
                "s3", region_name=self.region_name, endpoint_url=self.endpoint_url)
        except Exception as e:
            raise S3ConnectionError('Error during connection %s' % e)

    def load_tile_metadata(self, tile: Tile, dimensions=None):
        if tile.timestamp:
            return
        self.is_cached(tile, dimensions=dimensions)

    def _set_metadata(self, response, tile):
        if 'LastModified' in response:
            tile.timestamp = calendar.timegm(response['LastModified'].timetuple())
        if 'ContentLength' in response:
            tile.size = response['ContentLength']

    def is_cached(self, tile: Tile, dimensions=None) -> bool:
        if tile.is_missing():
            if self.use_http_get:
                try:
                    req = urllib2.Request(self.get_bucket_url(tile))
                    response = urllib2.urlopen(req)
                    self._set_metadata(response.info(), tile)
                except urllib2.HTTPError as e:
                    if e.code == 403:
                        return False
                    raise
            else:
                key = self.tile_key(tile)
                try:
                    r = self.conn().head_object(Bucket=self.bucket_name, Key=key)
                    self._set_metadata(r, tile)
                except botocore.exceptions.ClientError as e:
                    if e.response['Error']['Code'] in ('404', 'NoSuchKey'):
                        return False
                    raise

        return True

    def load_tiles(self, tiles: TileCollection, with_metadata=True, dimensions=None) -> bool:
        p = async_.Pool(min(4, len(tiles)))
        return all(p.map(self.load_tile, tiles))

    def load_tile(self, tile: Tile, with_metadata=True, dimensions=None) -> bool:
        if not tile.is_missing():
            return True

        key = self.tile_key(tile)
        log.debug('S3:load_tile, key: %s' % key)

        if self.use_http_get:
            try:
                req = urllib2.Request(self.get_bucket_url(tile))
                tile.image_result = ImageResult(urllib2.urlopen(req))
            except urllib2.HTTPError as e:
                if e.code == 403:
                    return False
                raise
        else:
            try:
                r = self.conn().get_object(Bucket=self.bucket_name, Key=key)
                self._set_metadata(r, tile)
                tile.image_result = ImageResult(r['Body'])
            except botocore.exceptions.ClientError as e:
                # moto get_object can return Error wrapped in Errors...
                error = e.response.get('Errors', e.response)['Error']
                if error['Code'] in ('404', 'NoSuchKey'):
                    return False
                raise

        return True

    def remove_tile(self, tile, dimensions=None):
        key = self.tile_key(tile)
        log.debug('remove_tile, key: %s' % key)
        self.conn().delete_object(Bucket=self.bucket_name, Key=key)

    def store_tiles(self, tiles, dimensions=None):
        p = async_.Pool(min(self._concurrent_writer, len(tiles)))
        p.map(self.store_tile, tiles)

    def store_tile(self, tile, dimensions=None):
        if tile.stored:
            return

        key = self.tile_key(tile)
        log.debug('S3: store_tile, key: %s' % key)

        extra_args = {}
        if self.file_ext in ('jpeg', 'png'):
            extra_args['ContentType'] = 'image/' + self.file_ext
        if self.access_control_list:
            extra_args['ACL'] = self.access_control_list
        with tile_buffer(tile) as buf:
            self.conn().upload_fileobj(
                NopCloser(buf),  # upload_fileobj closes buf, wrap in NopCloser
                self.bucket_name,
                key,
                ExtraArgs=extra_args)


class NopCloser(object):
    def __init__(self, wrapped):
        self.wrapped = wrapped

    def close(self):
        pass

    def __getattr__(self, name):
        return getattr(self.wrapped, name)
