import json
import os
import shutil
from fastapi import APIRouter, Depends, Request, UploadFile, File, Form
from sqlalchemy.ext.asyncio import AsyncSession
from config.get_db import get_db
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_admin.entity.vo.aichat_vo import *
from module_admin.service.login_service import LoginService
from module_admin.service.aichat_service import AiChatService
from utils.log_util import logger
from utils.response_util import ResponseUtil
from config.env import MinioConfig
from datetime import datetime
from utils.common_util import bytes2file_response
from utils.minio_util import Bucket

aichatController = APIRouter(prefix='/aichat', dependencies=[Depends(LoginService.get_current_user)])


@aichatController.get("/session/list/{sessionId}")
async def get_chat_session_list(request: Request,
                                sessionId: str,
                                query_db: AsyncSession = Depends(get_db),
                                current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    ai_session_list_result = await AiChatService.get_ai_session_list_services(query_db, sessionId, current_user)
    logger.info('获取成功')
    return ResponseUtil.success(data=ai_session_list_result)


@aichatController.get("/chat/list/{sessionId}")
async def get_chat_list(request: Request, sessionId: str, query_db: AsyncSession = Depends(get_db),
                        current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    ai_chat_list_result = await AiChatService.get_ai_chat_list_services(query_db, sessionId, current_user)
    logger.info('获取成功')
    return ResponseUtil.success(data=ai_chat_list_result)


@aichatController.post("/delete/session/{sessionId}")
async def delete_chat_session(request: Request, sessionId: str, query_db: AsyncSession = Depends(get_db),
                              current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    delete_chat_session_result = await AiChatService.delete_chat_session(query_db, sessionId, current_user)
    logger.info(delete_chat_session_result.message)
    return ResponseUtil.success(msg=delete_chat_session_result.message)


@aichatController.post("/add")
async def add_chat(request: Request, add_chat: AiChatModel, query_db: AsyncSession = Depends(get_db),
                   current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    operate_result = await AiChatService.add_chat(query_db, add_chat, current_user)
    logger.info(operate_result.message)
    return ResponseUtil.success(msg=operate_result.message)


@aichatController.post("/update")
async def update_chat(request: Request, update_chat: AiChatModel, query_db: AsyncSession = Depends(get_db),
                      current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    operate_result = await AiChatService.update_chat(query_db, update_chat)
    logger.info(operate_result.message)
    return ResponseUtil.success(msg=operate_result.message)


@aichatController.post("/upload")
async def upload_file(request: Request, sessionId: str = Form(), file: UploadFile = File(...),
                      current_user: CurrentUserModel = Depends(LoginService.get_current_user)):
    file_extension = os.path.splitext(file.filename)[1]  # 文件后缀
    file_prefix = os.path.splitext(file.filename)[0]  # 文件前缀
    bucket = Bucket(
        minio_address=MinioConfig.minio_address,
        minio_admin=MinioConfig.minio_admin,
        minio_password=MinioConfig.minio_password
    )
    file_id = datetime.now().strftime('%Y%m%d%H%M%S')
    file_location = file_prefix + "_" + file_id + file_extension
    with open(file_location, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # Remove the temporary file after upload (optional)
    def remove_file(file_path):
        os.remove(file_path)

    # Upload to MinIO
    bucket.upload_file_to_bucket(current_user.user.user_name, '/' + sessionId + '/' + file_location, file_location)

    # Clean up
    remove_file(file_location)
    return ResponseUtil.success(data={"is_success": True, "message": "上传成功", "file": file_location,
                                      "bucket": current_user.user.user_name},
                                msg="上传成功")


@aichatController.post("/file/download")
async def download_file(request: Request, download_file: DownloadFile = Form()):
    file_name = download_file.file
    bucket_name = download_file.bucket
    sessionId = download_file.sessionId
    bucket = Bucket(
        minio_address=MinioConfig.minio_address,
        minio_admin=MinioConfig.minio_admin,
        minio_password=MinioConfig.minio_password
    )

    # Remove the temporary file after upload (optional)
    def remove_file(file_path):
        os.remove(file_path)

    # Upload to MinIO
    bucket.download_file_from_bucket(bucket_name, '/' + sessionId + '/' + file_name, file_name)
    with open(file_name, "rb") as buffer:
        file_data = buffer.read()
    # Clean up
    remove_file(file_name)
    return ResponseUtil.streaming(data=bytes2file_response(file_data))