from __future__ import annotations

import hashlib
import re
from typing import Any, AnyStr, Dict, List, Optional, Tuple, Union

import dataiku
from backend.constants import CloudConnections
from backend.utils.api_clients import AH_GLOBAL_API_CLIENT
from backend.utils.dss_utils import compare_versions, with_impersonation
from backend.utils.logging_utils import get_logger
from backend.utils.semantic_models_utils import (
    DATASET_SQL_TOOL_ID,
    DATASET_SQL_TOOL_TYPE,
    SEMANTIC_MODEL_QUERY_TOOL_ID,
    SEMANTIC_MODEL_QUERY_TOOL_TYPE,
    DatasetSQLToolDetails,
    SemanticDetailsCollection,
    SemanticModelToolID,
    SemanticModelToolTypeLiteral,
)
from dataikuapi.dss.agent import DSSAgent
from dataikuapi.dss.agent_tool import DSSAgentTool
from dataikuapi.dss.managedfolder import DSSManagedFolder
from dataikuapi.dss.project import DSSProject
from dataikuapi.dss.recipe import DSSRecipe, SingleOutputRecipeCreator

logger = get_logger(__name__)

SEMANTIC_TOOL_TYPES = {
    SEMANTIC_MODEL_QUERY_TOOL_TYPE,
    DATASET_SQL_TOOL_TYPE,
}


def _build_embed_doc_recipe_rules(llm_id: str) -> List[Dict[str, Any]]:
    # Local import avoids a circular dependency: backend.config imports this module.
    from backend.config import get_image_validation

    return [
        {
            "filter": {
                "distinct": False,
                "enabled": True,
                "uiData": {
                    "mode": "||",
                    "conditions": [
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "pdf",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "docx",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "pptx",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "jpg",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "png",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                    ],
                    "$filterOptions": "rules",
                    "$latestOperator": "||",
                },
                "$status": {},
            },
            "actionToPerform": "VLM",
            "vlmSettings": {
                "customNbPages": 5,
                "customPagesOverlap": 1,
                "llmId": llm_id,
                "useCustomPrompt": False,
                "customPrompt": "Generate a concise summary, up to 3000 characters, derived from the screenshot(s) of document page(s) provided.\nBegin with a brief overview and highlight crucial words, facts, or concepts to enhance both semantic and keyword searchability.\nOmit any references to the original source.\n",
                "splitUnit": "CUSTOM",
            },
            "splittingSettings": {"chunkSizeCharacters": 3000, "chunkOverlapCharacters": 120},
            "storeInMultimodalColumn": "IMAGES",
            "reExtractUnmodifiedDocuments": False,
        },
        {
            "filter": {
                "distinct": False,
                "enabled": True,
                "uiData": {
                    "mode": "||",
                    "conditions": [
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "html",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "txt",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                        {
                            "input": "file extension",
                            "operator": "== [string]i",
                            "col": "file name",
                            "string": "md",
                            "num": 0,
                            "num2": 0,
                            "$showList": False,
                        },
                    ],
                    "$filterOptions": "rules",
                    "$latestOperator": "||",
                },
                "$status": {},
            },
            "actionToPerform": "STRUCTURED",
            "structuredSettings": {
                "maxSectionDepth": 6,
                "imageHandlingMode": "IGNORE",
                "ocrEngine": "AUTO",
                "ocrLanguages": "en",
                "llmId": "",
                "useCustomVLMPrompt": False,
                "customVLMPrompt": " You are analysing a document image to extract structured information for downstream processing and retrieval (RAG). Your goal is to extract all meaningful, reusable content in a markdown format up to 3000 characters.\n Instructions:\n 1. Classify the image type: diagram, chart, table, handwriting, printed text, or mixed.\n 2. Extract and describe based on type:\n - Diagram or Chart: Provide a detailed description of the structure, relationships, labels, axes, values, legends, and overall meaning. Convert all meaningful content into text. Summarize insights if relevant.\n - Handwritten or Printed Text: Transcribe all visible text accurately. Preserve formatting and structure. Use markdown headings, bullets, or code blocks as appropriate.\n - Tables: Reconstruct the table as markdown. Ensure all rows and columns are preserved. Retain headings and numerical precision.\n - Mixed Content: Separate and extract each part individually (e.g., transcribe text, describe chart, and format table).\n 3. Preserve structure and context: Use markdown to reflect hierarchy, formatting, and clarity.\n 4. Avoid assumptions or hallucinations. If any part is unclear, describe it as such.\n 5. Ensure each section is a reusable, standalone chunk.\n If the image is a logo, signature, or QR code, don't describe it and return an empty string.\n",
                "enableImageClassificationFiltering": get_image_validation(),
                "splitUnit": "SECTION",
            },
            "splittingSettings": {"chunkSizeCharacters": 3000, "chunkOverlapCharacters": 120},
            "storeInMultimodalColumn": "FULL_CONTENT",
            "reExtractUnmodifiedDocuments": False,
        },
    ]


