# File: backend/fetch_api.py

import hashlib
import json
import uuid
from copy import deepcopy
from io import BytesIO
from typing import Dict, List
from urllib.parse import unquote_plus

import dataiku
from flask import Blueprint, abort, g, jsonify, request, send_file

from backend.config import (
    _llm_friendly_names_by_purpose,
    _llm_friendly_short_names_by_purpose,
    get_default_embedding_llm,
    get_default_llm_id,
    get_llms,
    get_tools_details,
    get_uploads_managedfolder_id,
)
from backend.routes.admin import SQLConnectionType
from backend.routes.user import has_dataset_read_access
from backend.utils.dss_utils import with_impersonation
from backend.utils.logger_utils import log_http_request
from backend.utils.logging_utils import get_logger
from backend.utils.semantic_models_utils import (
    SemanticDatasourceInfo,
    normalize_to_semantic_datasource_info,
    DEFAULT_USER_SQL_EXECUTION,
)
from backend.utils.utils import (
    get_store,
)

logger = get_logger(__name__)


fetch_api = Blueprint("fetch_api", __name__, url_prefix="/api")


@fetch_api.route("/tools", methods=["GET"])
@log_http_request
def get_tools() -> dict:
    """Return the list of available tool definitions."""
    tool_details = get_tools_details(user=g.get("authIdentifier"))
    return jsonify(
        {
            "tools": [
                {
                    "id": tool["dss_id"],
                    "name": tool["name"],
                    "description": tool["description"],
                    "short_description": tool["short_description"],
                    "type": tool["type"],
                    "configured": tool["configured"],
                    "error": tool.get("error"),
                }
                for tool in tool_details
            ]
        }
    )


@fetch_api.route("/images/<folder_id>/<path:file_path>", methods=["GET"])
@log_http_request
def get_managed_folder_image(folder_id, file_path):
    """
    Stream an image stored in a DSS managed folder.

    • <folder_id>   is the projectKey.managedFolderId
    • <file_path>   is the internal path coming from the front-end

    Spaces are encoded as “+” (quote_plus) in the URL; convert them back
    before asking DSS.  Also try the raw variant for safety.
    """
    try:
        project_key, short_id = folder_id.split(".", 1)
        file_path = unquote_plus(file_path)

        folder = dataiku.Folder(short_id, project_key=project_key)
        try:
            with folder.get_download_stream(file_path) as stream:
                data = stream.read()
            return send_file(
                BytesIO(data),
                download_name=file_path.rsplit("/", 1)[-1],
                mimetype="image/png",  # all RAG screenshots are PNG
                as_attachment=False,
            )
        except Exception as e:
            logger.exception(f"Exception {e}")

        # If we reach here nothing matched
        abort(404, f"File not found in managed folder: {file_path}")

    except Exception as e:
        abort(404, str(e))


# -------------------------------------------------------------------#
#  LLMs
# -------------------------------------------------------------------#


@fetch_api.route("/llms", methods=["GET"])
@log_http_request
def list_llms():
    id_to_name = _llm_friendly_names_by_purpose("GENERIC_COMPLETION")
    id_to_short_name = _llm_friendly_short_names_by_purpose("GENERIC_COMPLETION")
    default_llm = get_default_llm_id()
    return jsonify(
        {
            "llms": get_llms(),
            "embedding_llms": [get_default_embedding_llm()],
            "ah_llms": [
                {
                    "id": default_llm,
                    "name": id_to_name.get(default_llm, default_llm),
                    "shortName": id_to_short_name.get(default_llm, default_llm),
                }
            ],
        }
    )


# -------------------------------------------------------------------#
#  User preferences
# -------------------------------------------------------------------#
@fetch_api.route("/preferences", methods=["GET", "PUT"])
@log_http_request
def user_preferences():
    store = get_store()
    if request.method == "GET":
        prefs = store.get_preferences() or {}
        return jsonify(prefs)

    payload = request.get_json(force=True) or {}
    updated = store.update_preferences({"prefs": payload})
    return jsonify(updated.model_dump(mode="json") if updated else {})


