from __future__ import annotations

from datetime import datetime
from typing import Optional, Tuple

import dataiku
from dataikuapi.dss.managedfolder import DSSManagedFolder
from dataikuapi.dss.project import DSSProject
from dataikuapi.dss.recipe import DSSRecipe
from dataikuapi.utils import DataikuException

from backend.config import (
    get_config,
    get_default_embedding_llm,
    get_my_agents_vision_llm,
)
from backend.constants import DRAFT_ZONE, PUBLISHED_ZONE, CloudConnections
from backend.schemas import schemas
from backend.schemas.schemas import SemanticDatasourceInfo
from backend.utils.api_clients import AH_GLOBAL_API_CLIENT
from backend.utils.docs_utils import (
    is_sharepoint_file_doc,
    sharepoint_doc_replacement_key,
    sharepoint_doc_source_folder_id,
    sharepoint_doc_source_path,
    with_sharepoint_source_folder_id,
)
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import (
    _normalize_agent_tool_id,
    build_sharepoint_managed_folder_name,
    compute_custom_tool_name,
    create_aug_llm,
    create_custom_tool,
    create_embed_doc_recipe,
    create_kb_tool,
    create_or_update_merge_folder_recipe,
    create_sync_scenario,
    create_visual_agent,
    edit_tool_settings,
    ensure_embed_doc_recipe_input,
    ensure_quick_agent_support,
    get_bound_semantic_tool_ids_in_zone,
    get_kb_in_zone,
    get_kb_tool_in_zone,
    get_managed_folder_in_zone,
    get_outputs_from_embed_docs_recipe,
    get_recipe_in_zone,
    get_semantic_tool_in_zone,
    get_ua_project,
    get_visual_agent_in_zone,
    get_zone_id_by_name,
    grant_project_access_to_user,
    set_additional_tool_description,
    set_embed_doc_recipe_vlm,
    update_visual_agent,
)
from backend.utils.semantic_models_utils import (
    DEFAULT_USER_SQL_EXECUTION,
    SEMANTIC_MODEL_TOOLS_META,
    SemanticModelToolID,
)
from backend.utils.utils import get_doc_name

logger = get_logger(__name__)


def _sync_kb_tool_description(kb_tool, knowledge_description: str, project_key: str, zone_name: str) -> None:
    if not kb_tool:
        return
    tool_settings = kb_tool.get_settings()
    tool_settings_raw = tool_settings.get_raw()
    tool_settings_raw["additionalDescriptionForLLM"] = knowledge_description
    tool_settings.save()
    logger.info(f"Updated kb_tool description, project {project_key}, zone {zone_name}")


def prepare_agent_tools(project: DSSProject, agent: dict, zone_name: str, with_semantic_tools: bool = True):
    docs = agent.get("documents", []) or []
    knowledge_docs = [doc for doc in docs if not (isinstance(doc, dict) and doc.get("deletePending"))]
    tools_ids = list(agent.get("tools", []))  # make sure keep it immutable
    kb_id = get_kb_in_zone(project, zone_name)
    knowledge_description = (agent.get("kb_description") or "").strip()
    kb_tool = get_kb_tool_in_zone(project, zone_name)
    semantic_datasource: SemanticDatasourceInfo = agent.get("semantic_datasource", {})

    if knowledge_docs:
        if not kb_tool:
            kb_tool = create_kb_tool(
                project=project,
                kb_id=kb_id,
                kb_desc=knowledge_description,
                maxDocuments=get_config().get("number_of_documents_to_retrieve", 4),
                zone_name=zone_name,
            )
            logger.info(
                f"create kb_tool {kb_tool.id} - project {project.project_key} - in zone {zone_name} with description '{knowledge_description}'"
            )
        else:
            _sync_kb_tool_description(kb_tool, knowledge_description, project.project_key, zone_name)
        kb_tool_id = f"{project.project_key}:{kb_tool.id}"
        tools_ids.append(kb_tool_id)
    # Keep list stable and unique
    tools_ids = list(dict.fromkeys(tools_ids))

    if semantic_datasource:
        # Retrieval of the project key and dataset ID (pattern of the ID is PROJECT_KEY.DATASET_ID)
        semantic_model_project_key, semantic_model_source_id = semantic_datasource["id"].split(".")
        semantic_datasource_type = semantic_datasource["type"]
        semantic_model_tool_id: SemanticModelToolID = SEMANTIC_MODEL_TOOLS_META[semantic_datasource_type]["tool_id"]
        semantic_model_tool_type = SEMANTIC_MODEL_TOOLS_META[semantic_datasource_type]["tool_type"]
        semantic_tool_name = compute_custom_tool_name(
            tool_id=semantic_model_tool_id, base_name=semantic_model_source_id, zone_name=zone_name
        )
        semantic_tool = get_semantic_tool_in_zone(
            project=project, tool_type=semantic_model_tool_type, source_id=semantic_model_source_id, zone_name=zone_name
        )
        sem_models_llm = agent.get("sem_models_llm", "")
        sem_models_interaction_mode = agent.get("sem_models_interaction_mode", "simple")
        agent_mode = False
        if sem_models_interaction_mode == "agentic":
            agent_mode = True

        tool_params = {}
        if semantic_model_tool_id == "semantic-model-query":
            tool_params = {
                "project_key": semantic_model_project_key,
                "semantic_model_id": semantic_model_source_id,
                "llm_id": sem_models_llm if sem_models_llm else agent.get("llmid"),
                "embedding_llm_id": get_default_embedding_llm(),
                "agent_mode": agent_mode,
                "enduser_sql_execution": DEFAULT_USER_SQL_EXECUTION,
            }
        elif semantic_model_tool_id == "dataset-sql-query":
            tool_params = {
                "project_key": semantic_model_project_key,
                "dataset_id": semantic_model_source_id,
                "llm_id": sem_models_llm if sem_models_llm else agent.get("llmid"),
                "agent_mode": agent_mode,
                "enduser_sql_execution": DEFAULT_USER_SQL_EXECUTION,
            }
        else:
            log_message = f"Get a request for a semantic tool creation with semantic_model_tool_id='{semantic_model_tool_id}' (Not supported)"
            logger.error(log_message)
            raise ValueError(log_message)

        if not semantic_tool:
            semantic_tool = create_custom_tool(
                project=project,
                tool_type=semantic_model_tool_type,
                tool_name=semantic_tool_name,
                tool_params=tool_params,
                zone_name=zone_name,
            )

        else:
            edit_tool_settings(project, semantic_tool.id, tool_params)

        additional_description = ""
        if description := semantic_datasource.get("description"):
            DESCRIPTION_PREFIX = "## Dataset description:"
            additional_description = description
            if DESCRIPTION_PREFIX not in description:
                additional_description = f"{DESCRIPTION_PREFIX} {description}\n\n"
            set_additional_tool_description(project, semantic_tool.id, additional_description)

        semantic_tool_id = f"{project.project_key}:{semantic_tool.id}"
        # The following allows to bind the semantic tool to the Agent
        if with_semantic_tools:
            tools_ids.append(semantic_tool_id)
        else:
            semantic_tool.delete()
    return tools_ids


