# @Time    : 2023/8/11 20:06
# @Author  : Lan
# @File    : storage.py
# @Software: PyCharm
import asyncio
import hashlib
import time

import aioboto3
from fastapi import UploadFile
from core.settings import data_root, settings
from apps.base.models import FileCodes


class SystemFileStorage:
    def __init__(self):
        self.chunk_size = 256 * 1024
        self.root_path = data_root
        self.token = '123456'

    def _save(self, file, save_path):
        with open(save_path, 'wb') as f:
            chunk = file.read(self.chunk_size)
            while chunk:
                f.write(chunk)
                chunk = file.read(self.chunk_size)

    async def save_file(self, file: UploadFile, save_path: str):
        save_path = self.root_path / save_path
        if not save_path.parent.exists():
            save_path.parent.mkdir(parents=True)
        await asyncio.to_thread(self._save, file.file, save_path)

    async def delete_file(self, file_code: FileCodes):
        save_path = self.root_path / await file_code.get_file_path()
        if save_path.exists():
            save_path.unlink()

    async def get_select_token(self, code):
        return hashlib.sha256(f"{code}{int(time.time() / 1000)}000{self.token}".encode()).hexdigest()

    async def get_file_url(self, file_code: FileCodes):
        return f'/share/download?key={await self.get_select_token(file_code.code)}&code={file_code.code}'


class S3FileStorage:
    def __init__(self):
        self.access_key_id = settings.s3_access_key_id
        self.secret_access_key = settings.s3_secret_access_key
        self.bucket_name = settings.s3_bucket_name
        self.endpoint_url = settings.s3_endpoint_url
        self.session = aioboto3.Session(
            aws_access_key_id=self.access_key_id, aws_secret_access_key=self.secret_access_key
        )

    async def save_file(self, file: UploadFile, save_path: str):
        async with self.session.client("s3", endpoint_url=self.endpoint_url) as s3:
            await s3.put_object(Bucket=self.bucket_name, Key=save_path, Body=await file.read(), ContentType=file.content_type)

    async def delete_file(self, file_code: FileCodes):
        async with self.session.client("s3", endpoint_url=self.endpoint_url) as s3:
            await s3.delete_object(Bucket=self.bucket_name, Key=await file_code.get_file_path())

    async def get_file_url(self, file_code: FileCodes):
        if file_code.prefix == '文本分享':
            return file_code.text
        async with self.session.client("s3", endpoint_url=self.endpoint_url) as s3:
            result = await s3.generate_presigned_url('get_object', Params={'Bucket': self.bucket_name, 'Key': await file_code.get_file_path()}, ExpiresIn=3600)
            return result


class FileStorageTemplate:
    def __init__(self):
        ...

    async def save_file(self, file: UploadFile, save_path: str):
        ...

    async def delete_file(self, file_code: FileCodes):
        ...

    async def get_file_url(self, file_code: FileCodes):
        ...


storages = {
    'local': SystemFileStorage,
    's3': S3FileStorage
}
file_storage = storages[settings.file_storage]()