from __future__ import annotations

from enum import Enum
from typing import Any, Dict, Set

import dataiku
from flask import Blueprint, jsonify, request
from werkzeug.exceptions import BadRequest, PreconditionFailed

from backend.services.admin_config_service import read_config, write_config
from backend.services.prompts_library_service import read_prompts, write_prompts
from backend.socket import socketio
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import list_project_agent_tools

admin_bp = Blueprint("admin", __name__, url_prefix="/admin")

logger = get_logger(__name__)


def _canonicalize_prompt_attrs(raw: Dict[str, Any], known_ids: Set[str]) -> Dict[str, Any]:
    cleaned: Dict[str, Any] = {}
    if not isinstance(raw, dict):
        return cleaned
    for pid, attrs in raw.items():
        # drop unknown prompt IDs and non-dict values
        if pid not in known_ids or not isinstance(attrs, dict):
            continue
        # drop keys with None values
        attrs2 = {k: v for k, v in attrs.items() if v is not None}
        # whitelist and types
        attrs2 = {k: v for k, v in attrs2.items() if k == "enabled" and isinstance(v, bool)}
        # keep sparse map: enabled:true alone is redundant
        if attrs2.get("enabled") is True and len(attrs2) == 1:
            continue
        # drop empty override objects
        if not attrs2:
            continue
        cleaned[pid] = attrs2
    return cleaned


def _canonicalize_prompts_library(payload: Dict[str, Any]) -> Dict[str, Any]:
    if not isinstance(payload, dict):
        return {"version": "1.0.0", "prompts": []}
    version = payload.get("version")
    prompts = payload.get("prompts")
    if not isinstance(prompts, list):
        prompts = []
    if not isinstance(version, str):
        version = "1.0.0"
    return {"version": version, "prompts": prompts}


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


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

    for connection_type in ConnectionType:
        try:
            dss_connections = client.list_connections_names(connection_type.value)
            if dss_connections:
                for connection_name in dss_connections:
                    connection_map.setdefault(connection_type.value, []).append(connection_name)
        except Exception as e:
            logger.error(f"Error processing connection type {connection_type.value}: {e}")
    return connection_map


# returns a list of workspaces, each workspace is a dict with:
# id: workspace["workspaceKey"]
# displayName: workspace["displayName"]
# description: workspace["description"]
def list_workspaces():
    return dataiku.api_client().list_workspaces()


# Recursively build folder hierarchy
def _build_folder_tree(folder, parent_path="", level=0):
    """
    Recursively build a flat list with hierarchy information.
    Each item contains: id, name, path, level, hasChildren
    """
    try:
        folder_id = folder.id
        folder_name = folder.get_name()
        current_path = f"{parent_path}/{folder_name}" if parent_path else folder_name

        result = [
            {
                "id": folder_id,
                "name": folder_name,
                "path": current_path,
                "level": level,
                "hasChildren": False,  # Will be updated if children exist
            }
        ]

        # Get children
        try:
            children = folder.list_child_folders()
            if children:
                result[0]["hasChildren"] = True
                for child in children:
                    result.extend(_build_folder_tree(child, current_path, level + 1))
        except Exception as e:
            logger.warning(f"Could not list children for folder {folder_name}: {e}")

        return result
    except Exception as e:
        logger.error(f"Error building folder tree: {e}")
        return []


def list_projects_folders():
    """
    Returns a flat list of all folders with hierarchy information.
    Each folder has: id, name, path, level, hasChildren
    """
    try:
        client = dataiku.api_client()
        root_folder = client.get_root_project_folder()

        # Start with root
        folders = [{"id": "ROOT", "name": "Root", "path": "/", "level": 0, "hasChildren": True}]

        # Add all children recursively
        try:
            children = root_folder.list_child_folders()
            for child in children:
                folders.extend(_build_folder_tree(child, "", 1))
        except Exception as e:
            logger.warning(f"Could not list root children: {e}")

        return folders
    except Exception as e:
        logger.error(f"Error listing project folders: {e}")
        return []


