from enum import Enum

from typing import Any, Dict, List, Literal, Optional, Union

from typing_extensions import NotRequired, TypedDict

# =============================================================================
#  Plugin and components
# =============================================================================
SEMANTIC_MODELS_PUGIN_ID = "semantic-models-lab"

# Semantic Model tools
SEMANTIC_MODEL_QUERY_TOOL_ID = "semantic-model-query"
DATASET_SQL_TOOL_ID = "dataset-sql-query"
SEMANTIC_MODEL_QUERY_TOOL_TYPE = "Custom_agent_tool_semantic-models-lab_semantic-model-query"
DATASET_SQL_TOOL_TYPE = "Custom_agent_tool_semantic-models-lab_dataset-sql-query"
SemanticModelToolID = Literal["semantic-model-query", "dataset-sql-query"]
SemanticDatasourceTypeLiteral = Literal["SEMANTIC_MODEL", "DATASET"]
SemanticModelToolTypeLiteral = Literal[
    "Custom_agent_tool_semantic-models-lab_semantic-model-query",
    "Custom_agent_tool_semantic-models-lab_dataset-sql-query",
]
SemanticModelInteractionMode = Literal["agentic", "simple"]


END_USER_SQL_EXECUTION_PARAM = "enduser_sql_execution"


class EndUserSQLExecution(str, Enum):
    """
    Plugin security options handling the end user running SQL queries ('Access datasets as').
    They are associated with the parameter 'enduser_sql_execution'. The available options are:
    - 'TOOL_USER' to use directly the calling the tool to run the SQL
    - 'END_USER' to leverage the end user 'dkuCallerTicket' passed through the Agent's context
    - 'END_USER_IF_AVAILABLE' to leverage 'END_USER' if passed, otherwise, 'TOOL_USER'
    """

    TOOL_USER = "tool_user"
    END_USER = "enduser"
    END_USER_IF_AVAILABLE = "enduser_available"


DEFAULT_USER_SQL_EXECUTION: EndUserSQLExecution = EndUserSQLExecution.END_USER.value


class SemanticDatasourceOrigin(TypedDict):
    id: str
    name: str
    type: Literal["WORKSPACE", "DATA_COLLECTION"]


# 'NotRequired' flags are set to prevent pydantic errors when 'semantic-models-lab' features are not involved
class SemanticDatasourceCore(TypedDict):
    id: NotRequired[str]
    name: NotRequired[str]
    type: NotRequired[SemanticDatasourceTypeLiteral]
    origins: NotRequired[List[SemanticDatasourceOrigin]]


class SemanticDatasourceConfig(SemanticDatasourceCore):
    endUserSQLExecution: NotRequired[EndUserSQLExecution]


class SemanticDatasourceInfo(SemanticDatasourceConfig):
    description: NotRequired[Optional[str]]
    lastModifiedOn: NotRequired[Optional[int]]
    lastModifiedBy: NotRequired[Optional[str]]


class SemanticToolMetadata(TypedDict):
    tool_id: SemanticModelToolID
    tool_type: SemanticModelToolTypeLiteral


SEMANTIC_MODEL_TOOLS_META: Dict[SemanticDatasourceTypeLiteral, SemanticToolMetadata] = {
    "SEMANTIC_MODEL": {"tool_id": SEMANTIC_MODEL_QUERY_TOOL_ID, "tool_type": SEMANTIC_MODEL_QUERY_TOOL_TYPE},
    "DATASET": {"tool_id": DATASET_SQL_TOOL_ID, "tool_type": DATASET_SQL_TOOL_TYPE},
}


# Details types to be used in 'get_agent_details':
class DatasetSQLToolDetails(TypedDict):
    semantic_datasource: SemanticDatasourceInfo
    sem_models_llm: Optional[str]
    sem_models_interaction_mode: Optional[SemanticModelInteractionMode]


SemanticDetailsCollection = Dict[SemanticModelToolID, List[Union[DatasetSQLToolDetails]]]


def normalize_to_semantic_datasource_core(raw_info: Dict[str, Any]) -> SemanticDatasourceCore:
    """Normalizes API-provided semantic datasource metadata into a SemanticDatasourceCore object by validating required fields
    and prioritizing displayName over name when present.
    """
    normalized_info: SemanticDatasourceCore = {}
    for key in ["id", "name", "type", "origins"]:
        assert key in raw_info, f"The key '{key}' is missing (raw_info={raw_info})"
        normalized_info[key] = raw_info[key]
    # The 'displayName' overrides the Name if it exists
    if "displayName" in raw_info:
        normalized_info["name"] = raw_info["displayName"]
    return normalized_info


def normalize_to_semantic_datasource_info(raw_info: Dict[str, Any]) -> SemanticDatasourceInfo:
    """Extends 'normalize_to_semantic_datasource_core' by deriving a prioritized description (description over shortDesc)
    and populating modification metadata with safe defaults.
    """

    normalized_info: SemanticDatasourceInfo = normalize_to_semantic_datasource_core(raw_info)
    # The description is the 'shortDesc', updated by 'description' if they exist
    description = ""
    for key in ["shortDesc", "description"]:
        if key in raw_info and raw_info[key]:
            description = raw_info[key]
    normalized_info["description"] = description
    for key in ["lastModifiedOn", "lastModifiedBy"]:
        normalized_info[key] = raw_info.get(key, "")
    if enduser_sql_execution := raw_info.get("endUserSQLExecution", "None"):
        normalized_info["endUserSQLExecution"] = enduser_sql_execution
    return normalized_info