def create_quick_agent(project: DSSProject, agent: dict, project_name: str, zone_name: str):
    logger.info(f"Bundling user agent {agent.get('id')} in zone {zone_name}")
    try:
        tools_ids = prepare_agent_tools(project, agent, zone_name)
        llm_id = agent.get("llmid")
        prompt = agent.get("system_prompt", "")
        # Create visual agent in case we have tools
        return create_visual_agent(
            project=project,
            agent_name=f"{project_name} {zone_name}",
            tools_ids=tools_ids,
            llm_id=llm_id,
            prompt=prompt,
        )
    except Exception as e:
        logger.exception(f"Failed to create visual agent in zone {zone_name}: {e}")


def create_agent_project(project_key, project_name, project_folder_id, agent, owner):
    from backend.utils.client_utils import current_request_login

    client = AH_GLOBAL_API_CLIENT
    connection = get_config().get("default_fs_connection")

    project: DSSProject = client.create_project(
        project_key=project_key, owner=owner, name=project_name, project_folder_id=project_folder_id
    )
    # project settings
    settings = project.get_settings()
    psettings = settings.get_raw()
    # Set project to archived to avoid it showing up in the project list and being editable
    psettings["projectStatus"] = "Archived"
    # Bypass AI Consumer license check for creating the Quick Agent scenario by flagging the project as Quick Agent project
    psettings["projectType"] = "AGENT_HUB_QUICK_AGENT_SUPPORT"
    settings.save()

    # Give user permissions to the project
    grant_project_access_to_user(project, current_request_login())

    llm_id = agent.get("llmid")
    vision_llm_id = get_my_agents_vision_llm()
    embedding_llm = get_default_embedding_llm()

    logger.info(
        f"Created project {project_key} for user agent, owner: {owner}, llm: {llm_id}, embedding llm: {embedding_llm}"
    )
    try:
        # Prepare flow zones
        for zone_name in [DRAFT_ZONE, PUBLISHED_ZONE]:
            logger.info(f"Creating managed folder with connection {connection} in zone {zone_name}")
            # Prepare folders, recipes & ra models
            folder = project.create_managed_folder(f"documents_{zone_name}", connection_name=connection)
            merged_folder = project.create_managed_folder(
                _merged_docs_folder_name(zone_name), connection_name=connection
            )
            merge_recipe = create_or_update_merge_folder_recipe(
                project=project,
                zone_name=zone_name,
                input_folder_ids=[folder.id],
                output_folder_id=merged_folder.id,
                recipe_name=_merge_docs_recipe_name(zone_name),
            )
            r = create_embed_doc_recipe(
                project=project,
                zone=zone_name,
                folder_id=merged_folder.id,
                llm_id=vision_llm_id,
                embedding_llm=embedding_llm,
            )
            # Finding the KB ID
            new_kb_id = r.get_settings().data["recipe"]["outputs"]["knowledge_bank"]["items"][0]["ref"]
            ra_model = create_aug_llm(project=project, name=f"ra_model_{zone_name}", kb_id=new_kb_id, llm_id=llm_id)
            # We need to retrieve the DSSSavedModel for the ra_model
            sm = project.get_saved_model(ra_model.id)
            zone = project.get_flow().create_zone(zone_name)
            zone.add_items([folder, merged_folder, merge_recipe, r, sm])
            va = create_quick_agent(project, agent, project_name, zone_name)
            if va:
                zone.add_item(va)
        meta = project.get_metadata()
        meta["tags"] = ["user_agent", current_request_login()]
        project.set_metadata(meta)
        return project
    except Exception as e:
        logger.exception("Failed to create all elements in the project {e}")
        logger.info("Deleting project")
        project.delete()
        return None


def get_ua_project_folder(agent_id: str) -> Tuple[DSSProject, DSSManagedFolder]:
    project = get_ua_project(agent_id)
    documents_folder = get_managed_folder_in_zone(project, f"documents_{DRAFT_ZONE}", DRAFT_ZONE)
    if not documents_folder:
        raise Exception("Draft documents folder not found in template")
    return project, documents_folder


def _merged_docs_folder_name(zone_name: str) -> str:
    return f"merged_documents_{zone_name}"


def _merge_docs_recipe_name(zone_name: str) -> str:
    return f"merge_documents_{zone_name}"


def _sharepoint_selected_files_folder_name(zone_name: str) -> str:
    return f"sharepoint_selected_files_{zone_name}"


def _sharepoint_selected_paths_dataset_name(zone_name: str) -> str:
    return f"sharepoint_selected_paths_{zone_name}"


def _sharepoint_selected_files_recipe_name(zone_name: str) -> str:
    return f"merge_sharepoint_selected_files_{zone_name}"


def _cloud_docs_folder_name(connection_type: str, zone_name: str) -> str:
    if connection_type == CloudConnections.SHAREPOINT_ONLINE.value:
        return f"sharepoint_documents_{zone_name}"
    return f"{connection_type}_documents_{zone_name}".lower()


def _sp_folder_name_for_doc(doc: dict, zone_name: str) -> Optional[str]:
    if not isinstance(doc, dict):
        return None
    connection = doc.get("connection") or ""
    site_id = doc.get("siteId") or doc.get("site") or ""
    drive_id = doc.get("driveId") or doc.get("drive") or ""
    path = sharepoint_doc_source_path(doc) or "/"
    if not connection:
        return None
    return build_sharepoint_managed_folder_name(zone_name, connection, site_id, drive_id, path)


def _sharepoint_file_docs(docs: list[dict]) -> list[dict]:
    return [
        doc
        for doc in docs or []
        if is_sharepoint_file_doc(doc) and not doc.get("deletePending") and sharepoint_doc_source_folder_id(doc)
    ]


def _unique_by_id(items: list) -> list:
    unique_items = []
    seen = set()
    for item in items:
        item_id = getattr(item, "id", item)
        if not item_id or item_id in seen:
            continue
        seen.add(item_id)
        unique_items.append(item)
    return unique_items


def _matches_connection_doc(doc, connection_type: str, doc_name: Optional[str] = None) -> bool:
    if not isinstance(doc, dict) or doc.get("connectionType") != connection_type:
        return False
    if doc_name is None:
        return True
    return get_doc_name(doc).lower() == doc_name.lower()


def _delete_managed_folder_object(folder: Optional[DSSManagedFolder], zone_name: str, folder_label: str) -> None:
    if not folder:
        return

    try:
        # Remove the managed-folder object from DSS flow metadata.
        # The backing cloud storage is not cleared here.
        folder.delete()
        logger.info("Deleted managed folder object '%s' in zone '%s'", folder_label, zone_name)
    except Exception as e:
        logger.warning("Failed to delete managed folder object '%s' in zone '%s': %s", folder_label, zone_name, e)


