import bson
import datetime
import pydantic

from storage.mongo import database
from utils import time
from bson import codec_options


collection = database.get_collection("tasks", codec_options=codec_options.CodecOptions(tz_aware=True))


class Task(pydantic.BaseModel):
    queue: str
    payload: dict
    put_at: pydantic.AwareDatetime
    available_from: pydantic.AwareDatetime
    seconds_to_execute: int
    _id: bson.ObjectId|None = None
    taken_at: pydantic.AwareDatetime|None = None
    attempts: int = 0


async def add_task(task: Task) -> str:
    result = await collection.insert_one(task.model_dump())
    return result.inserted_id


async def take_task(queue: str) -> Task|None:
    now = time.now()
    async for raw_task in collection.find({'queue': queue, 'available_from': {'$lte': now}}):
        task = Task.model_validate(raw_task)
        task._id = raw_task['_id']
        if not task.taken_at or (task.taken_at + datetime.timedelta(seconds=task.seconds_to_execute)) < now:
            await collection.update_one({'_id': task._id}, {'$set': {'taken_at': now, 'attempts': task.attempts + 1}})
            return task
    return None


async def finish_task(id: bson.ObjectId) -> bool:
    result = await collection.delete_one({'_id': id})
    return result