import logging
from enum import Enum
from typing import Dict, List

import dataiku
from backend.utils.agents_utils import list_augmented_llms_by_project, map_aug_llms_id_name
from dataikuapi.dss.project import DSSProject


class ConnectionType(Enum):
    FILESYSTEM = "Filesystem"
    EC2 = "EC2"
    FTP = "FTP"
    SSH = "SSH"
    AZURE = "Azure"
    GCS = "GCS"
    HDFS = "HDFS"
    SHAREPOINT_ONLINE = "SharePointOnline"


def get_project_agents(current_project: DSSProject) -> List[Dict[str, str]]:
    agents_map = get_project_agent_mapping(current_project)
    return [
        {
            "value": current_project.project_key + ":" + agent_id,
            "label": agent_name,
            "description": "",  # description will be filled manually by admin
        }
        for agent_id, agent_name in agents_map.items()
    ]


def get_project_agent_mapping(current_project: DSSProject):
    return {
        llm.get("id"): llm.get("friendlyName")
        for llm in current_project.list_llms()
        if llm.get("id").startswith("agent")
    }


def list_selected_agents(payload):
    selected_agents = payload["rootModel"].get("agents_ids", [])
    selected_projects = payload["rootModel"].get("projects_keys", [])
    agents_map = map_agents_id_name(selected_projects, with_project_key=True)
    return [{"value": agent, "label": agents_map.get(agent)} for agent in selected_agents]


def list_selected_augmented_llms(payload):
    selected_aug_llms = payload["rootModel"].get("augmented_llms_ids", [])
    selected_projects = payload["rootModel"].get("projects_keys", [])
    aug_llms_map = map_aug_llms_id_name(selected_projects, with_project_key=True)
    return [{"value": llm, "label": aug_llms_map.get(llm)} for llm in selected_aug_llms]


def list_augmented_llms(payload):
    selected_projects = payload["rootModel"].get("projects_keys", [])
    aug_llms_by_project = list_augmented_llms_by_project(selected_projects)
    choices = []
    for project_key, llms in aug_llms_by_project.items():
        for llm in llms:
            choices.append(
                {
                    "value": llm["value"],
                    "label": f"[{project_key}] {llm['label']}",
                    "description": llm["description"],
                }
            )
    return {"choices": choices}


def list_workspaces():
    choices = []
    client = dataiku.api_client()
    for workspace in client.list_workspaces():
        # TODO should we check permissions ?
        choices.append(
            {
                "value": workspace["workspaceKey"],
                "label": f"{workspace['displayName']}",
                "description": workspace["description"],
            }
        )
    return {"choices": choices}


def list_projects_folders():
    choices = []
    client = dataiku.api_client()
    root_folder = client.get_root_project_folder()
    for folder in root_folder.list_child_folders():
        choices.append(
            {
                "value": folder.id,
                "label": f"{folder.get_name()}",
            }
        )
    return {"choices": choices}


def list_agents_by_project(selected_projects) -> Dict[str, List[Dict[str, str]]]:
    client = dataiku.api_client()
    agents_by_project: Dict[str, List[Dict[str, str]]] = {}
    if not selected_projects:
        return agents_by_project
    for project_key in selected_projects:
        project_ob: DSSProject = client.get_project(project_key)
        agents = get_project_agents(project_ob)
        if agents:
            agents_by_project[project_key] = agents
    return agents_by_project


def map_agents_id_name(selected_projects, with_project_key: bool = False) -> Dict[str, str]:
    agents_by_project = list_agents_by_project(selected_projects)
    agents_map = {}
    for project_key, agents in agents_by_project.items():
        for agent in agents:
            agents_map[agent["value"]] = agent["label"] if not with_project_key else f"[{project_key}] {agent['label']}"
    return agents_map


def list_projects():
    projects = dataiku.api_client().list_projects()
    return {
        "choices": [
            {"value": project["projectKey"], "label": f"[{project['projectKey']}] {project['name']}"}
            for project in projects
        ]
    }


def list_projects_name():
    projects = dataiku.api_client().list_projects()
    return {"choices": [{"value": project["name"], "label": project["name"]} for project in projects]}