@fetch_api.route("/plugins/<plugin_id>/info", methods=["GET"])
@log_http_request
def get_plugin_version(plugin_id):
    """Return the installed version of the plugin, or 404 if not installed."""
    client = dataiku.api_client()
    plugins = client.list_plugins()
    plugin = next((p for p in plugins if p.get("id") == plugin_id), None)

    if not plugin or not plugin.get("installed", True):
        logger.error(f"Plugin {plugin_id} is not installed.")
        return jsonify({"error": f'Plugin "{plugin_id}" is not installed.'}), 404

    version = None
    # Try installedDesc first, fallback to top-level version
    if plugin.get("installedDesc") and plugin["installedDesc"].get("desc"):
        version = plugin["installedDesc"]["desc"].get("version")
    elif plugin.get("version"):
        version = plugin["version"]

    if version:
        return jsonify({"version": version})
    else:
        logger.error(f"Plugin {plugin_id} has no version information.")
        return jsonify({"error": f'Plugin "{plugin_id}" version information is unavailable.'}), 404


@fetch_api.route("/workspaces", methods=["GET"])
@log_http_request
@with_impersonation
def list_workspaces():
    """
    Lists the workspaces accessible by the end-user
    """
    try:
        client = dataiku.api_client()
        results = []
        for item in client.list_workspaces():

            results.append(
                {
                    "id": item.get("workspaceKey"),
                    "name": item.get("displayName"),
                    "description": item.get("description", ""),
                }
            )
        return jsonify({"data": results}), 200

    except Exception as e:
        logger.exception(f"Error listing workspaces for the user: {e}")
        return jsonify({"error": str(e)}), 500


@fetch_api.route("/semantic-sources", methods=["POST"])
@log_http_request
@with_impersonation
def get_semantic_sources() -> List[SemanticDatasourceInfo]:
    """
    Fetches the semantic datasources (datasets, semantic models) from collections/workspaces
        accessible by an end-user.
    """
    semantic_sources_dict: Dict[str, SemanticDatasourceInfo] = {}
    user_id = g.get("authIdentifier")
    supported_connections = [conn.value for conn in SQLConnectionType]
    try:
        payload = request.get_json(force=True)
        workspaces_list = payload.get("workspaceIds")
        collections_list = payload.get("collectionIds")
        client = dataiku.api_client()

        for collection_id in collections_list:
            collection = client.get_data_collection(collection_id)
            for object in collection.list_objects():
                raw_object = object.get_raw()
                dataset_name = deepcopy(raw_object["id"])
                if (raw_object["type"] == "DATASET") and (
                    has_dataset_read_access(raw_object["projectKey"], dataset_name, user_id)
                ):
                    dataset_settings = object.get_as_dataset().get_settings().get_raw()
                    raw_object["name"] = dataset_name
                    instance_id = f'{raw_object["projectKey"]}.{dataset_name}'
                    raw_object["id"] = instance_id
                    connection_type = dataset_settings["type"]
                    if connection_type in supported_connections:
                        raw_object["description"] = dataset_settings.get("description", "")
                        raw_object["shortDesc"] = dataset_settings.get("shortDesc", "")
                        raw_object["lastModifiedOn"] = dataset_settings.get("versionTag", {}).get("lastModifiedOn", "")
                        raw_object["lastModifiedBy"] = (
                            dataset_settings.get("versionTag", {}).get("lastModifiedBy", {}).get("login", "")
                        )
                        # We use '.setdefault' and '.update' as 1 dataset can be shared into many collections
                        semantic_sources_dict.setdefault(instance_id, raw_object).update(raw_object)
                        semantic_sources_dict[instance_id].setdefault("origins", []).append(
                            {
                                "id": collection_id,
                                "name": collection.get_settings().display_name,
                                "type": "DATA_COLLECTION",
                            }
                        )
                        # We want the enduser_sql_execution to be set to the 'enduser':
                        semantic_sources_dict[instance_id]["endUserSQLExecution"] = DEFAULT_USER_SQL_EXECUTION

        for workspace_key in workspaces_list:
            workspace = client.get_workspace(workspace_key)
            for object in workspace.list_objects():
                raw_object: dict = object.data
                # 'reference' and 'subType' are not consistent depending on the workspace objects:
                type_candidate = raw_object.get("reference", {}).get("type", "")
                subtype_candidate = raw_object.get("subType", "")
                if (type_candidate == "DATASET") and (subtype_candidate in supported_connections):
                    # The top-level 'id' is specific to the workspace
                    del raw_object["id"]
                    # Flattening operation as the workspace objects 'projectKey', 'type' and 'id' are within the 'reference' key
                    raw_object.update(raw_object.get("reference", {}))
                    del raw_object["reference"]
                    dataset_name = deepcopy(raw_object["id"])
                    raw_object["name"] = dataset_name
                    instance_id = f'{raw_object["projectKey"]}.{dataset_name}'
                    raw_object["id"] = instance_id
                    last_modified_on = ""
                    last_modified_by = ""
                    if metadata := raw_object.get("metadata"):
                        last_modified_on = (
                            metadata.get("referencedObjectLastModifiedOn", "")
                            if metadata.get("referencedObjectLastModifiedOn", "")
                            else metadata.get("lastModifiedOn", "")
                        )
                        last_modified_by = (
                            metadata.get("referencedObjectCreatedBy", "")
                            if metadata.get("referencedObjectCreatedBy", "")
                            else metadata.get("lastModifiedBy", "")
                        )
                    if last_modified_on and (raw_object.get("lastModifiedOn") is None):
                        raw_object["lastModifiedOn"] = last_modified_on
                    if last_modified_by and (raw_object.get("lastModifiedBy") is None):
                        raw_object["lastModifiedBy"] = last_modified_by

                    if has_dataset_read_access(raw_object["projectKey"], dataset_name, user_id):
                        # We use '.setdefault' and '.update' as 1 dataset can be shared into many workspaces and collections
                        semantic_sources_dict.setdefault(instance_id, raw_object).update(raw_object)
                        semantic_sources_dict[instance_id].setdefault("origins", []).append(
                            {"id": workspace_key, "name": workspace.get_settings().display_name, "type": "WORKSPACE"}
                        )
                        # We want the enduser_sql_execution to be set to the 'enduser':
                        semantic_sources_dict[instance_id]["endUserSQLExecution"] = DEFAULT_USER_SQL_EXECUTION
        results: List[SemanticDatasourceInfo] = [
            normalize_to_semantic_datasource_info(value) for _, value in semantic_sources_dict.items()
        ]
        return jsonify({"data": results}), 200

    except Exception as e:
        logger.exception(f"Error fetching semantic sources for the user: {e}")
        return jsonify({"error": str(e)}), 500