def _get_tool_item_value(tool: Any, key: str) -> Any:
    if isinstance(tool, dict):
        return tool.get(key)
    return getattr(tool, key, None)


def _normalize_agent_tool_id(tool_id: Optional[str]) -> str:
    if not tool_id:
        return ""
    if ":" in tool_id:
        return tool_id.split(":", 1)[1]
    if "." in tool_id:
        return tool_id.split(".", 1)[1]
    return tool_id


def list_project_agent_tools(project: DSSProject):
    # Check DSS version to determine if we can include shared tools
    client = AH_GLOBAL_API_CLIENT
    dss_version = client.get_instance_info().raw.get("dssVersion", "0.0.0")
    if dss_version == "0.0.0":
        logger.warning("Could not retrieve DSS version")
    if (compare_versions(dss_version, "14.3.0")) >= 0:
        return project.list_agent_tools(include_shared=True, as_type="listitems")
    return [
        {"id": tool.id, "name": tool.name, "type": tool.type, "projectKey": project.project_key}
        for tool in project.list_agent_tools(as_type="listitems")
    ]


def get_ua_project(agent_id: str) -> DSSProject:
    client = AH_GLOBAL_API_CLIENT
    if "ac_user_agent_" not in agent_id:
        raise Exception("Agent project key not found")
    logger.info(f"Using existing project {agent_id}")
    return client.get_project(agent_id)


def create_embed_doc_recipe(
    project: DSSProject,
    zone: str,
    folder_id: str,
    llm_id: str,
    embedding_llm: str,
    recipe_name: Optional[str] = None,
    kb_name: Optional[str] = None,
):
    recipe_name = recipe_name or f"compute_documents_embedded_{zone}"
    kb_name = kb_name or f"documents_embedded_{zone}"
    rc = project.new_recipe("embed_documents", recipe_name)
    rc.with_input(folder_id)
    rc.with_output_knowledge_bank(kb_name, embedding_llm)
    rc.with_vlm(llm_id)  # should  this be an option text extraction or vlm?
    r = rc.create()
    recipe_settings = r.get_settings()
    recipe_params = recipe_settings.get_recipe_params()
    recipe_params["extractionMode"] = "CUSTOM_RULES"
    recipe_params["defaultVlmId"] = llm_id
    recipe_params["defaultImageHandlingMode"] = "OCR_SCANNED"
    recipe_params["rules"] = _build_embed_doc_recipe_rules(llm_id)
    recipe_settings.save()
    # r.run()
    return r


def ensure_embed_doc_recipe_input(
    project: DSSProject,
    zone: str,
    folder_id: str,
) -> Optional[DSSRecipe]:
    from dataikuapi.utils import DataikuException

    recipe_name = f"compute_documents_embedded_{zone}"
    try:
        recipe = project.get_recipe(recipe_name)
    except DataikuException:
        logger.warning("Embed documents recipe '%s' not found in zone '%s'", recipe_name, zone)
        return None

    settings = recipe.get_settings()
    definition = settings.get_recipe_raw_definition()
    inputs = definition.setdefault("inputs", {})
    main_input = inputs.setdefault("main", {})
    items = main_input.setdefault("items", [])

    current_ref = items[0].get("ref") if items else None
    if current_ref == folder_id and len(items) == 1:
        return recipe

    logger.info(
        "Repointing embed documents recipe '%s' in zone '%s' to input folder '%s'",
        recipe_name,
        zone,
        folder_id,
    )
    preserved_item = dict(items[0]) if items else {}
    preserved_item["ref"] = folder_id
    main_input["items"] = [preserved_item]
    settings.save()
    return recipe


def configure_merge_folder_recipe(
    recipe: DSSRecipe,
    explicit_paths_dataset_id: Optional[str] = None,
    explicit_mode_missing_path_handling: str = "FAIL",
) -> None:
    dp = recipe.get_definition_and_payload()
    definition = dp.get_recipe_raw_definition()
    params = definition.setdefault("params", {})
    # can be OVERWRITE, SUFFIX, or FAIL
    params["conflictHandling"] = "SUFFIX"
    inputs = definition.setdefault("inputs", {})
    if explicit_paths_dataset_id:
        inputs["explicitPaths"] = {"items": [{"ref": explicit_paths_dataset_id}]}
        params["explicitModeMissingPathHandling"] = explicit_mode_missing_path_handling
    else:
        inputs.pop("explicitPaths", None)
        params.pop("explicitModeMissingPathHandling", None)
    recipe.set_definition_and_payload(dp)