def list_agent_tools():
    from backend.utils.project_utils import list_project_agent_tools
    # TODO when we have tool sharing we can handle tools from different projects
    tool_defs = list_project_agent_tools(dataiku.api_client().get_default_project())
    return {"choices": [{"value": tool_def["id"], "label": tool_def["name"]} for tool_def in tool_defs]}


def list_agents(payload):
    selected_projects = payload["rootModel"].get("projects_keys", [])
    agents_by_project = list_agents_by_project(selected_projects)
    choices = []
    for project_key, agents in agents_by_project.items():
        for agent in agents:
            choices.append(
                {
                    "value": agent["value"],
                    "label": f"[{project_key}] {agent['label']}",
                    "description": agent["description"],
                }
            )
    return {"choices": choices}

def list_sql_connections(client):
    conns = client.list_connections_names("sql")
    return {"choices": [{"value": conn, "label": conn} for conn in conns]}

def list_llms_by_connections_type(current_project, purpose="GENERIC_COMPLETION"):
    llm_by_type = {}
    types_to_exclude = {"SAVED_MODEL_AGENT", "RETRIEVAL_AUGMENTED"}

    # Get LLMs based on the specified purpose
    all_llms = current_project.list_llms(purpose=purpose)

    # Group LLMs by their type
    for llm in all_llms:
        llm_type = llm.get("type")

        if llm_type in types_to_exclude:
            continue

        # Initialize the list for this type if it doesn't exist
        if llm_type not in llm_by_type:
            llm_by_type[llm_type] = []

        # Add the LLM to the appropriate group
        llm_by_type[llm_type].append({"value": llm.get("id"), "label": llm.get("friendlyName"), "type": llm_type})

    return {"choices": llm_by_type}


def list_connections_by_type(client):
    connection_map = {}
    client = dataiku.api_client()

    # Loop through all connection types in the enum
    for connection_type in ConnectionType:
        try:
            dss_connections = client.list_connections_names(connection_type.value)
            if dss_connections:
                # dss_connections is a list of connection names
                for connection_name in dss_connections:
                    connection_map.setdefault(connection_type.value, []).append(connection_name)
        except Exception as e:
            logging.error(f"Error processing connection type {connection_type.value}: {e}")

    return {"choices": connection_map}


def list_workspaces_folders(client, payload):
    project_key = payload["projectKey"]
    # Get a list of managed folders
    managed_folders = client.get_project(project_key).list_managed_folders()
    choices = []

    for folder in managed_folders:
        choices.append(
            {
                "value": folder.get("id"),  # or folder.get("id")
                "label": folder.get("name"),  # or folder.get("name")
            }
        )
    return {"choices": choices}


def do(payload, config, plugin_config, inputs):  # noqa: PLR0911
    parameter_name = payload["parameterName"]
    client = dataiku.api_client()
    current_project = client.get_default_project()
    if parameter_name == "projects_keys":
        return list_projects()
    elif parameter_name == "agents_ids" or parameter_name == "agent_id":
        return list_agents(payload)
    elif parameter_name == "augmented_llms_ids":
        return list_augmented_llms(payload)
    elif parameter_name == "agent_id":
        return {"choices": list_selected_agents(payload)}
    elif parameter_name == "augmented_llm_id":
        return {"choices": list_selected_augmented_llms(payload)}
    elif parameter_name == "tools":
        return list_agent_tools()
    elif parameter_name == "stories_workspace":
        return list_workspaces()
    elif parameter_name == "template_project_key":
        return list_projects_name()
    elif parameter_name == "default_fs_connection":
        return list_connections_by_type(client)
    elif parameter_name == "db_connection":
        return list_sql_connections(client)
    elif parameter_name == "agents_folder":
        return list_projects_folders()
    elif parameter_name == "llm_id" or parameter_name == "embedding_llms" or "charts_generation_llm_id":
        purpose = payload.get("purpose", "GENERIC_COMPLETION")
        return list_llms_by_connections_type(current_project, purpose=purpose)
    elif parameter_name == "workspace_folder":
        return list_workspaces_folders(client, payload)
    else:
        return {
            "choices": [
                {
                    "value": "wrong",
                    "label": f"Problem getting the name of the parameter.",
                }
            ]
        }