def _resolve_cloud_doc_folder(project: DSSProject, doc: dict, zone_name: str) -> Optional[DSSManagedFolder]:
    folder = None
    if doc.get("folderId"):
        try:
            folder = project.get_managed_folder(doc["folderId"])
        except Exception as e:
            logger.warning(
                "Could not resolve cloud folder '%s' in zone '%s': %s",
                doc.get("folderId"),
                zone_name,
                e,
            )
    if folder:
        return folder

    connection_type = doc.get("connectionType")
    if connection_type == CloudConnections.SHAREPOINT_ONLINE.value:
        folder_name = _sp_folder_name_for_doc(doc, zone_name)
        if folder_name:
            folder = get_managed_folder_in_zone(project, folder_name, zone_name)
            if folder:
                return folder

    if connection_type:
        return get_managed_folder_in_zone(project, _cloud_docs_folder_name(connection_type, zone_name), zone_name)
    return None


def _cloud_doc_folders(project: DSSProject, docs: list[dict], zone_name: str) -> list[DSSManagedFolder]:
    folders = []
    for doc in docs:
        if not isinstance(doc, dict) or doc.get("deletePending") or not doc.get("connectionType"):
            continue
        if is_sharepoint_file_doc(doc):
            continue
        folder = _resolve_cloud_doc_folder(project, doc, zone_name)
        if folder:
            folders.append(folder)
    return _unique_by_id(folders)


def _has_cloud_documents(docs: list[dict], connection_type: Optional[str] = None) -> bool:
    return any(
        isinstance(doc, dict)
        and doc.get("connectionType")
        and (connection_type is None or doc.get("connectionType") == connection_type)
        and not doc.get("deletePending")
        for doc in docs or []
    )


def _dataset_item_name(item) -> str:
    if isinstance(item, dict):
        return item.get("name") or item.get("datasetName") or item.get("id") or ""
    return getattr(item, "name", "") or getattr(item, "id", "")


def _dataset_exists(project: DSSProject, dataset_name: str) -> bool:
    try:
        return any(_dataset_item_name(item) == dataset_name for item in project.list_datasets())
    except Exception as e:
        logger.warning("Could not list datasets in project '%s': %s", project.project_key, e)
        return False


def _ensure_dataset_in_zone(project: DSSProject, dataset, zone_name: str) -> None:
    zone_id = get_zone_id_by_name(project, zone_name)
    if not zone_id:
        return
    try:
        project.get_flow().get_zone(zone_id).add_item(dataset)
    except Exception as e:
        logger.info("Could not add dataset '%s' to zone '%s': %s", getattr(dataset, "name", dataset), zone_name, e)


def _ensure_sharepoint_selected_paths_dataset(project: DSSProject, zone_name: str):
    dataset_name = _sharepoint_selected_paths_dataset_name(zone_name)
    if not _dataset_exists(project, dataset_name):
        project.create_dataset(dataset_name, "Inline", params={})
    dataset = project.get_dataset(dataset_name)
    dataset.set_schema(
        {
            "columns": [
                {"name": "folder_id", "type": "string"},
                {"name": "path", "type": "string"},
            ]
        }
    )
    _ensure_dataset_in_zone(project, dataset, zone_name)
    return dataset


def _write_sharepoint_selected_paths_dataset(project: DSSProject, dataset_name: str, rows: list[dict]) -> None:
    dataset = dataiku.Dataset(dataset_name, project_key=project.project_key, ignore_flow=True)

    with dataset.get_writer() as writer:
        for row in rows:
            writer.write_row_dict(
                {
                    "folder_id": row.get("folder_id", ""),
                    "path": row.get("path", ""),
                }
            )


def _ensure_sharepoint_selected_files_folder(project: DSSProject, zone_name: str) -> DSSManagedFolder:
    folder_name = _sharepoint_selected_files_folder_name(zone_name)
    selected_files_folder = get_managed_folder_in_zone(project, folder_name, zone_name)
    if selected_files_folder:
        return selected_files_folder

    selected_files_folder = project.create_managed_folder(
        folder_name,
        connection_name=get_config().get("default_fs_connection"),
    )
    zone_id = get_zone_id_by_name(project, zone_name)
    if zone_id:
        project.get_flow().get_zone(zone_id).add_item(selected_files_folder)
    logger.info("Created SharePoint selected files folder '%s'", folder_name)
    return selected_files_folder


def _selected_file_rows(file_docs: list[dict]) -> list[dict]:
    rows = []
    seen = set()
    for doc in file_docs:
        source_folder_id = sharepoint_doc_source_folder_id(doc)
        relative_path = doc.get("relativePath") or doc.get("path")
        if not source_folder_id or not relative_path:
            continue
        if not relative_path.startswith("/"):
            relative_path = f"/{relative_path}"
        key = (source_folder_id, relative_path)
        if key in seen:
            continue
        seen.add(key)
        rows.append({"folder_id": source_folder_id, "path": relative_path})
    return rows


def ensure_sharepoint_selected_files_assets(
    project: DSSProject,
    agent: dict,
    zone_name: str,
) -> Optional[DSSManagedFolder]:
    file_docs = _sharepoint_file_docs(agent.get("documents") or [])
    if not file_docs:
        return None

    selected_files_folder = _ensure_sharepoint_selected_files_folder(project, zone_name)
    _ensure_sharepoint_selected_paths_dataset(project, zone_name)
    selected_paths_dataset_name = _sharepoint_selected_paths_dataset_name(zone_name)
    rows = _selected_file_rows(file_docs)
    _write_sharepoint_selected_paths_dataset(project, selected_paths_dataset_name, rows)

    input_folder_ids = _unique_by_id(
        [sharepoint_doc_source_folder_id(doc) for doc in file_docs if sharepoint_doc_source_folder_id(doc)]
    )
    create_or_update_merge_folder_recipe(
        project=project,
        zone_name=zone_name,
        input_folder_ids=input_folder_ids,
        output_folder_id=selected_files_folder.id,
        recipe_name=_sharepoint_selected_files_recipe_name(zone_name),
        explicit_paths_dataset_id=selected_paths_dataset_name,
    )

    logger.info(
        "Ensured SharePoint selected files assets in zone '%s' with %s source folders and %s selected paths",
        zone_name,
        len(set(input_folder_ids)),
        len(rows),
    )
    return selected_files_folder


