import contextlib
import tempfile

import boto3
import botocore.exceptions
from flask import request, send_file

from .base import BaseStorage

class Storage(BaseStorage):
    def __init__(self, app):
        super().__init__(app)
        self.client = boto3.resource('s3', **self.app.config['S3_CONFIG'])

    def _get_object_key(self, file_hash, user_id):
        return '{}_{}'.format(file_hash, user_id)

    def get_object_key(self, f, thumb=False):
        key = self._get_object_key(f.hash, f.user_id if f.user_id else 0)
        if thumb:
            key += '_thumb'
        return key

    def store_file(self, uploaded_file, file_hash, user, ip):
        bucket = self.client.Bucket(self.app.config['S3_BUCKET'])
        key = self._get_object_key(file_hash, user.id if user else 0)
        obj = bucket.upload_fileobj(Fileobj=uploaded_file.stream, Key=key)
        size = uploaded_file.content_length
        if not size:
            obj = self.client.ObjectSummary(self.app.config['S3_BUCKET'], key)
            size = obj.size
        try:
            return self.add_file(file_hash, uploaded_file.filename, size, user, ip)
        except:
            obj.delete()
            raise

    def get_file(self, f, thumb=False):
        key = self.get_object_key(f, thumb=thumb)
        if thumb:
            bucket = self.app.config['S3_THUMB_BUCKET']
        else:
            bucket = self.app.config['S3_BUCKET']
        obj = self.client.Object(bucket, key)
        kwargs = {}
        if 'Range' in request.headers:
            kwargs['Range'] = request.headers['Range']
        data = obj.get(**kwargs)
        rv = send_file(data['Body'], attachment_filename=f.filename)
        rv.headers['Content-Length'] = data['ContentLength']
        rv.headers['Accept-Ranges'] = data['AcceptRanges']
        if 'ContentRange' in data:
            rv.headers['Content-Range'] = data['ContentRange']
            rv.status_code = 206
        return rv

    def delete_file(self, f):
        obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f))
        obj.delete()
        obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f, thumb=True))
        obj.delete()

    @contextlib.contextmanager
    def temp_file(self, f):
        obj = self.client.Object(self.app.config['S3_BUCKET'], self.get_object_key(f))
        with tempfile.NamedTemporaryFile() as f:
            obj.download_fileobj(f)
            f.seek(0)
            yield f

    def get_thumbnail(self, f):
        try:
            return self.get_file(f, thumb=True)
        except botocore.exceptions.ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                # If thumbnail does not exist, just return None.
                return
            raise

    def store_thumbnail(self, f, stream):
        bucket = self.client.Bucket(self.app.config['S3_THUMB_BUCKET'])
        key = self.get_object_key(f, thumb=True)
        obj = bucket.upload_fileobj(Fileobj=stream, Key=key)