from config import logger, states, tasks


def save(
    doc_id: str,
    *,
    task: str | None = None,
    workspace: str | None = None,
    filetime: str | None = None,
    config_mtime: float | 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 config_mtime is not None:
        updates["config_mtime"] = config_mtime

    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 fields.
    If the document doesn't exist, all fields 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,
        "config_mtime": doc.get("config_mtime") if doc else None,
    }


def sync_desktop_tasks(work_desktop_tasks: dict):
    """
    Sync work_desktop_tasks from config file to state
    """
    update_dict = {str(k): v for k, v in work_desktop_tasks.items()}
    states.update_one(
        {"_id": "work_desktop_tasks"},
        {"$set": update_dict},
        upsert=True,
    )


def retrieve_desktop_state():
    """
    Get work_desktop_tasks mapping from state
    """
    doc = states.find_one({"_id": "work_desktop_tasks"})
    if not doc:
        return {}
    # Convert string keys to int and exclude _id
    return {int(k): v for k, v in doc.items() if k != "_id"}