def ensure_merged_documents_assets(project: DSSProject, agent: dict, zone_name: str):
    local_folder = get_managed_folder_in_zone(project, f"documents_{zone_name}", zone_name)
    if not local_folder:
        raise Exception(f"documents_{zone_name} folder not found")

    merged_folder = get_managed_folder_in_zone(project, _merged_docs_folder_name(zone_name), zone_name)
    if not merged_folder:
        connection = get_config().get("default_fs_connection")
        merged_folder = project.create_managed_folder(_merged_docs_folder_name(zone_name), connection_name=connection)
        zone_id = get_zone_id_by_name(project, zone_name)
        if zone_id:
            project.get_flow().get_zone(zone_id).add_item(merged_folder)
        logger.info("Created merged documents folder '%s'", _merged_docs_folder_name(zone_name))

    selected_files_folder = ensure_sharepoint_selected_files_assets(project, agent, zone_name)
    cloud_folders = _cloud_doc_folders(project, agent.get("documents") or [], zone_name)
    input_folder_ids = _unique_by_id(
        [
            local_folder.id,
            *([selected_files_folder.id] if selected_files_folder else []),
            *[folder.id for folder in cloud_folders],
        ]
    )
    merged_recipe = create_or_update_merge_folder_recipe(
        project=project,
        zone_name=zone_name,
        input_folder_ids=input_folder_ids,
        output_folder_id=merged_folder.id,
        recipe_name=_merge_docs_recipe_name(zone_name),
    )
    ensure_embed_doc_recipe_input(project, zone_name, merged_folder.id)
    return merged_recipe


def ensure_ua_project(agent: dict, owner) -> DSSProject:
    if "ac_user_agent_" not in agent.get("id"):
        project_key = f"ac_user_agent_{agent['id']}"
        logger.info(f"Creating new project {project_key} owner: {agent.get('owner')}")
        logger.debug(f"ensure_ua_project, Agent details: {agent}")
        agent_hub_folder_id = get_config().get("agents_folder")
        logger.info(f"Using agent hub folder id: {agent_hub_folder_id}")
        return create_agent_project(
            project_key=project_key,
            project_name=agent.get("name", project_key),
            project_folder_id=agent_hub_folder_id,
            agent=agent,
            owner=owner,
        )
    return get_ua_project(agent.get("id"))


def upload_documents(agent: dict, files) -> None:
    """Upload documents to the draft folder."""
    project, folder = get_ua_project_folder(agent.get("id"))
    ensure_merged_documents_assets(project, agent, DRAFT_ZONE)
    vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
    if not vis_agent:
        raise Exception("Visual agent not found in draft zone")
    tools_ids = prepare_agent_tools(project, agent, DRAFT_ZONE)
    update_visual_agent(
        agent=vis_agent, tools_ids=tools_ids, llm_id=agent.get("llmid"), prompt=agent.get("system_prompt", "")
    )
    for f in files:
        try:
            f.stream.seek(0)
            folder.put_file(f.filename, f.stream)
            logger.info(f"Uploaded document: {f.filename}")
        except Exception as e:
            logger.error(f"Failed to upload {f.filename}: {e}")
            raise


def remove_connection(agent: schemas.AgentRead, connection_type: str, doc_name: Optional[str] = None) -> None:
    """Remove a document or connection folder from the draft folder."""
    project, folder = get_ua_project_folder(agent.id)
    docs = agent.documents or []
    connection_to_delete = next((doc for doc in docs if _matches_connection_doc(doc, connection_type, doc_name)), None)
    if connection_to_delete:
        vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
        if not vis_agent:
            raise Exception("Visual agent not found in draft zone")
        logger.info(f"Removing kb tool related to the connection {connection_type}")
        copy_agent = agent.model_dump()
        copy_agent["documents"] = [
            doc for doc in copy_agent["documents"] if not _matches_connection_doc(doc, connection_type, doc_name)
        ]
        ensure_merged_documents_assets(project, copy_agent, DRAFT_ZONE)
        # only update when no more documents to remove kb tool
        tools_ids = prepare_agent_tools(project, copy_agent, DRAFT_ZONE)
        update_visual_agent(agent=vis_agent, tools_ids=tools_ids, llm_id=agent.llmid, prompt=agent.system_prompt)
        if is_sharepoint_file_doc(connection_to_delete):
            source_folder_id = sharepoint_doc_source_folder_id(connection_to_delete)
            has_remaining_source_file_docs = any(
                is_sharepoint_file_doc(doc)
                and not doc.get("deletePending")
                and sharepoint_doc_source_folder_id(doc) == source_folder_id
                for doc in copy_agent.get("documents", [])
            )
            if source_folder_id and not has_remaining_source_file_docs:
                try:
                    source_folder = project.get_managed_folder(source_folder_id)
                except Exception as e:
                    logger.warning(
                        "Could not resolve SharePoint selected-file source folder '%s' for deletion: %s",
                        source_folder_id,
                        e,
                    )
                    source_folder = None
                _delete_managed_folder_object(source_folder, DRAFT_ZONE, source_folder_id)
        else:
            folder_label = connection_to_delete.get("folderId") or get_doc_name(connection_to_delete) or connection_type
            folder = None
            if connection_to_delete.get("folderId"):
                try:
                    folder = project.get_managed_folder(connection_to_delete["folderId"])
                except Exception as e:
                    logger.warning(
                        "Could not resolve managed folder '%s' for deletion in zone '%s': %s",
                        connection_to_delete.get("folderId"),
                        DRAFT_ZONE,
                        e,
                    )
            _delete_managed_folder_object(folder, DRAFT_ZONE, folder_label)
        has_remaining_connection_docs = any(
            isinstance(doc, dict) and doc.get("connectionType") == connection_type and not doc.get("deletePending")
            for doc in copy_agent.get("documents", [])
        )
        if connection_type == CloudConnections.SHAREPOINT_ONLINE.value and not has_remaining_connection_docs:
            cleanup_sharepoint_assets(project, DRAFT_ZONE)

        # TODO: handle guardrails _cleanup_extracted_document_cache ?


def remove_document(agent: schemas.AgentRead, filename: str) -> None:
    """Remove a document from the draft folder."""
    project, folder = get_ua_project_folder(agent.id)
    if agent.documents and len(agent.documents) == 1:
        vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
        if not vis_agent:
            raise Exception("Visual agent not found in draft zone")
        logger.info("No more documents, removing kb tool")
        copy_agent = agent.model_dump()
        copy_agent["documents"] = []
        # only update when no more documents to remove kb tool
        tools_ids = prepare_agent_tools(project, copy_agent, DRAFT_ZONE)
        update_visual_agent(agent=vis_agent, tools_ids=tools_ids, llm_id=agent.llmid, prompt=agent.system_prompt)
    try:
        folder.delete_file(filename)
        logger.info(f"Deleted document: {filename}")
    except Exception as e:
        logger.error(f"Failed to delete document {filename}: {e}")
        # Don't raise - file might already be gone

    # Clean up extracted documents cache
    _cleanup_extracted_document_cache(project, filename, DRAFT_ZONE)


def _cleanup_extracted_document_cache(project, filename: str, zone: str) -> None:
    """
    Clean up cached extraction and guardrails files for a deleted document.

    Args:
        project: DSSProject instance
        filename: Name of the deleted document
        zone: Zone name (draft or published)
    """
    extracted_folder = get_managed_folder_in_zone(project, f"extracted_documents_{zone}", zone)
    if not extracted_folder:
        return

    # Delete cached extraction and guardrails result files
    for suffix in ["_extracted.json", "_guardrails.json"]:
        cache_file = f"{filename}{suffix}"
        try:
            extracted_folder.delete_file(cache_file)
            logger.info(f"Deleted cached file: {cache_file}")
        except Exception:
            pass  # File doesn't exist, skip