def list_agent_tools():
    try:
        project = dataiku.api_client().get_default_project()
        tool_defs = []
        tools_objects = list_project_agent_tools(project)

        for tool in tools_objects:
            tool_type = tool.get("type")
            tool_name = tool.get("name")
            project_key = tool.get("projectKey")
            tool_id = tool.get("id")
            tool_description = tool.get("additionalDescriptionForLLM", "")
            tool_defs.append({
                "projectKey": project_key,
                "id": tool_id,
                "type": tool_type,
                "name": tool_name,
                "description": tool_description
            })

        return tool_defs
    except Exception as e:
        logger.error(f"Error listing agent tools: {e}")
        return []


def list_managed_folders():
    """
    Returns a list of managed folders in the current project.
    """
    try:
        client = dataiku.api_client()
        project = client.get_default_project()
        folders = []

        for folder in project.list_managed_folders():
            folder_id = folder.get("id")
            folder_name = folder.get("name", folder_id)
            folders.append({"id": folder_id, "name": folder_name})

        return folders
    except Exception as e:
        logger.error(f"Error listing managed folders: {e}")
        return []


@admin_bp.route("/config", methods=["GET"])
@log_http_request
def get_admin_config():
    config, etag = read_config()
    payload = {**config, "_meta": {"etag": etag}}
    return jsonify({"data": payload}), 200


@admin_bp.route("/config", methods=["PUT"])
@log_http_request
def put_admin_config():
    body = request.get_json(silent=True) or {}
    config = body.get("config")
    etag = body.get("etag")
    socket_id = body.get("socketId")

    if not isinstance(config, dict):
        raise BadRequest("invalid config")

    # Canonicalize prompt attributes using current prompts definitions (cached)
    try:
        defs, _ = read_prompts()
        known_ids = {p.get("id") for p in defs.get("prompts", []) if p.get("id")}
    except Exception:
        known_ids = set()
    attrs = config.get("promptsAttributesById") or {}
    config["promptsAttributesById"] = _canonicalize_prompt_attrs(attrs, known_ids)

    try:
        new_etag = write_config(config, expected_etag=etag)
    except PreconditionFailed as e:
        return jsonify({"error": str(e)}), 412
    except Exception as e:
        return jsonify({"error": str(e)}), 500

    fresh, fresh_etag = read_config()
    payload = {**fresh, "_meta": {"etag": fresh_etag}}
    
    # Notify all connected clients that admin config has changed (except the specific socket that made the change)
    try:
        socketio.emit("admin:config_updated", {"timestamp": fresh_etag}, skip_sid=socket_id)
    except Exception as e:
        logger.error(f"Error emitting admin:config_updated event: {e}")
    
    return jsonify({"data": payload}), 200


@admin_bp.route("/enterprise/projects", methods=["GET"])
@log_http_request
def list_projects():
    try:
        projects = dataiku.api_client().list_projects()
        workspaces = list_workspaces()

        payload = {
            "projects": [
                {
                    "key": p.get("projectKey"),
                    "name": p.get("name"),
                }
                for p in projects
                if p.get("projectKey")
            ],
            "workspaces": [
                {
                    "id": ws.get("workspaceKey"),
                    "label": ws.get("displayName") or ws.get("workspaceKey"),
                    "description": ws.get("description", ""),
                }
                for ws in workspaces
            ],
        }
        return jsonify({"data": payload}), 200
    except Exception as e:
        logger.error(f"Error listing projects: {e}")
        return jsonify({"error": str(e)}), 500


@admin_bp.route("/enterprise/projects/<project_key>/agents", methods=["GET"])
@log_http_request
def list_project_agents(project_key: str):
    try:
        proj = dataiku.api_client().get_project(project_key)
        agents = [{"id": agent.get("id"), "name": agent.get("name"), "type": "agent"} for agent in proj.list_agents()]
        aug_llms = [
            {"id": llm.get("id"), "name": llm.get("friendlyName"), "type": "augmented_llm"}
            for llm in proj.list_llms()
            if llm.get("id").startswith("retrieval-augmented-llm")
        ]
        return jsonify(
            {"data": {"project_key": project_key, "agents": sorted(agents + aug_llms, key=lambda x: x["name"].lower())}}
        ), 200
    except Exception as e:
        logger.error(f"Error listing agents for project {project_key}: {e}")
        return jsonify({"error": str(e)}), 500