@with_impersonation
def create_or_update_merge_folder_recipe(
    project: DSSProject,
    zone_name: str,
    input_folder_ids: list[str],
    output_folder_id: str,
    recipe_name: Optional[str] = None,
    explicit_paths_dataset_id: Optional[str] = None,
) -> DSSRecipe:
    recipe_name = recipe_name or f"merge_documents_{zone_name}"
    unique_input_folder_ids = list(dict.fromkeys(input_folder_ids))
    expected_explicit_path_ids = [explicit_paths_dataset_id] if explicit_paths_dataset_id else []
    existing_recipe = None
    for item in project.list_recipes():
        if item.get("name") == recipe_name:
            existing_recipe = project.get_recipe(recipe_name)
            break

    if existing_recipe:
        existing_inputs = existing_recipe.get_settings().get_recipe_inputs().get("main", {}).get("items", [])
        existing_input_ids = [item.get("ref") for item in existing_inputs]
        existing_explicit_inputs = (
            existing_recipe.get_settings().get_recipe_inputs().get("explicitPaths", {}).get("items", [])
        )
        existing_explicit_input_ids = [item.get("ref") for item in existing_explicit_inputs]
        existing_outputs = existing_recipe.get_settings().get_recipe_outputs().get("main", {}).get("items", [])
        existing_output_ids = [item.get("ref") for item in existing_outputs]
        if (
            existing_input_ids == unique_input_folder_ids
            and existing_output_ids == [output_folder_id]
            and existing_explicit_input_ids == expected_explicit_path_ids
        ):
            configure_merge_folder_recipe(existing_recipe, explicit_paths_dataset_id)
            return existing_recipe
        logger.info(
            "Recreating merge folder recipe '%s' in zone '%s' with inputs %s, explicit paths %s and output %s",
            recipe_name,
            zone_name,
            unique_input_folder_ids,
            expected_explicit_path_ids,
            output_folder_id,
        )
        existing_recipe.delete()

    builder = SingleOutputRecipeCreator("merge_folder", recipe_name, project)
    for folder_id in unique_input_folder_ids:
        builder.with_input(folder_id)
    builder.with_existing_output(output_folder_id)
    recipe = builder.create()
    zone_id = get_zone_id_by_name(project, zone_name)
    if zone_id:
        project.get_flow().get_zone(zone_id).add_item(recipe)
    configure_merge_folder_recipe(recipe, explicit_paths_dataset_id)
    return recipe


def set_embed_doc_recipe_vlm(recipe: DSSRecipe, llm_id: str) -> None:
    recipe_settings = recipe.get_settings()
    recipe_params = recipe_settings.get_recipe_params()
    recipe_params["defaultVlmId"] = llm_id
    recipe_params["rules"] = _build_embed_doc_recipe_rules(llm_id)
    recipe_settings.save()


def create_aug_llm(project, name, kb_id, llm_id):
    return project.create_retrieval_augmented_llm(name, kb_id, llm_id)


def create_kb_tool(
    project: DSSProject,
    kb_id: str,
    kb_desc: str = "",
    maxDocuments: int = 4,
    zone_name="DRAFT",
    tool_name: Optional[str] = None,
) -> DSSAgent:
    logger.info(f"Creating KB tool for KB id {kb_id}")
    vector_search_tool_creator = project.new_agent_tool("VectorStoreSearch", tool_name or ("KB Tool " + zone_name))
    vector_search_tool_creator.with_knowledge_bank(kb_id)
    vector_search_tool = vector_search_tool_creator.create()
    # Update settings after creation
    if kb_desc or maxDocuments != 4:
        tool_settings = vector_search_tool.get_settings()
        tool_settings_raw = tool_settings.get_raw()
        tool_settings_raw["additionalDescriptionForLLM"] = kb_desc
        tool_settings_raw["params"]["maxDocuments"] = maxDocuments
        tool_settings.save()

    return vector_search_tool


def compute_custom_tool_name(tool_id: Union[SemanticModelToolID], base_name, zone_name: str = "draft"):
    MAX_BASE_NAME_CHARS = 20  # Maximum characters to prevent error with tool calling issues (ex: Max chars is 64 for OpenAI function calls)
    name_prefix = ""
    if tool_id in [SEMANTIC_MODEL_QUERY_TOOL_ID, DATASET_SQL_TOOL_ID]:
        name_prefix = f"{tool_id} "
    return f"{name_prefix}{base_name[0:MAX_BASE_NAME_CHARS]} {zone_name}"


def create_custom_tool(
    project: DSSProject,
    tool_type: Union[SemanticModelToolTypeLiteral, AnyStr],
    tool_name: str,
    tool_params: Dict[str, Any],
    zone_name: str = "draft",
) -> DSSAgentTool:
    """Creates a custom agent tool (i.e. a tool from a plugin) in a DSS project"""
    logger.info(f"Creating a custom tool '{tool_type}' in zone '{zone_name}' (params={tool_params})")
    tool = project.new_agent_tool(type=tool_type, name=tool_name).create()
    tool_settings = tool.get_settings()
    tool_settings.params["config"].update(tool_params)
    tool_settings.save()
    return tool