def delete_agent_project(agent_id: str) -> None:
    """Delete the entire agent project."""
    client = dataiku.api_client()
    try:
        project = client.get_project(agent_id)
        project.delete(clear_managed_datasets=True, clear_output_managed_folders=True, clear_job_and_scenario_logs=True)
        logger.info(f"Deleted project {agent_id}")
    except DataikuException as e:
        logger.warning(f"Project {agent_id} not found or already deleted: {e}")


def launch_index_job(
    agent_id: str, zone: str = DRAFT_ZONE, include_local: bool = True, include_cloud: bool = True
) -> Optional[str]:
    project = get_ua_project(agent_id)
    _run_merge_recipe_if_exists(project, zone)
    merged_folder = get_managed_folder_in_zone(project, _merged_docs_folder_name(zone), zone)
    if merged_folder:
        ensure_embed_doc_recipe_input(project, zone, merged_folder.id)
    if not include_local and not include_cloud:
        return None

    if job_id := _run_recipe_if_exists(project, f"compute_documents_embedded_{zone}", agent_id, zone):
        return job_id

    recipe = get_recipe_in_zone(project, "embed_documents", zone)
    if not recipe:
        raise Exception(f"No embed_documents recipe found in zone '{zone}'")
    set_embed_doc_recipe_vlm(recipe, get_my_agents_vision_llm())
    job = recipe.run(wait=False)
    return job.id


def _collect_indexing_statuses(
    agent_id: str, job_id: Optional[str], scenario_runs: Optional[list[dict]]
) -> Optional[list[str]]:
    statuses: list[str] = []

    if job_id:
        live = get_job_status(job_id, agent_id)
        if not live:
            return None
        statuses.append(live.get("status"))

    if not scenario_runs:
        return statuses

    for run in scenario_runs:
        scenario_id = (run or {}).get("scenarioId")
        run_id = (run or {}).get("runId")
        if not scenario_id or not run_id:
            return None

    project = get_ua_project(agent_id)
    for run in scenario_runs:
        scenario_id = (run or {}).get("scenarioId")
        run_id = (run or {}).get("runId")
        status = _get_scenario_run_status(project, scenario_id, run_id)
        if not status:
            return None
        statuses.append(status)

    return statuses


def _status_from_indexing_statuses(statuses: list[str]) -> Optional[dict]:
    if not statuses:
        return None
    if "failure" in statuses:
        return {"status": "failure"}
    if "running" in statuses:
        return {"status": "running"}
    if "pending" in statuses:
        return {"status": "pending"}
    if all(status == "success" for status in statuses):
        return {"status": "success"}
    return None


def compute_indexing_status(
    agent_id: str, job_id: Optional[str], scenario_runs: Optional[list[dict]] = None
) -> Optional[dict]:
    statuses = _collect_indexing_statuses(agent_id, job_id, scenario_runs)
    if statuses is None:
        return None
    return _status_from_indexing_statuses(statuses)


def ensure_sharepoint_assets(project: DSSProject, agent: dict, zone_name: str) -> Optional[DSSRecipe]:
    sp_folder = _find_sp_folder(project, zone_name)
    if not sp_folder:
        logger.warning(
            "No SharePoint managed folder found for zone '%s' (name: sharepoint_documents_%s)",
            zone_name,
            zone_name,
        )
        return None

    settings_raw = sp_folder.get_settings().get_raw()
    if settings_raw.get("type") != CloudConnections.SHAREPOINT_ONLINE.value:
        logger.warning(
            "SharePoint folder found but type is '%s' instead of %s (zone '%s')",
            settings_raw.get("type"),
            CloudConnections.SHAREPOINT_ONLINE.value,
            zone_name,
        )
        return None

    ensure_merged_documents_assets(project, agent, zone_name)
    recipe_name = f"compute_documents_embedded_{zone_name}"
    try:
        recipe = project.get_recipe(recipe_name)
    except DataikuException:
        recipe = get_recipe_in_zone(project, "embed_documents", zone_name)
    kb_id = _extract_kb_id_from_recipe(recipe) if recipe else get_kb_in_zone(project, zone_name)
    if not kb_id:
        logger.warning(
            "Could not resolve merged documents KB id from recipe '%s' in zone '%s'",
            recipe_name,
            zone_name,
        )
        return recipe

    kb_tool = get_kb_tool_in_zone(project, zone_name)
    if kb_id and not kb_tool:
        create_kb_tool(
            project=project,
            kb_id=kb_id,
            kb_desc=(agent.get("kb_description") or "").strip(),
            maxDocuments=get_config().get("number_of_documents_to_retrieve", 4),
            zone_name=zone_name,
        )

    _sync_visual_agent_tools(project, agent, zone_name)
    return recipe


def _run_recipe_if_exists(project: DSSProject, recipe_name: str, agent_id: str, zone: str) -> Optional[str]:
    try:
        recipe = project.get_recipe(recipe_name)
        set_embed_doc_recipe_vlm(recipe, get_my_agents_vision_llm())
        job = recipe.run(wait=False)
        logger.info(f"Started indexing job {job.id} for recipe '{recipe_name}' in zone '{zone}' for agent {agent_id}")
        return job.id
    except DataikuException as e:
        # Recipe may not exist (e.g. no SharePoint folder configured)
        if "not found" in str(e).lower() or "does not exist" in str(e).lower():
            return None
        raise


def _run_merge_recipe_if_exists(project: DSSProject, zone: str) -> None:
    for recipe_name in [_merge_docs_recipe_name(zone)]:
        try:
            recipe = project.get_recipe(recipe_name)
            recipe.run(wait=True)
            logger.info("Ran merge recipe '%s' in zone '%s'", recipe_name, zone)
        except DataikuException as e:
            if "not found" in str(e).lower() or "does not exist" in str(e).lower():
                if recipe_name == _merge_docs_recipe_name(zone):
                    logger.warning(
                        "Merge recipe not found in zone '%s'; indexing will use current embed recipe input", zone
                    )
                    return
                continue
            raise


def run_cloud_sync_scenarios(project: DSSProject, cloud_docs: list[dict], zone_name: str = DRAFT_ZONE) -> list[dict]:
    scenario_runs = []
    for scenario in ensure_cloud_sync_scenarios(project, cloud_docs, zone_name):
        if not scenario:
            continue
        run_id = _run_scenario_and_get_run_id(scenario)
        if run_id:
            scenario_runs.append({"scenarioId": scenario.id, "runId": run_id})
    return scenario_runs


