# -*- coding:utf-8 -*-
"""
@Author  : xuxingchen
@Contact : xuxingchen@sinochem.com
@Desc    : api后端服务程序启动入口
"""
import argparse
import os

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import aiofiles
import uvicorn

from config import VERSION, PREFIX_PATH, FACES_DIR_PATH, IMAGE_SERVER_PORT
from utils import logger
from routers import edge_simulation_api


class ImageServer:
    def __init__(self, app_args):
        self.port = app_args.port
        self.app = None
        self.prepare()
        self.server_main()

    def prepare(self):
        """准备处理，一些适配和配置加载工作"""
        # 清理图片库，移除不存在数据库中的图片
        logger.Logger.init("人脸图片库执行初始化校验 ...")
        for filename in os.listdir(FACES_DIR_PATH):
            if filename.startswith("_"):
                os.remove(os.path.join(FACES_DIR_PATH, filename))
        logger.Logger.init("人脸图片库完成初始化校验 ✅")

        self.app = FastAPI(
            title="Image Server",
            description="图片后端服务",
            version=VERSION,
            docs_url=f"{PREFIX_PATH}/docs"  # 设为None则会关闭
        )

        # self.app.mount(f"{PREFIX_PATH}/faces", StaticFiles(directory=FACES_DIR_PATH), name="static")
        self.app.include_router(edge_simulation_api.router, tags=["边缘端请求接口"])

        @self.app.get(f"{PREFIX_PATH}/faces/{{image_filename}}", summary="返回指定图片")
        async def get_image_data(request: Request, image_filename: str):
            """返回指定图片"""
            image_path = f"{FACES_DIR_PATH}/{image_filename}"
            if not os.path.exists(image_path):
                raise HTTPException(status_code=404, detail="Image not found")

            async def image_streamer(file_path: str):
                async with aiofiles.open(file_path, mode='rb') as f:
                    while chunk := await f.read(1024):  # 每次读取 1024 字节
                        yield chunk

            logger.Logger.info(f"{PREFIX_PATH}/faces -> {image_path}")
            return StreamingResponse(image_streamer(image_path), media_type="image/jpeg")

    def server_main(self):
        """主服务程序"""
        uvicorn.run(self.app, host="0.0.0.0", port=self.port, access_log=False)


def main(port: str = IMAGE_SERVER_PORT):
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument(
        "-p",
        "--port",
        type=int,
        default=port,
        help=f"Port of server, default {logger.new_dc(port)}",
    )
    app_args = parser.parse_args()
    ImageServer(app_args)


if __name__ == "__main__":
    main()