def edit_tool_settings(project: DSSProject, tool_id: str, tool_params: Dict[str, Any]):
    logger.debug(f"Editing the tool `{tool_id}` settings...")
    tool_settings = project.get_agent_tool(tool_id).get_settings()
    tool_settings.params["config"].update(tool_params)
    tool_settings.save()


def set_additional_tool_description(project: DSSProject, tool_id: str, tool_description: str):
    logger.debug(f"Editing the tool `{tool_id}` description...")
    tool_settings = project.get_agent_tool(tool_id).get_settings()
    tool_settings.get_raw()["additionalDescriptionForLLM"] = tool_description
    tool_settings.save()


def get_semantic_tool_in_zone(
    project: DSSProject,
    tool_type: SemanticModelToolTypeLiteral,
    source_id: str,
    zone_name: str,
):
    """Retrieves a custom tool 'from a zone'. The tool must be named using `def compute_custom_tool_name`"""
    try:
        for tool in list_project_agent_tools(project):
            tool_name = tool.name
            logger.info(f"get_semantic_tool_in_zone, checking tool '{tool_name}' ({tool.id}) of type {tool.type}")
            # A regex must be used in addition to the check on `tool.type` as it is possible that a semantic datasource contains the `zone_name`
            is_candidate = all(
                [
                    re.search(rf"{re.escape(zone_name)}$", tool_name) is not None,
                    tool.type in [SEMANTIC_MODEL_QUERY_TOOL_TYPE, DATASET_SQL_TOOL_TYPE],
                ]
            )
            if is_candidate:
                datasource_key = ""
                if tool_type == SEMANTIC_MODEL_QUERY_TOOL_TYPE:
                    datasource_key = "semantic_model_id"
                elif tool_type == DATASET_SQL_TOOL_TYPE:
                    datasource_key = "dataset_id"
                agent_tool = tool.to_agent_tool()
                tool_source_id = agent_tool.get_settings().params["config"].get(datasource_key, "")
                if (tool_source_id) and (tool_source_id == source_id):
                    logger.info(
                        f"get_semantic_tool_in_zone, found custom tool for the source '{source_id}' in zone '{zone_name}': {tool.id}"
                    )
                    return agent_tool
        logger.warning(
            f"get_semantic_tool_in_zone, No custom tool found for the source '{source_id}' in zone '{zone_name}'"
        )
        return None

    except Exception as e:
        logger.error(f"Error getting custom tool in zone '{zone_name}': {e}")
        return None


def get_bound_semantic_tool_ids_in_zone(project: DSSProject, zone_name: str) -> List[str]:
    agent = get_visual_agent_in_zone(project, zone_name)
    if not agent:
        return []

    semantic_tool_ids = {
        normalized_tool_id
        for tool in list_project_agent_tools(project)
        if _get_tool_item_value(tool, "type") in SEMANTIC_TOOL_TYPES
        if (normalized_tool_id := _normalize_agent_tool_id(_get_tool_item_value(tool, "id")))
    }
    if not semantic_tool_ids:
        return []

    bound_semantic_tool_ids = []
    agent_version_tools = agent.get_settings().get_version_settings("v1").tools
    for tool in agent_version_tools:
        normalized_tool_id = _normalize_agent_tool_id(_get_tool_item_value(tool, "toolRef"))
        if normalized_tool_id and normalized_tool_id in semantic_tool_ids:
            bound_semantic_tool_ids.append(normalized_tool_id)

    return list(dict.fromkeys(bound_semantic_tool_ids))


def get_tools_by_type(project, tool_type: str) -> List[Dict[str, Any]]:
    return [tool for tool in project.list_agent_tools() if tool["type"] == tool_type]


def get_dataset_sql_tool_details(project: DSSProject, tool_id: str) -> DatasetSQLToolDetails:
    tool = project.get_agent_tool(tool_id)
    tool_settings = tool.get_settings().get_raw()
    logger.info(f"Fetching '{DATASET_SQL_TOOL_ID}' tool datasource details for '{tool_settings['name']}'")
    tool_params = tool_settings["params"]["config"]
    project_key = tool_params["project_key"]
    dataset_name = tool_params["dataset_id"]
    tool_details: DatasetSQLToolDetails = {
        "semantic_datasource": {
            "id": f"{project_key}.{dataset_name}",
            "name": dataset_name,
            "type": "DATASET",
            "description": tool_settings.get("additionalDescriptionForLLM", ""),
        },
        "sem_models_llm": tool_params.get("llm_id", ""),
        "sem_models_interaction_mode": "agentic" if tool_params.get("agent_mode", False) else "simple",
    }
    return tool_details


