from asyncio import Lock, to_thread
from pathlib import Path

from pydantic import BaseModel
from fastapi import BackgroundTasks
from tcvectordb.model.collection_view import SplitterProcess, ParsingProcess

from settings import settings
from main import app
from services import cos, collection_view

class FileNotify(BaseModel):
	key: str

lock = Lock()

def upload_to_vdb(local_path: str):
	collection_view.load_and_split_text(
		local_file_path=local_path,
		metadata={},
		splitter_process=SplitterProcess(
			append_title_to_chunk=False,
			append_keywords_to_chunk=True,
			chunk_splitter=None,
		),
		parsing_process=ParsingProcess(
			parsing_type='VisionModelParsing',
		),
	)

async def process_uploaded_file(notify: FileNotify):
	async with lock:
		if not await to_thread(cos.object_exists, settings.tencent_cloud.bucket, notify.key):
			return

		local_path = f'/tmp/{notify.key}'
		await to_thread(cos.download_file, settings.tencent_cloud.bucket, notify.key, local_path)
		await to_thread(upload_to_vdb, local_path)

		Path(local_path).unlink()

@app.post('/notify-upload-complete')
def notify_upload_complete(notify: FileNotify, background_tasks: BackgroundTasks):
	background_tasks.add_task(process_uploaded_file, notify)
	return