@fetch_api.route("/collections", methods=["GET"])
@log_http_request
@with_impersonation
def list_collections():
    """
    Lists the collections  accessible by the end-user
    """
    try:
        client = dataiku.api_client()

        results = []
        for item in client.list_data_collections(as_type="dict"):

            results.append(
                {"id": item.get("id"), "name": item.get("displayName"), "description": item.get("description", "")}
            )
        return jsonify({"data": results}), 200

    except Exception as e:
        logger.exception(f"Error listing collections for the user: {e}")
        return jsonify({"error": str(e)}), 500


@fetch_api.route("/enterprise-asset-library/prompts", methods=["GET"])
@log_http_request
@with_impersonation
def list_enterprise_asset_library_prompts():
    """
    List prompts from the Enterprise Asset Library.
    Optionally filter by collection IDs via query parameter.

    Query params:
        collections: comma-separated list of collection IDs to filter by
    """
    try:
        client = dataiku.api_client()
        asset_library_client = client.get_enterprise_asset_library()

        # Get optional collection filter
        collections_param = request.args.get("collections")

        # If no collections are passed, return empty list
        if not collections_param:
            return jsonify({"prompts": []}), 200

        restrict_collections = [c.strip() for c in collections_param.split(",") if c.strip()]

        # Fetch prompts and collections
        prompts = asset_library_client.list_prompts(restrict_collections=restrict_collections)
        collections_list = asset_library_client.list_collections()

        # Create a lookup map: { collectionId: collectionName }
        collection_map = {c.get("id"): c.get("name") for c in collections_list}

        # Transform the response to a consistent format
        result = []
        for prompt in prompts:
            c_id = prompt.get("collectionId")

            result.append(
                {
                    "id": prompt.get("id"),
                    "name": prompt.get("name"),
                    "description": prompt.get("description", ""),
                    "content": prompt.get("content", ""),
                    "category": prompt.get("category", "Enterprise"),
                    "tags": prompt.get("tags", []),
                    "collectionId": c_id,
                    "collectionName": collection_map.get(c_id),
                    "isCreatedByDataiku": c_id == "DATAIKU",
                }
            )

        return jsonify({"prompts": result}), 200
    except AttributeError as e:
        logger.exception(
            "Enterprise Asset Library API not available - DSS version may be too old: %s",
            e,
        )
        return jsonify({"error": "Enterprise Asset Library is not available in this DSS version"}), 501
    except Exception as e:
        logger.exception(f"Error listing prompts from Enterprise Asset Library: {e}")
        return jsonify({"error": str(e)}), 500