def ensure_cloud_sync_scenarios(project: DSSProject, cloud_docs: list[dict], zone_name: str = DRAFT_ZONE) -> list:
    frequency = _get_sharepoint_sync_frequency()
    interval = _get_sharepoint_sync_interval()
    scenarios = []
    for connection_type in _get_cloud_connection_types(cloud_docs):
        scenario = ensure_cloud_sync_scenario(project, zone_name, connection_type, frequency, interval)
        if scenario:
            scenarios.append(scenario)
    return scenarios


def ensure_cloud_sync_scenario(
    project: DSSProject,
    zone_name: str,
    connection_type: str,
    frequency: str = "daily",
    interval: int = 1,
):
    from backend.utils.client_utils import current_request_login

    grant_project_access_to_user(project, current_request_login())  # Ensure permissions are set for scenario creation
    scenario_name = _get_cloud_scenario_name(zone_name, connection_type)
    scenario = _get_cloud_scenario(project, zone_name, connection_type)
    kb_id, imgs_id = _merged_embed_outputs(project, zone_name)
    if scenario and not _scenario_builds_outputs(scenario, kb_id, imgs_id):
        logger.info("Recreating cloud sync scenario '%s' to build merged document outputs", scenario_name)
        scenario.delete()
        scenario = None
    if not scenario:
        if not kb_id or not imgs_id:
            logger.warning(
                "Cannot create cloud sync scenario '%s': missing KB or images outputs in zone '%s'",
                scenario_name,
                zone_name,
            )
            return None
        ensure_quick_agent_support(project)
        scenario = create_sync_scenario(project, scenario_name, kb_id, imgs_id, frequency, interval)
    else:
        _configure_cloud_scenario(schedule_scenario=scenario, frequency=frequency, interval=interval)
    return scenario


def _scenario_builds_outputs(scenario, kb_id: Optional[str], imgs_id: Optional[str]) -> bool:
    if not kb_id or not imgs_id:
        return False
    try:
        settings = scenario.get_settings()
        raw = settings.get_raw()
        steps = raw.get("params", {}).get("steps", [])
        build_item_ids = {
            build.get("itemId")
            for step in steps
            for build in step.get("params", {}).get("builds", [])
            if isinstance(build, dict)
        }
        return {kb_id, imgs_id}.issubset(build_item_ids)
    except Exception as e:
        logger.warning("Could not inspect cloud sync scenario '%s': %s", getattr(scenario, "id", ""), e)
        return False


def _get_sharepoint_sync_frequency() -> str:
    from backend.config import get_config

    return (get_config().get("sharepoint_sync_frequency") or "daily").lower()


def _get_sharepoint_sync_interval() -> int:
    from backend.config import get_config

    raw_value = get_config().get("sharepoint_sync_interval", 1)
    try:
        interval = int(raw_value)
    except (TypeError, ValueError):
        return 1
    return interval if interval >= 1 else 1


def _get_cloud_connection_types(cloud_docs: list[dict]) -> list[str]:
    connection_types = {
        doc.get("connectionType")
        for doc in cloud_docs
        if isinstance(doc, dict) and doc.get("connectionType") and not doc.get("deletePending")
    }
    return list(connection_types)


def _configure_cloud_scenario(schedule_scenario, frequency: str, interval: int = 1) -> None:
    settings = schedule_scenario.get_settings()
    settings.active = True
    settings.raw_triggers.clear()
    if frequency == "hourly":
        settings.add_hourly_trigger(0, repeat_every=interval)
    elif frequency == "weekly":
        settings.add_daily_trigger(hour=23, minute=59, days=["Monday"], repeat_every=interval)
    elif frequency == "monthly":
        settings.add_monthly_trigger(day=1, hour=23, minute=59, repeat_every=interval)
    else:
        settings.add_daily_trigger(hour=23, minute=59, repeat_every=interval)
    settings.save()


def _get_scenario_by_name(project: DSSProject, scenario_name: str):
    for sc in project.list_scenarios():
        if sc.get("name") == scenario_name:
            return project.get_scenario(sc.get("id"))
    return None


def _get_cloud_scenario_name(zone_name: str, connection_type: str) -> str:
    return f"Build {connection_type} {zone_name}"


def _get_cloud_scenario(project: DSSProject, zone_name: str, connection_type: Optional[str] = None):
    return _get_scenario_by_name(project, _get_cloud_scenario_name(zone_name, connection_type))


def _run_scenario_and_get_run_id(scenario) -> Optional[str]:
    current = scenario.get_current_run()
    if current:
        return current.id
    trigger = scenario.run()
    run = trigger.wait_for_scenario_run(no_fail=True)
    return run.id if run else None


def _get_scenario_run_status(project: DSSProject, scenario_id: str, run_id: str) -> Optional[str]:
    try:
        scenario = project.get_scenario(scenario_id)
        run = scenario.get_run(run_id)
        if run.running:
            return "running"
        outcome = run.outcome
        if outcome in ("SUCCESS", "WARNING"):
            return "success"
        if outcome in ("FAILED", "ABORTED"):
            return "failure"
        return "pending"
    except Exception as e:
        logger.error("Failed to get scenario run status for %s/%s: %s", scenario_id, run_id, e)
        return None


def _find_sp_folder(project: DSSProject, zone_name: str) -> Optional[DSSManagedFolder]:
    folder_prefix = f"sharepoint_documents_{zone_name}"
    flow = project.get_flow()

    for folder_info in project.list_managed_folders():
        folder_name = folder_info.get("name", "")
        if not folder_name.startswith(folder_prefix):
            continue
        try:
            folder = project.get_managed_folder(folder_info["id"])
        except Exception:
            continue

        try:
            folder_zone = flow.get_zone_of_object(folder)
            if folder_zone and folder_zone.name != zone_name:
                continue
        except Exception:
            pass

        try:
            if folder.get_settings().get_raw().get("type") != CloudConnections.SHAREPOINT_ONLINE.value:
                continue
        except Exception:
            continue

        return folder

    return None


def _merged_embed_outputs(project: DSSProject, zone_name: str) -> tuple[Optional[str], Optional[str]]:
    recipe_name = f"compute_documents_embedded_{zone_name}"
    try:
        recipe = project.get_recipe(recipe_name)
    except DataikuException as e:
        logger.warning(
            "Merged documents embed recipe '%s' not found in zone '%s': %s",
            recipe_name,
            zone_name,
            e,
        )
        return None, None
    return get_outputs_from_embed_docs_recipe(recipe)


def _extract_kb_id_from_recipe(recipe) -> Optional[str]:
    try:
        recipe_raw = recipe.get_settings().get_raw()
        return (
            recipe_raw.get("recipe", {}).get("outputs", {}).get("knowledge_bank", {}).get("items", [{}])[0].get("ref")
        )
    except Exception:
        return None


def _sync_visual_agent_tools(project: DSSProject, agent: dict, zone_name: str) -> None:
    vis_agent = get_visual_agent_in_zone(project, zone_name)
    if not vis_agent:
        return
    tools_ids = prepare_agent_tools(project, agent, zone_name)
    update_visual_agent(
        agent=vis_agent,
        tools_ids=tools_ids,
        llm_id=agent.get("llmid"),
        prompt=agent.get("system_prompt", ""),
    )