@admin_bp.route("/myagents/options", methods=["GET"])
@log_http_request
def myagents_options():
    try:
        client = dataiku.api_client()

        # Get real filesystem connections
        connections_map = get_connections_map(client)
        filesystem_connections = []

        # Add filesystem connections
        if "Filesystem" in connections_map:
            for conn_name in connections_map["Filesystem"]:
                filesystem_connections.append({"id": conn_name, "label": conn_name, "type": "Filesystem"})

        # Add other connection types if needed
        for conn_type in ["EC2", "FTP", "SSH", "Azure", "GCS", "HDFS", "SharePointOnline"]:
            if conn_type in connections_map:
                for conn_name in connections_map[conn_type]:
                    filesystem_connections.append(
                        {"id": conn_name, "label": f"{conn_name} ({conn_type})", "type": conn_type}
                    )

        # Get real project folders
        folders = list_projects_folders()

        # Get real agent tools
        tools = list_agent_tools()

        data = {
            "fileSystemConnections": filesystem_connections,
            "folders": folders,
            "managedTools": tools,
        }

        return jsonify({"data": data}), 200
    except Exception as e:
        logger.error(f"Error getting myagents options: {e}")
        return jsonify({"error": str(e)}), 500


@admin_bp.route("/managed-folders", methods=["GET"])
@log_http_request
def get_managed_folders():
    """Get list of managed folders for documents uploads."""
    try:
        folders = list_managed_folders()
        return jsonify({"data": folders}), 200
    except Exception as e:
        logger.error(f"Error getting managed folders: {e}")
        return jsonify({"error": str(e)}), 500


TYPES_TO_EXCLUDE = {"SAVED_MODEL_AGENT", "RETRIEVAL_AUGMENTED"}


def list_llms_by_purpose(project, purpose):
    llms = []
    all_llms = project.list_llms(purpose=purpose)

    for llm in all_llms:
        llm_type = llm.get("type")
        if llm_type in TYPES_TO_EXCLUDE:
            continue

        llms.append({"id": llm.get("id"), "name": llm.get("friendlyName"), "type": llm_type})

    return llms


@admin_bp.route("/llms/available", methods=["GET"])
@log_http_request
def list_available_llms():
    try:
        client = dataiku.api_client()
        project = client.get_project(dataiku.default_project_key())

        text_models = list_llms_by_purpose(project, "GENERIC_COMPLETION")
        embedding_models = list_llms_by_purpose(project, "TEXT_EMBEDDING_EXTRACTION")

        data = {"textCompletionModels": text_models, "embeddingModels": embedding_models}

        return jsonify({"data": data}), 200

    except Exception as e:
        logger.exception("Erreur lors de la récupération des LLMs")
        return jsonify({"error": str(e)}), 500


@admin_bp.route("/prompts/library", methods=["GET"])
@log_http_request
def get_prompts_library():
    """Get the complete prompts library with etag."""
    try:
        data, etag = read_prompts()
        payload = {**data, "_meta": {"etag": etag}}
        return jsonify({"data": payload}), 200
    except Exception as e:
        logger.error(f"Error getting prompts library: {e}")
        return jsonify({"error": str(e)}), 500


@admin_bp.route("/prompts/library", methods=["PUT"])
@log_http_request
def put_prompts_library():
    """Save the entire prompts library with optimistic locking."""
    body = request.get_json(silent=True) or {}
    prompts_data = body.get("data")
    etag = body.get("etag")
    socket_id = body.get("socketId")

    if not isinstance(prompts_data, dict):
        raise BadRequest("invalid prompts data")

    # Basic validation
    if "prompts" not in prompts_data or not isinstance(prompts_data["prompts"], list):
        raise BadRequest("prompts data must contain a 'prompts' array")

    prompts_data = _canonicalize_prompts_library(prompts_data)

    try:
        new_etag = write_prompts(prompts_data, expected_etag=etag)
    except PreconditionFailed as e:
        return jsonify({"error": str(e)}), 412
    except Exception as e:
        logger.error(f"Error saving prompts library: {e}")
        return jsonify({"error": str(e)}), 500

    # Return fresh data with new etag
    fresh, fresh_etag = read_prompts()
    payload = {**fresh, "_meta": {"etag": fresh_etag}}
    
    # Notify all connected clients that prompts library has changed (except the specific socket that made the change)
    try:
        socketio.emit("admin:prompts_library_updated", {"timestamp": fresh_etag}, skip_sid=socket_id)
    except Exception as e:
        logger.error(f"Error emitting admin:prompts_library_updated event: {e}")
    
    return jsonify({"data": payload}), 200