def get_all_semantic_details(
    project: DSSProject, zone_name: str = "draft"
) -> Tuple[SemanticDetailsCollection, List[str]]:
    """Retrieves all the 'SemanticDatasourceInfo' details, from the Agent project"""
    # TODO: Ensure alignment with current semantic model features in Agent Hub
    logger.debug(f"Fetching all semantic datasource details in zone '{zone_name}' ...")
    agent = get_visual_agent_in_zone(project, zone_name)
    agent_version_tools = agent.get_settings().get_version_settings("v1").tools
    semantic_details_collection: SemanticDetailsCollection = {}
    semantic_tool_ids = []

    # 'dataset_sql_tool_details' handled using a list as 'tabular document uploads' or later features might require to bind several 'dataset-sql-query' to the agent
    dataset_sql_tool_details = []

    proj_dataset_sql_tools = get_tools_by_type(project, DATASET_SQL_TOOL_TYPE)
    proj_dataset_sql_tool_ids = [tool["id"] for tool in proj_dataset_sql_tools]
    agent_dataset_sql_tools = [
        tool["toolRef"] for tool in agent_version_tools if tool["toolRef"] in proj_dataset_sql_tool_ids
    ]
    for tool_id in agent_dataset_sql_tools:
        dataset_sql_tool_details.append(get_dataset_sql_tool_details(project, tool_id))
        semantic_tool_ids.append(tool_id)
    semantic_details_collection[DATASET_SQL_TOOL_ID] = dataset_sql_tool_details

    return semantic_details_collection, semantic_tool_ids


def get_kb_tool_in_zone(project: DSSProject, zone_name: str, conn_type: str = None):
    try:
        for tool in list_project_agent_tools(project):
            logger.info(f"get_kb_tool_in_zone, checking tool {tool.id} of type {tool.type}")
            if tool.type == "VectorStoreSearch":
                tool_name = (tool.name or "").lower()
                has_zone = zone_name.lower() in tool_name
                connection = conn_type.lower() in tool_name if conn_type else tool_name.startswith("kb tool")
                if has_zone and connection:
                    logger.info(f"get_kb_tool_in_zone, found KB tool in zone '{zone_name}': {tool.id}")
                    return tool.to_agent_tool()
        logger.warning(f"No KB tool found in zone '{zone_name}'")
        return None

    except Exception as e:
        logger.error(f"Error getting KB tool in zone '{zone_name}': {e}")
        return None


def get_agent_tool_display_name(tool: DSSAgentTool) -> str:
    try:
        tool_settings = tool.get_settings()
        tool_settings_raw = tool_settings.get_raw() if tool_settings else {}
    except Exception:
        tool_settings_raw = {}
    return tool_settings_raw.get("name")


def create_visual_agent(project: DSSProject, agent_name: str, tools_ids: list[str], llm_id: str, prompt: str):
    logger.info(f"Creating visual agent {agent_name} with tools {tools_ids} and llm {llm_id}")
    agent = project.create_agent(agent_name, "TOOLS_USING_AGENT")
    return update_visual_agent(agent, tools_ids, llm_id, prompt)


def update_visual_agent(agent: DSSAgent, tools_ids: list[str], llm_id: str, prompt: str):
    client = AH_GLOBAL_API_CLIENT
    logger.info(f"Updating visual agent {agent.id} with tools {tools_ids} and llm {llm_id}")
    # Set the agent settings
    agent_settings = agent.get_settings()
    s_raw = agent_settings.get_raw()
    for v in s_raw["versions"]:
        if v["versionId"] == "v1":
            # no direct setter for prompt for now
            v["toolsUsingAgentSettings"]["systemPromptAppend"] = prompt
            v["toolsUsingAgentSettings"]["tools"] = []
    vsettings = agent_settings.get_version_settings("v1")
    vsettings.llm_id = llm_id
    for tool_id in tools_ids:
        if ":" in tool_id:
            project_id, agent_tool_id = tool_id.split(":", 1)
            source_project = client.get_project(project_id)
        else:
            agent_tool_id = tool_id
            source_project = client.get_default_project()

        tool = source_project.get_agent_tool(agent_tool_id)
        vsettings.add_tool(tool)
    agent_settings.save()

    return agent