def cleanup_sharepoint_assets(project: DSSProject, zone_name: str) -> None:
    sharepoint_kb_tool = get_kb_tool_in_zone(project, zone_name, conn_type=CloudConnections.SHAREPOINT_ONLINE.value)
    if sharepoint_kb_tool:
        try:
            sharepoint_kb_tool.delete()
            logger.info("Deleted SharePoint KB tool in zone '%s'", zone_name)
        except Exception as e:
            logger.warning("Failed to delete SharePoint KB tool in zone '%s': %s", zone_name, e)

    scenario_name = _get_cloud_scenario_name(zone_name, CloudConnections.SHAREPOINT_ONLINE.value)
    scenario = _get_scenario_by_name(project, scenario_name)
    if scenario:
        try:
            scenario.delete()
            logger.info("Deleted SharePoint sync scenario '%s' in zone '%s'", scenario_name, zone_name)
        except Exception as e:
            logger.warning(
                "Failed to delete SharePoint sync scenario '%s' in zone '%s': %s", scenario_name, zone_name, e
            )


def get_job_status(job_id: str, agent_id: str) -> Optional[dict]:
    try:
        client = dataiku.api_client()
        project = client.get_project(agent_id)
        job = project.get_job(job_id)

        base_status = job.get_status().get("baseStatus", {})
        state = base_status.get("state", "")
        # Map DSS states to our states
        status_map = {
            "NOT_STARTED": "pending",
            "RUNNING": "running",
            "DONE": "success",
            "FAILED": "failure",
            "ABORTED": "failure",
        }

        if state == "FAILED":
            activities = base_status.get("activities")
            if not activities or not isinstance(activities, dict) or len(activities) == 0:
                logger.error("No activities found")
            first_key = next(iter(activities))
            first_activity = activities[first_key]
            if first_activity and first_activity.get("firstFailure"):
                error_message = first_activity.get("firstFailure").get("detailedMessage")
                logger.error(f"Job {job_id} failed: {error_message}")

        return {
            "status": status_map.get(state, "pending"),
            "updatedAt": datetime.utcnow().isoformat(),
        }

    except Exception as e:
        logger.error(f"Failed to get job status for {job_id}: {e}")
        return None


def sync_agent_to_published(agent: dict) -> None:
    project = get_ua_project(agent.get("id"))
    has_cloud_docs = _has_cloud_documents(agent.get("documents") or [])
    vis_agent = get_visual_agent_in_zone(project, PUBLISHED_ZONE)
    if vis_agent:
        logger.info(f"Updating existing Agent {vis_agent.id} in zone {PUBLISHED_ZONE}")
        previous_semantic_tool_ids = set(get_bound_semantic_tool_ids_in_zone(project, PUBLISHED_ZONE))
        tools_ids = prepare_agent_tools(project, agent, PUBLISHED_ZONE)
        update_visual_agent(
            agent=vis_agent,
            tools_ids=tools_ids,
            llm_id=agent.get("llmid"),
            prompt=agent.get("system_prompt", ""),
        )
        desired_tool_ids = {_normalize_agent_tool_id(tool_id) for tool_id in tools_ids}
        stale_semantic_tool_ids = previous_semantic_tool_ids - desired_tool_ids
        for tool_id in stale_semantic_tool_ids:
            try:
                project.get_agent_tool(tool_id).delete()
                logger.info(f"Deleted stale semantic tool {tool_id} from zone {PUBLISHED_ZONE}")
            except Exception as e:
                logger.warning(f"Failed to delete stale semantic tool {tool_id} from zone {PUBLISHED_ZONE}: {e}")
    else:
        logger.info(f"Creating new Agent {agent.get('name', 'Agent')} in zone {PUBLISHED_ZONE}")
        create_quick_agent(project=project, agent=agent, project_name=project.name, zone_name=PUBLISHED_ZONE)

    if not has_cloud_docs:
        cleanup_sharepoint_assets(project, PUBLISHED_ZONE)


def sync_folders_to_published(agent: dict) -> None:
    project, draft_folder = get_ua_project_folder(agent.get("id"))
    published_folder, draft_type = _sync_documents_folder_settings(project, draft_folder)
    documents = list(agent.get("documents") or [])
    has_cloud_docs = _has_cloud_documents(documents)
    has_sharepoint_docs = _has_cloud_documents(documents, CloudConnections.SHAREPOINT_ONLINE.value)

    published_agent = _build_published_agent_for_documents(
        project=project,
        agent=agent,
        has_cloud_docs=has_cloud_docs,
        has_sharepoint_docs=has_sharepoint_docs,
    )
    ensure_merged_documents_assets(project, published_agent, PUBLISHED_ZONE)
    if has_sharepoint_docs:
        ensure_sharepoint_assets(project, published_agent, PUBLISHED_ZONE)
    if has_cloud_docs:
        ensure_cloud_sync_scenarios(project, published_agent.get("documents") or [], PUBLISHED_ZONE)

    if draft_type == CloudConnections.SHAREPOINT_ONLINE.value:
        logger.info(
            f"Published documents folder configured to {CloudConnections.SHAREPOINT_ONLINE.value} from draft settings"
        )
        return

    _sync_local_documents_to_published(draft_folder, published_folder, agent)
    _sync_extracted_folder_to_published(project, _get_active_document_names(agent))


def _sync_documents_folder_settings(
    project: DSSProject, draft_folder: DSSManagedFolder
) -> tuple[DSSManagedFolder, str | None]:
    published_folder = get_managed_folder_in_zone(project, f"documents_{PUBLISHED_ZONE}", PUBLISHED_ZONE)

    draft_settings = draft_folder.get_settings()
    draft_settings_raw = draft_settings.get_raw()
    draft_type = draft_settings_raw.get("type")
    draft_params = draft_settings_raw.get("params") or {}

    if not published_folder:
        connection = draft_params.get("connection") or get_config().get("default_fs_connection")
        published_folder = project.create_managed_folder(f"documents_{PUBLISHED_ZONE}", connection_name=connection)
        zone_id = get_zone_id_by_name(project, PUBLISHED_ZONE)
        if zone_id:
            zone = project.get_flow().get_zone(zone_id)
            zone.add_item(published_folder)
        logger.info(f"Created missing documents_{PUBLISHED_ZONE} folder")

    published_settings = published_folder.get_settings()
    published_settings_raw = published_settings.get_raw()
    published_settings_raw["type"] = draft_type
    published_settings_raw["params"] = dict(draft_params)
    published_settings.save()
    return published_folder, draft_type


