from config import logger, states, tasks


def save(
    doc_id: str,
    *,
    task: str | None = None,
    workspace: str | None = None,
    filetime: str | None = None,
) -> None:
    """
    Upsert a document with _id=doc_id, setting any of the provided fields.
    Leave fields you don’t pass unchanged.
    """
    updates: dict = {}
    if task is not None:
        updates["task"] = task
    if workspace is not None:
        updates["workspace"] = workspace
    if filetime is not None:
        updates["filetime"] = filetime

    if updates:
        states.update_one(
            {"_id": doc_id},
            {"$set": updates},
            upsert=True,
        )


def retrieve(doc_id: str) -> dict[str, str | None]:
    """
    Fetches the document with _id=doc_id and returns its 'task' and 'workspace'.
    If the document doesn’t exist, both will be None.
    """
    doc = states.find_one({"_id": doc_id})
    return {
        "task": doc.get("task") if doc else None,
        "workspace": doc.get("workspace") if doc else None,
        "filetime": doc.get("filetime") if doc else None,
    }


# just
def update_work_state(work: str, task_id: str):
    """
    update work state
    """
    states.update_one({"_id": "work"}, {"$set": {work: task_id}})


def init_work_state(wd: dict):
    """
    init work states with default values
    """
    if not states.find_one({"_id": "work"}):
        states.insert_one(
            {
                "_id": "work",
                **{
                    wd[k]: tasks.find_one({"path": f"work/{wd[k]}"})["task_id"]
                    for k in wd
                },
            }
        )


def retrieve_work_state():
    return states.find_one({"_id": "work"})