def get_agent_details(project: DSSProject, zone_name: str) -> dict:
    """get details from visual agent and kb tool in the zone

    Args:
        project (DSSProject): project containing the visual agent
        zone_name (str): zone name

    Returns:
        dict: dict with keys: llm_id, tools (list of tool ids), kb_description, system_prompt, name
    """
    agent = get_visual_agent_in_zone(project, zone_name)
    settings = agent.get_settings()
    v = settings.get_version_settings("v1")
    s_raw = settings.get_raw()
    details = {
        "name": s_raw["name"].strip().replace(f" {zone_name}", ""),
    }
    details["llm_id"] = v.llm_id

    details["kb_description"] = ""
    kb_tool = get_kb_tool_in_zone(project, zone_name)
    sharepoint_kb_tool = get_kb_tool_in_zone(project, zone_name, conn_type=CloudConnections.SHAREPOINT_ONLINE.value)
    kb_tool_ids = {_normalize_agent_tool_id(tool.id) for tool in [kb_tool, sharepoint_kb_tool] if tool}

    if kb_tool:
        kb_tool_settings_raw = kb_tool.get_settings().get_raw()
        details["kb_description"] = kb_tool_settings_raw.get("additionalDescriptionForLLM", "")
    else:
        details["kb_description"] = ""

    semantic_details_collection, semantic_tool_ids = get_all_semantic_details(project, zone_name)
    hidden_tool_ids = kb_tool_ids | {_normalize_agent_tool_id(tool_id) for tool_id in semantic_tool_ids}
    details["tools"] = [
        tool["toolRef"].replace(".", ":", 1)
        for tool in v.tools
        if _normalize_agent_tool_id(tool["toolRef"]) not in hidden_tool_ids
    ]

    for version in s_raw["versions"]:
        if version["versionId"] == "v1":
            # no direct access for prompt for now
            details["system_prompt"] = version.get("toolsUsingAgentSettings", {}).get("systemPromptAppend", "")

    # Semantic details loading
    dataset_sql_details: List[DatasetSQLToolDetails] = semantic_details_collection[DATASET_SQL_TOOL_ID]
    # TODO: Ensure alignment with current semantic model features in Agent Hub
    if dataset_sql_details:
        details.update(dataset_sql_details[0])

    logger.info(f"{zone_name} details loaded: {details} (agent '{agent.id}')")
    return details


def get_cloud_docs_folder_details(
    project: DSSProject,
    zone_name: str,
    folder_id: str,
    conn_type=str,
) -> Optional[dict]:
    """
    Return Cloud documents folder params for a zone if that zone's documents folder
    is configured as conn_type. Otherwise return None.
    """

    def _extract(folder: Optional[DSSManagedFolder]) -> Optional[dict]:
        if not folder:
            return None
        settings_raw = folder.get_settings().get_raw()
        if settings_raw.get("type") != conn_type:
            return None

        params = settings_raw.get("params") or {}
        path = (params.get("path") or "/").strip() or "/"
        if not path.startswith("/"):
            path = f"/{path}"

        details = {
            "connection": params.get("connection", ""),
            "site": params.get("site", ""),
            "drive": params.get("drive", ""),
            "path": path,
        }
        if params.get("siteId"):
            details["siteId"] = params.get("siteId")
        if params.get("driveId"):
            details["driveId"] = params.get("driveId")
        return details

    try:
        # Preferred location: dedicated SharePoint folder
        dedicated = get_folder_by_id(project, folder_id=folder_id)
        details = _extract(dedicated)
        if details:
            return details

    except Exception as e:
        logger.warning(f"Failed reading folder {folder_id} settings for zone '{zone_name}': {e}")
        return None


def get_visual_agent_in_zone(project: DSSProject, zone_name: str) -> Optional[DSSAgent]:
    try:
        if not (zone_id := get_zone_id_by_name(project, zone_name)):
            return None

        flow, fallback = project.get_flow(), None
        for mod in project.list_saved_models():
            if mod["savedModelType"] != "TOOLS_USING_AGENT":
                continue
            try:
                agent = project.get_agent(mod["id"])
                if (z := flow.get_zone_of_object(agent)) and z.id == zone_id:
                    # Exact match
                    if zone_name in mod["name"]:
                        return agent
                    # First found fallback
                    fallback = fallback or agent
            except Exception:
                continue
        return fallback
    except Exception as e:
        logger.error(f"Error getting visual agent in zone '{zone_name}': {e}")
        return None


def get_zone_id_by_name(project: DSSProject, zone_name: str) -> Optional[str]:
    try:
        flow = project.get_flow()
        zones = flow.list_zones()

        for zone in zones:
            if zone.name == zone_name:
                logger.info(f"get_zone_id_by_name, found zone '{zone_name}' with ID: {zone.id}")
                return zone.id

        logger.warning(f"Zone '{zone_name}' not found in project")
        return None

    except Exception as e:
        logger.error(f"Error getting zone ID for '{zone_name}': {e}")
        return None


def get_folder_by_id(project: DSSProject, folder_id: str) -> Optional[DSSManagedFolder]:
    try:
        return project.get_managed_folder(folder_id)
    except Exception as e:
        logger.error(f"Error getting folder with ID '{folder_id}': {e}")
        return None


def get_managed_folder_size(agent_id: str, folder_id: str) -> int | None:
    try:
        project = get_ua_project(agent_id)
        folder = get_folder_by_id(project, folder_id)
        if not folder:
            return None
        items = (folder.list_contents() or {}).get("items", [])
        return int(sum(int((item or {}).get("size", 0) or 0) for item in items))
    except Exception as e:
        logger.warning("Could not compute %s managed folder size for %s: %s", folder_id, agent_id, e)
        return None