def _build_published_agent_for_documents(
    project: DSSProject, agent: dict, has_cloud_docs: bool, has_sharepoint_docs: bool
) -> dict:
    published_agent = dict(agent)
    documents = list(agent.get("documents") or [])

    if has_sharepoint_docs:
        published_sharepoint_docs = _sync_sharepoint_folders_to_published(project, agent)
        if published_sharepoint_docs:
            published_agent["documents"] = [
                (
                    published_sharepoint_docs.get(sharepoint_doc_replacement_key(doc), doc)
                    if isinstance(doc, dict)
                    else doc
                )
                for doc in documents
            ]
        else:
            published_agent["documents"] = documents
    else:
        published_agent["documents"] = documents
        if not has_cloud_docs:
            cleanup_sharepoint_assets(project, PUBLISHED_ZONE)

    return published_agent


def _sync_local_documents_to_published(
    draft_folder: DSSManagedFolder, published_folder: DSSManagedFolder, agent: dict
) -> None:
    logger.info(f"Syncing documents from draft to published for agent {agent['id']}")
    active_docs = _get_active_document_names(agent)

    _clear_managed_folder_contents(published_folder)

    copied_count = 0
    for doc_name in active_docs:
        try:
            response = draft_folder.get_file(doc_name)
            published_folder.put_file(doc_name, response.content)
            logger.info(f"document {doc_name} Copied to published folder")
            copied_count += 1
        except Exception as e:
            logger.error(f"Failed to copy document {doc_name}: {e}")

    logger.info(f"Synced {copied_count} documents to published folder")


def _get_active_document_names(agent: dict) -> set[str]:
    return {
        d["name"]
        for d in agent.get("documents", [])
        if isinstance(d, dict)
        and d.get("source") not in {"connection_folder"}
        and d.get("active")
        and not d.get("deletePending")
    }


def _clear_managed_folder_contents(folder: DSSManagedFolder) -> None:
    try:
        for item in folder.list_contents().get("items", []):
            path = item.get("path")
            folder.delete_file(path)
            logger.info(f"Deleted old file from published: {path}")
    except Exception as e:
        logger.error(f"Error clearing published folder: {e}")


def _resolve_sharepoint_source_folder(project: DSSProject, doc: dict) -> Optional[DSSManagedFolder]:
    folder_id = sharepoint_doc_source_folder_id(doc)
    if folder_id:
        try:
            return project.get_managed_folder(folder_id)
        except Exception as e:
            logger.warning("Could not resolve draft SharePoint folder '%s': %s", folder_id, e)

    draft_name = _sp_folder_name_for_doc(doc, DRAFT_ZONE)
    if draft_name:
        return get_managed_folder_in_zone(project, draft_name, DRAFT_ZONE)
    return None


def _ensure_published_sharepoint_folder(
    project: DSSProject, dst_name: str, src_params: dict, zone_id: str | None
) -> Optional[DSSManagedFolder]:
    dst_folder = get_managed_folder_in_zone(project, dst_name, PUBLISHED_ZONE)
    if not dst_folder:
        connection = src_params.get("connection") or get_config().get("default_fs_connection")
        dst_folder = project.create_managed_folder(dst_name, connection_name=connection)
        if zone_id:
            project.get_flow().get_zone(zone_id).add_item(dst_folder)
        logger.info("Created missing %s folder", dst_name)

    dst_settings = dst_folder.get_settings()
    dst_raw = dst_settings.get_raw()
    dst_raw["type"] = CloudConnections.SHAREPOINT_ONLINE.value
    dst_raw["params"] = dict(src_params)
    dst_settings.save()
    return dst_folder


def _sync_sharepoint_folders_to_published(project: DSSProject, agent: dict) -> dict[str, dict]:
    sharepoint_docs = [
        doc
        for doc in (agent.get("documents") or [])
        if isinstance(doc, dict)
        and doc.get("connectionType") == CloudConnections.SHAREPOINT_ONLINE.value
        and not doc.get("deletePending")
    ]
    if not sharepoint_docs:
        return {}

    zone_id = get_zone_id_by_name(project, PUBLISHED_ZONE)
    docs_by_src: dict[str, dict] = {}
    for doc in sharepoint_docs:
        src_folder = _resolve_sharepoint_source_folder(project, doc)
        if not src_folder:
            continue

        src_raw = src_folder.get_settings().get_raw()
        if src_raw.get("type") != CloudConnections.SHAREPOINT_ONLINE.value:
            continue

        src_params = src_raw.get("params") or {}
        dst_name = _sp_folder_name_for_doc(doc, PUBLISHED_ZONE)
        if not dst_name:
            continue

        dst_folder = _ensure_published_sharepoint_folder(project, dst_name, src_params, zone_id)
        if not dst_folder:
            continue

        key = sharepoint_doc_replacement_key(doc)
        docs_by_src[key] = with_sharepoint_source_folder_id(doc, dst_folder.id)
    logger.info("Synced %s SharePoint managed folders to published zone", len(sharepoint_docs))
    return docs_by_src


def _sync_extracted_folder_to_published(project, active_docs: set) -> None:
    """
    Sync the extracted_documents folder from draft to published zone.

    This copies cached text extractions and guardrails results for active documents.

    Args:
        project: DSSProject instance
        active_docs: Set of active document names to sync
    """
    draft_extracted = get_managed_folder_in_zone(project, f"extracted_documents_{DRAFT_ZONE}", DRAFT_ZONE)
    if not draft_extracted:
        logger.info("No extracted_documents_draft folder found, skipping extraction sync")
        return

    # Get or create published extracted folder
    published_extracted = get_managed_folder_in_zone(project, f"extracted_documents_{PUBLISHED_ZONE}", PUBLISHED_ZONE)
    if not published_extracted:
        try:
            from backend.config import get_config
            from backend.utils.project_utils import get_zone_id_by_name

            connection = get_config().get("default_fs_connection")
            published_extracted = project.create_managed_folder(
                f"extracted_documents_{PUBLISHED_ZONE}", connection_name=connection
            )
            # Add to published zone
            zone_id = get_zone_id_by_name(project, PUBLISHED_ZONE)
            if zone_id:
                flow = project.get_flow()
                zone = flow.get_zone(zone_id)
                zone.add_item(published_extracted)
                logger.info(f"Created extracted_documents_{PUBLISHED_ZONE} folder")
        except Exception as e:
            logger.warning(f"Failed to create published extracted folder: {e}")
            return

    # Clear existing files in published extracted folder
    try:
        for item in published_extracted.list_contents().get("items", []):
            path = item.get("path")
            published_extracted.delete_file(path)
    except Exception as e:
        logger.warning(f"Error clearing published extracted folder: {e}")

    # Copy extraction and guardrails cache files for active documents
    copied_count = 0
    for doc_name in active_docs:
        # Copy extracted text cache
        for suffix in ["_extracted.json", "_guardrails.json"]:
            cache_file = f"{doc_name}{suffix}"
            try:
                response = draft_extracted.get_file(cache_file)
                published_extracted.put_file(cache_file, response.content)
                copied_count += 1
            except Exception:
                pass  # File doesn't exist, skip

    logger.info(f"Synced {copied_count} extraction cache files to published folder")