@fetch_api.route("/enterprise-asset-library/collections", methods=["GET"])
@log_http_request
@with_impersonation
def list_enterprise_asset_library_collections():
    """
    List collections from the Enterprise Asset Library that the user has read access to.

    This implementation derives the collection list from available prompts to ensure
    both ID and Name are returned (as list_collections() might only return IDs).
    """
    try:
        client = dataiku.api_client()
        asset_library_client = client.get_enterprise_asset_library()

        result = asset_library_client.list_collections()

        return jsonify({"collections": result}), 200

    except AttributeError:
        # API method not available (DSS version too old)
        logger.exception("Enterprise Asset Library API not available - DSS version may be too old")
        return jsonify({"error": "Enterprise Asset Library is not available in this DSS version"}), 501
    except Exception as e:
        logger.exception(f"Error listing collections from Enterprise Asset Library: {e}")
        return jsonify({"error": str(e)}), 500


@fetch_api.route("/conversations/messages/<message_id>/trace", methods=["GET"])
@log_http_request
def get_message_trace(message_id):
    """
    Retrieve the trace data stored for the given message.
    If draft=true in request args, retrieves trace data from draft conversation.
    """
    try:
        store = get_store()

        # Check if draft parameter is present and true
        draft = request.args.get("draft", "").lower() == "true"

        if draft:
            # Get the agent ID from request args
            aid = request.args.get("agent_id")
            if not aid:
                logger.error("Agent ID is required for draft trace retrieval but not provided")
                return jsonify({"error": "Agent ID is required for draft trace retrieval"}), 400

            # Get draft conversation data
            convo = store.get_draft_conversation(aid)
            if not convo or not convo.convo:
                logger.error("Draft conversation not found for agent ID %s", aid)
                return jsonify({"error": "Draft conversation not found"}), 404
            # Find the message with matching ID and return its event log
            for m in convo.convo.get("messages", []):
                if m["id"] == message_id:
                    # Return only the trace field from the message
                    return jsonify({"messageId": message_id, "trace": m.get("trace", {})})

            # Message not found in draft conversation
            return jsonify({"messageId": message_id, "trace": {}}), 404
        else:
            # Original behavior - get trace from regular messages
            trace = store.get_message_trace(message_id)
            if not trace:
                logger.error("Trace not found for message ID %s", message_id)
                return jsonify({"error": "Trace data not found"}), 404
            return jsonify(trace)

    except Exception as e:
        logger.exception("Failed to retrieve trace for message %s: %s", message_id, e)
        return jsonify({"error": "Failed to retrieve trace"}), 500