def get_managed_folder_in_zone(
    project: DSSProject, folder_name: str, zone_name: Optional[str] = None
) -> Optional[DSSManagedFolder]:
    try:
        zone_id = None
        if zone_name:
            zone_id = get_zone_id_by_name(project, zone_name)
            if not zone_id:
                logger.warning(f"Zone '{zone_name}' not found, cannot get folder")
                return None

        flow = project.get_flow()

        for folder_info in project.list_managed_folders():
            if folder_info["name"] == folder_name:
                folder = project.get_managed_folder(folder_info["id"])

                if not zone_name:
                    return folder

                try:
                    folder_zone = flow.get_zone_of_object(folder)
                    if folder_zone and folder_zone.id == zone_id:
                        logger.info(f"get_managed_folder_in_zone, found folder '{folder_name}' in zone '{zone_name}'")
                        return folder
                except Exception as e:
                    logger.info(f"Could not get zone for folder {folder_name}: {e}")

        logger.warning(f"Folder '{folder_name}' not found in zone '{zone_name}'")
        return None

    except Exception as e:
        logger.error(f"Error getting managed folder: {e}")
        return None


def ensure_quick_agent_support(project: DSSProject):
    s = project.get_settings()
    psettings = s.get_raw()
    # 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"
    s.save()


def grant_project_access_to_user(project: DSSProject, user: str):
    logger.info(
        f"Granting read and write permissions for user '{user}' on its quick agent project '{project.project_key}'"
    )
    project_permissions = project.get_permissions()
    project_permissions["permissions"].append({"user": user, "readProjectContent": True, "writeProjectContent": True})
    project.set_permissions(project_permissions)


@with_impersonation
def create_sharepoint_managed_folder(
    project: DSSProject,
    folder_name: str,
    connection: str,
    site: str,
    site_id: str,
    drive: str,
    drive_id: str,
    path: str = "/",
    zone_name: str | None = None,
) -> dict[str, Any]:

    logger.info(
        f"Creating SharePoint managed folder '{folder_name}' in project '{project.project_key}'"
        + (f", zone '{zone_name}'" if zone_name else "")
    )
    folder = get_managed_folder_in_zone(project, folder_name, zone_name)
    if not folder:
        folder = project.create_managed_folder(folder_name, connection_name=connection)
        zone_id = get_zone_id_by_name(project, zone_name)
        if zone_id:
            zone = project.get_flow().get_zone(zone_id)
            zone.add_item(folder)

    settings = folder.get_settings()
    settings_raw = settings.get_raw()
    settings_raw["type"] = CloudConnections.SHAREPOINT_ONLINE.value
    params = {
        "connection": connection,
        "site": site,
        "drive": drive,
        "path": path,
    }
    if site_id:
        params["siteId"] = site_id
    if drive_id:
        params["driveId"] = drive_id
    settings_raw["params"] = params
    settings.save()
    return folder


def build_sharepoint_managed_folder_name(
    zone_name: str,
    connection: str,
    site_id: str,
    drive_id: str,
    path: str,
) -> str:
    normalized_path = (path or "/").strip() or "/"
    stable_key = "|".join(
        [
            (connection or "").strip().lower(),
            (site_id or "").strip().lower(),
            (drive_id or "").strip().lower(),
            normalized_path.lower(),
        ]
    )
    suffix = hashlib.sha1(stable_key.encode("utf-8")).hexdigest()[:10]
    return f"sharepoint_documents_{zone_name}_{suffix}"


def get_recipe_in_zone(project: DSSProject, recipe_type: str, zone_name: str) -> Optional[DSSRecipe]:
    zone_id = None

    try:
        flow = project.get_flow()
        zones = flow.list_zones()

        # Find zone ID by name
        for zone in zones:
            if zone.name == zone_name:
                zone_id = zone.id
                break

        if not zone_id:
            logger.warning(f"Zone '{zone_name}' not found")
            return None

        # Find recipe in the zone
        for item in project.list_recipes():
            if item["type"] == recipe_type:
                recipe = project.get_recipe(item["name"])
                recipe_zone = flow.get_zone_of_object(recipe)

                if recipe_zone and recipe_zone.id == zone_id:
                    logger.info(f"get_recipe_in_zone, found {recipe_type} recipe '{recipe.name}' in zone '{zone_name}'")
                    return recipe

        logger.warning(f"No {recipe_type} recipe found in zone '{zone_name}'")
        return None

    except Exception as e:
        logger.error(f"Error finding recipe in zone '{zone_name}': {e}")
        return None


