from typing import Any

from dataikuapi.dss.agent_tool import DSSAgentTool

from backend.constants import DRAFT_ZONE, PUBLISHED_ZONE
from backend.schemas.schemas import AgentRead
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import compute_custom_tool_name, get_ua_project, list_project_agent_tools
from backend.utils.semantic_models_utils import (
    DEFAULT_USER_SQL_EXECUTION,
    END_USER_SQL_EXECUTION_PARAM,
    SEMANTIC_MODEL_TOOLS_META,
    SemanticDatasourceInfo,
    SemanticModelToolID,
)

logger = get_logger(__name__)


def upgrade_semantic_datasources(agent: AgentRead) -> tuple[dict[str, Any], bool]:
    """Upgrades Agent semantic tools to enforce the setting 'enduser_sql_execution' to 'enduser'"""
    datasources_upgraded = False
    draft_agent_dict = agent.model_dump(mode="json")
    published_version = draft_agent_dict.get("published_version", {}) or {}
    if published_version and isinstance(published_version, dict):
        published_datasources = published_version.get("datasources", []) or []
    else:
        published_datasources = []
    zone_datasources: dict[str, list[SemanticDatasourceInfo]] = {
        DRAFT_ZONE: draft_agent_dict.get("datasources", []) or [],
        PUBLISHED_ZONE: published_datasources,
    }
    updated_zone_datasources: dict[str, list[SemanticDatasourceInfo]] = {
        DRAFT_ZONE: [],
        PUBLISHED_ZONE: [],
    }
    project = get_ua_project(agent.id)
    # 'tool_by_names' is initialized empty to prevent looping on tools if no updgrade is needed
    tool_by_names: dict[str, DSSAgentTool] = {}
    for zone_name, semantic_datasources in zone_datasources.items():
        for semantic_datasource in semantic_datasources:
            if semantic_datasource.get("endUserSQLExecution", "") != DEFAULT_USER_SQL_EXECUTION:
                if not tool_by_names:
                    tool_by_names = {tool.name: tool.to_agent_tool() for tool in list_project_agent_tools(project)}
                semantic_model_source_id = semantic_datasource["id"].split(".")[-1]
                semantic_datasource_type = semantic_datasource["type"]
                semantic_model_tool_id: SemanticModelToolID = SEMANTIC_MODEL_TOOLS_META[semantic_datasource_type][
                    "tool_id"
                ]
                # TODO-Chat-with-tabular-files / update with the new functions to handle semantic tools:
                semantic_tool_name = compute_custom_tool_name(
                    tool_id=semantic_model_tool_id,
                    base_name=semantic_model_source_id,
                    zone_name=zone_name,
                )
                try:
                    semantic_tool = tool_by_names[semantic_tool_name]
                    tool_settings = semantic_tool.get_settings()
                    enduser_sql_execution = tool_settings.params["config"].get(END_USER_SQL_EXECUTION_PARAM, "")
                    if enduser_sql_execution != DEFAULT_USER_SQL_EXECUTION:
                        logger.info(
                            f"updating the Agent '{agent.id}' tool '{semantic_tool_name}' param '{END_USER_SQL_EXECUTION_PARAM}' with value '{DEFAULT_USER_SQL_EXECUTION}' (zone='{zone_name}')"
                        )
                        tool_settings.params["config"][END_USER_SQL_EXECUTION_PARAM] = DEFAULT_USER_SQL_EXECUTION
                        tool_settings.save()
                        semantic_datasource["endUserSQLExecution"] = DEFAULT_USER_SQL_EXECUTION
                        updated_zone_datasources[zone_name].append(semantic_datasource)
                        datasources_upgraded = True
                except Exception as e:
                    logger.exception(
                        f"Error updating the Agent '{agent.id}' '{END_USER_SQL_EXECUTION_PARAM}' in zone {zone_name}: {e} (semantic_datasource = {semantic_datasource}, semantic_tool_name = {semantic_tool_name}, tool_by_names = {tool_by_names})"
                    )

    draft_agent_dict["datasources"] = updated_zone_datasources[DRAFT_ZONE]
    if not draft_agent_dict["published_version"]:
        draft_agent_dict["published_version"] = {}
    draft_agent_dict["published_version"]["datasources"] = updated_zone_datasources[PUBLISHED_ZONE]
    return draft_agent_dict, datasources_upgraded