@fetch_api.route("/upload", methods=["POST"])
@log_http_request
def upload_files():
    """Upload files to managed folder and store metadata in database."""
    try:
        # Get form data
        conv_id = request.form.get("conv_id")
        user_id = g.get("authIdentifier")
        message_id = request.form.get("message_id")

        if not conv_id or not user_id:
            logger.error("Missing required parameters for upload: conv_id=%s, user_id=%s", conv_id, user_id)
            return jsonify({"error": "Missing required parameters: conv_id, user_id"}), 400
        # message_id is now optional!
        # Check if files were uploaded
        if "files" not in request.files:
            logger.error("No files part in the request for conversation %s", conv_id)
            return jsonify({"error": "No files provided"}), 400

        files = request.files.getlist("files")
        if not files or all(f.filename == "" for f in files):
            logger.error("No files provided for conversation %s", conv_id)
            return jsonify({"error": "No files selected"}), 400

        # Get managed folder
        folder_id = get_uploads_managedfolder_id()
        if not folder_id:
            logger.error("Managed folder not configured for uploads")
            return jsonify({"error": "Uploads managed folder not configured"}), 500
        client = dataiku.api_client()
        project = client.get_default_project()

        try:
            folder = project.get_managed_folder(folder_id)
        except Exception as e:
            logger.exception(f"Failed to access managed folder with ID {folder_id}: {e}")
            return jsonify({"error": f"Failed to access managed folder: {str(e)}"}), 500

        # Create hierarchical path
        base_path = f"inputs/{user_id}/{conv_id}"
        # Ensure hierarchy exists by creating a .keep file
        try:
            keep_path = f"{base_path}/.keep"
            folder.put_file(keep_path, b"")
        except Exception as e:
            logger.error(f"Could not create .keep file: {e}")

        attachments = []

        # Process each uploaded file
        for file in files:
            if file.filename == "":
                continue

            filename = file.filename
            # If no message_id, just use a temp seed for hash
            unique_seed = f"{message_id or str(uuid.uuid4())}:{filename}:{uuid.uuid4()}".encode("utf-8")
            hash_prefix = hashlib.sha256(unique_seed).hexdigest()[:16]
            stored_filename = f"{hash_prefix}_{filename}"
            file_path = f"{base_path}/{stored_filename}"
            try:
                # Upload file to managed folder
                folder.put_file(file_path, file.stream)
                # Get file size
                file.stream.seek(0, 2)  # Seek to end
                file_size = file.stream.tell()
                file.stream.seek(0)  # Reset to beginning
                # Create attachment metadata (id will be replaced with database ID after upsert)
                attachment = {
                    "id": f"temp_{uuid.uuid4()}",  # Temporary ID; will be replaced with database ID after upsert
                    "document_name": filename,
                    "size": file_size,
                    "type": file.content_type or "application/octet-stream",
                    "document_path": file_path,
                    "uploadStatus": "uploaded",
                }
                attachments.append(attachment)
                logger.info(f"Uploaded file: {filename} -> {file_path}")
            except Exception as e:
                logger.error(f"Failed to upload {filename}: {e}")
                return jsonify({"error": f"Failed to upload {filename}: {str(e)}"}), 500

        # Update database
        if attachments:
            store = get_store()
            try:
                # Always upsert derived_docs for each file (associate with conversation and user)
                # and update attachment IDs with the database IDs
                for attachment in attachments:
                    derived_doc = store.upsert_derived_document(
                        conv_id,
                        attachment["document_name"],
                        attachment["document_path"],
                        {"snapshots": [], "size": attachment["size"], "type": attachment["type"], "status": "uploaded"},
                    )
                    # Use database ID instead of UUID for consistent ID matching across frontend/backend
                    attachment["id"] = derived_doc.id
                # Only fill message_attachments if message_id is present
                if message_id:
                    attachments_json = json.dumps(attachments)
                    store.insert_or_update_message_attachments(message_id, attachments_json)
                logger.info(f"Stored {len(attachments)} uploads for conversation {conv_id} (message_id={message_id})")
            except Exception as e:
                logger.error(f"Failed to update database: {e}")
                return jsonify({"error": f"Files uploaded but database update failed: {str(e)}"}), 500

        return (
            jsonify(
                {
                    "success": True,
                    "attachments": attachments,
                    "messageId": message_id,
                    "message": f"Successfully uploaded {len(attachments)} files",
                }
            ),
            200,
        )

    except Exception as e:
        logger.error(f"Upload error: {e}")
        return jsonify({"error": f"Upload failed: {str(e)}"}), 500


@fetch_api.route("/uploads", methods=["DELETE"])
@log_http_request
def delete_upload():
    try:
        payload = request.get_json(force=True, silent=True) or {}
        source_path = payload.get("path") or payload.get("source_path")
        if not source_path:
            logger.error("Missing required parameter 'path' for delete upload")
            return jsonify({"error": "Missing required parameter: path"}), 400

        # Delete DB row first
        store = get_store()
        deleted = store.delete_derived_document_by_source_path(source_path)

        # Try deleting the file from managed folder
        try:
            folder_id = get_uploads_managedfolder_id()
            client = dataiku.api_client()
            project = client.get_default_project()
            folder = project.get_managed_folder(folder_id)
            # Best-effort remove
            if hasattr(folder, "delete_path"):
                folder.delete_path(source_path)
            elif hasattr(folder, "delete_file"):
                folder.delete_file(source_path)
        except Exception as e:
            # Log but don't fail the whole operation
            logger.exception(f"Failed to delete file {source_path} from managed folder: {e}")

        return jsonify({"success": True, "deleted": bool(deleted), "path": source_path}), 200
    except Exception as e:
        logger.exception(f"Failed to delete upload: {e}")
        return jsonify({"error": "Failed to delete upload"}), 500