def get_augmented_llm_in_zone(project: DSSProject, zone_name: str) -> Optional[str]:
    try:
        # List all LLMs in the project
        for llm_def in project.list_llms():
            llm_id = llm_def.get("id", "")

            # Augmented LLMs typically have "retrievalaugmented" in their ID
            if "retrieval-augmented-llm" in llm_id:
                # Check if this LLM matches our zone
                # This assumes naming convention like:
                # - retrievalaugmented_draft_xxx
                # - retrievalaugmented_published_xxx
                # OR check the KB name it references
                llm_name = llm_def.description
                if zone_name.lower() in llm_name.lower():
                    full_id = f"{project.project_key}:{llm_id}"
                    logger.info(f"get_augmented_llm_in_zone, found augmented LLM for zone '{zone_name}': {full_id}")
                    return full_id

                # Alternative: check which KB it uses
                # You could inspect the LLM definition to see which KB it references
                # and match that to the zone

        # If no zone-specific LLM found, log warning
        logger.warning(f"No augmented LLM found for zone '{zone_name}'")

        return None

    except Exception as e:
        logger.error(f"Error finding augmented LLM in zone '{zone_name}': {e}")
        return None


def get_kb_in_zone(project: DSSProject, zone_name: str) -> Optional[str]:
    try:
        target_name = f"documents_embedded_{zone_name}".lower()
        # List all knowledge banks in the project
        for kb_def in project.list_knowledge_banks():
            kb_id = kb_def.get("id", "")
            kb_name = (kb_def.get("name", "") or "").lower()

            # Check if this KB matches our zone
            if kb_name == target_name:
                return kb_id
        logger.warning(f"No knowledge bank found for zone '{zone_name}'")
        return None

    except Exception as e:
        logger.error(f"Error finding knowledge bank in zone '{zone_name}': {e}")
        return None


def update_project_name(project: DSSProject, new_name: str):
    logger.info(f"Renaming project {project.project_key} to '{new_name}'")
    meta = project.get_metadata()
    meta["label"] = new_name
    project.set_metadata(meta)


def update_agent_name_in_zone(project: DSSProject, zone_name: str, new_name: str):
    logger.info(f"Renaming agent in zone '{zone_name}' to '{new_name}'")
    for mod in project.list_saved_models():
        if mod["savedModelType"] == "TOOLS_USING_AGENT" and zone_name in mod["name"]:
            sm = project.get_saved_model(mod["id"])
            ds = sm.get_settings()
            ds_raw = ds.get_raw()
            ds_raw["name"] = new_name + " " + zone_name
            ds.save()
            logger.info(f"Updated agent name: '{new_name}'")
            break


def current_project_owner_login() -> Optional[str]:
    """
    Returns the login of the current project owner.
    """
    try:
        return (dataiku.api_client().get_default_project().get_permissions() or {}).get("owner")
    except Exception:
        return None


def get_outputs_from_embed_docs_recipe(embed_recipe: DSSRecipe):
    outputs = embed_recipe.get_settings().get_recipe_outputs()
    imgs_items = outputs.get("images", {}).get("items")
    imgs = imgs_items[0].get("ref") if imgs_items else None
    kb_items = outputs.get("knowledge_bank", {}).get("items")
    kb = kb_items[0].get("ref") if kb_items else None
    return kb, imgs


@with_impersonation
def create_sync_scenario(
    project: DSSProject,
    name: str,
    kb: str,
    imgs: str,
    frequency: str = "daily",
    interval: int = 1,
):
    logger.info(f"Creating sync scenario '{name}' with KB '{kb}' and images folder '{imgs}'")
    # 2) Define the scenario settings
    params_def = {
        "params": {
            "steps": [
                {
                    "id": f"build_0_true_f_{kb}_u_{imgs}",
                    "type": "build_flowitem",
                    "name": f"build folder {kb} {imgs}",
                    "enabled": True,
                    "alwaysShowComment": False,
                    "runConditionType": "RUN_IF_STATUS_MATCH",
                    "runConditionStatuses": ["SUCCESS", "WARNING"],
                    "resetScenarioStatus": False,
                    "delayBetweenRetries": 10,
                    "maxRetriesOnFail": 0,
                    "params": {
                        "builds": [
                            {"type": "RETRIEVABLE_KNOWLEDGE", "itemId": kb},
                            {"type": "MANAGED_FOLDER", "itemId": imgs},
                        ],
                        "jobType": "RECURSIVE_BUILD",
                        "autoUpdateSchemaBeforeEachRecipeRun": False,
                        "stopAtFlowZoneBoundary": False,
                        "refreshHiveMetastore": True,
                        "handleWarningsAs": "WARNING",
                        "proceedOnFailure": False,
                    },
                }
            ]
        }
    }

    # A new client must be created so that it impersonates the end user when creating the 'sync_scenario'.
    impersonated_project = dataiku.api_client().get_project(project.project_key)

    scenario = impersonated_project.create_scenario(name, "step_based", params_def)
    settings = 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()
    return scenario
