# File: backend/fetch_api.py

import hashlib
import json
import uuid
from copy import deepcopy
from datetime import datetime
from io import BytesIO
from urllib.parse import unquote, unquote_plus

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

from backend.config import (
    _llm_friendly_names_by_purpose,
    _llm_friendly_short_names_by_purpose,
    get_config,
    get_default_embedding_llm,
    get_default_llm_id,
    get_llms,
    get_tools_details,
    get_uploads_managedfolder_id,
)
from backend.constants import DRAFT_ZONE
from backend.indexing_monitor import wake_monitor
from backend.services.agent_assets import (
    delete_agent_project,
    ensure_ua_project,
    get_job_status,
    get_ua_project_folder,
    get_visual_agent_in_zone,
    launch_index_job,
    prepare_agent_tools,
    remove_document,
    update_visual_agent,
    upload_documents,
)
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 (
    get_agent_details,
    get_managed_folder_in_zone,
    get_ua_project,
    update_agent_name_in_zone,
    update_project_name,
)
from backend.utils.utils import (
    _agent_version_rows,
    get_file_size,
    get_store,
    get_user_and_ent_agents,
    make_json_serializable,
)

logger = get_logger(__name__)

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

publishing_service = None


def enrich_agents_with_tools(agents: list[dict]) -> None:
    tool_details = get_tools_details(user=g.get("authIdentifier"))
    tool_map = {t["dss_id"]: t for t in tool_details}

    def _to_ids(value):
        if not value:
            return []
        if isinstance(value[0], str):
            return value
        # assume list[dict]
        return [v.get("id") for v in value if isinstance(v, dict)]

    for agent in agents:
        raw_ids = _to_ids(agent.get("tools", []))
        agent["tools"] = raw_ids
        agent["toolDetails"] = [
            {
                "id": i,
                "name": tool_map.get(i, {}).get("name", "Unknown"),
                "description": tool_map.get(i, {}).get("description", ""),
            }
            for i in raw_ids
            if i in tool_map
        ]

        # same treatment for the published snapshot
        pv = agent.get("published_version")
        if isinstance(pv, dict):
            pv_ids = _to_ids(pv.get("tools", []))
            pv["tools"] = pv_ids
            pv["toolDetails"] = [
                {
                    "id": i,
                    "name": tool_map.get(i, {}).get("name", "Unknown"),
                    "description": tool_map.get(i, {}).get("description", ""),
                }
                for i in pv_ids
                if i in tool_map
            ]


# --- Agents CRUD (user-scoped via UserStore) ---

# ──────────────────────────────────────────────────────────────────────────────
#  Principal search  (users / groups)
# ──────────────────────────────────────────────────────────────────────────────


@fetch_api.route("/principals/search", methods=["GET"])
@log_http_request
def search_principals():
    """
    Query params:
      • q             (string, mandatory)
      • type          = 'user' | 'group' | ''  (optional filter)
      • excludeGroups (string, optional) - comma-separated list of group names to exclude
    Returns: [{ "principal": "alice", "displayName": "Alice Doe", "type":"user" }, …]
    Capped to 20 results.
    """
    term = (request.args.get("q") or "").lower()
    kind = (request.args.get("type") or "").lower()
    auth_identifier = g.get("authIdentifier")

    exclude_groups_param = request.args.get("excludeGroups") or ""
    excluded_groups = set(g.strip() for g in exclude_groups_param.split(",") if g.strip())

    enable_groups_restriction_for_sharing = get_config().get("enable_groups_restriction_for_sharing", False)

    if not term:
        logger.warning("Missing 'q' parameter in principals search")
        return jsonify({"error": "q parameter required"}), 400

    client = dataiku.api_client()
    hits = []

    filtered_users = []
    groups = []
    users = client.list_users_info()

    if enable_groups_restriction_for_sharing:
        # Fetch only the current user's groups
        user = client.get_user(auth_identifier)
        user_groups = user.get_info().groups

        # Filter out excluded groups
        if excluded_groups:
            user_groups = [g for g in user_groups if g not in excluded_groups]

        # Create a set of group names for efficient lookup
        group_names = set(user_groups)
        groups = [{"name": group_name} for group_name in user_groups]
        # Filter users to only those who share at least one group with the current user
        for u in users:
            user_data = u.get_raw()
            user_member_groups = set(user_data.get("groups", []))
            # Check if user belongs to any of our filtered groups
            if user_member_groups & group_names:
                filtered_users.append(u)
    else:
        filtered_users = users
        all_groups = client.list_groups()
        for g_data in all_groups:
            g_name = g_data.get("name")
            if excluded_groups and g_name in excluded_groups:
                continue
            groups.append({"name": g_name})

    if kind in ("", "user"):
        for u in filtered_users:  # only search among filtered users
            user = u.get_raw()
            login = user["login"].lower()
            display = user.get("displayName") or user["login"]
            if term in login or term in display.lower():
                hits.append({"principal": user["login"], "displayName": display, "type": "user"})
            if len(hits) >= 20:
                break

    if kind in ("", "group") and len(hits) < 20:
        for group in groups:
            gid = group["name"].lower()
            if term in gid:
                hits.append({"principal": group["name"], "displayName": group["name"], "type": "group"})
            if len(hits) >= 20:
                break

    return jsonify({"results": hits[:20]})


@fetch_api.route("/groups/search", methods=["GET"])
@log_http_request
def search_groups():
    """
    Returns all groups the current user belongs to.
    Returns: [{ "principal": "group1", "displayName": "Group 1", "type":"group" }, …]
    """
    auth_identifier = g.get("authIdentifier")
    client = dataiku.api_client()
    hits = []

    # Fetch only the current user's groups
    user = client.get_user(auth_identifier)
    user_groups = user.get_info().groups

    # Return groups in the same format as search_principals
    for group_name in user_groups:
        try:
            # Try to get display name from group definition
            group = client.get_group(group_name)
            group_def = group.get_definition()
            display_name = group_def.get("displayName") or group_name
        except Exception:
            # Fallback to group name if lookup fails
            display_name = group_name

        hits.append({"principal": group_name, "displayName": display_name, "type": "group"})

    return jsonify({"results": hits[:50]})


# ──────────────────────────────────────────────────────────────────────────────
#  Agent-sharing CRUD
# ──────────────────────────────────────────────────────────────────────────────


def _enrich_display_names(items: list[dict]) -> list[dict]:
    client = dataiku.api_client()
    enriched: list[dict] = []

    for it in items:
        dn = it["principal"]  # default fallback
        try:
            if it["type"] == "user":
                raw = client.get_user(it["principal"]).get_settings().get_raw()
                dn = raw.get("displayName") or raw.get("fullName") or dn
            else:  # group
                raw = client.get_group(it["principal"]).get_definition()
                dn = raw.get("displayName") or raw.get("name") or dn
        except Exception:
            # ignore lookup errors – keep fallback
            pass

        enriched.append({**it, "displayName": dn})
    return enriched


@fetch_api.route("/agents/<agent_id>/shares", methods=["GET", "PUT"])
@log_http_request
def agent_shares(agent_id):
    store = get_store()
    logger.info("%s /agents/%s/shares", request.method, agent_id)

    agent = store.get_agent(agent_id)
    if not agent or not agent.get("published_version"):
        return jsonify({"error": "Agent must be published before sharing"}), 403

    if request.method == "GET":
        shares = store.get_agent_shares(agent_id)
        return jsonify({"shares": _enrich_display_names(shares)})

    # PUT → replace full list
    payload = request.get_json(force=True) or {}
    principals_list = payload.get("shares") or []
    if not isinstance(principals_list, list) or any("principal" not in s or "type" not in s for s in principals_list):
        return jsonify({"error": "Body must be a list of {principal,type}"}), 400

    try:
        old_shares = store.get_agent_shares(agent_id)  # ← fetch the “before” list
        old_principals = {s["principal"] for s in old_shares}
        store.replace_agent_shares(agent_id, principals_list)
        logger.info("Replaced shares for agent %s: %r", agent_id, principals_list)
        agent = store.get_agent(agent_id)
        agent_name = agent.get("name", agent_id)
        new_count = len(store.get_agent_shares(agent_id))
        for share in principals_list:
            principal = share["principal"]
            logger.info(f"emitting access granted notification to room : user:{principal}")
            socketio.emit(
                "agent:access_granted",
                {"agentId": agent_id, "agentName": agent_name, "shareCount": new_count},
                room=f"user:{principal}",
            )
        # 3) you may also want to notify users who were removed:
        old_principals = {s["principal"] for s in old_shares}
        new_principals = {s["principal"] for s in principals_list}
        for revoked in old_principals - new_principals:
            logger.info(f"emitting agent:access_revoked notification to room : user:{revoked}")
            socketio.emit(
                "agent:access_revoked",
                {"agentId": agent_id, "agentName": agent_name, "shareCount": new_count},
                room=f"user:{revoked}",
            )
    except PermissionError:
        logger.warning("Unauthorized share replace attempt on %s by %s", agent_id, g.authIdentifier)
        return jsonify({"error": "Only the owner can modify sharing"}), 403

    return jsonify({"shares": store.get_agent_shares(agent_id)})


@fetch_api.route("/agents", methods=["GET"])
@log_http_request
def get_agents():
    logger.info("Listing all agents for user")
    result = get_user_and_ent_agents()

    # Enrich all user agents with tool details
    # enrich_agents_with_tools(result["userAgents"])

    return jsonify(result)


def update_draft_agent(agent: dict, tools_explicitly_sent: bool) -> dict:
    store = get_store()
    logger.info(f"updating agent {agent['id']}")
    project = get_ua_project(agent)
    vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
    # Update visual agent in DRAFT zone based on the form data
    if tools_explicitly_sent:
        tools_ids = prepare_agent_tools(project, agent, DRAFT_ZONE)
    else:
        tools = get_agent_details(vis_agent)["tools"]
        tools_ids = prepare_agent_tools(project, {**agent, "tools": tools}, DRAFT_ZONE)

    update_visual_agent(
        agent=vis_agent,
        tools_ids=tools_ids,
        llm_id=agent.get("llmid"),
        prompt=agent.get("system_prompt", ""),
    )
    saved = store.update_agent(agent["id"], agent)
    return saved


@fetch_api.route("/agents", methods=["POST"])
@log_http_request
def create_or_update_agent():
    store = get_store()
    # Check if the form submission explicitly includes tool data
    # This could be either selected tools OR an empty indicator
    tools_explicitly_sent = False
    # ───────────── multipart (draft autosave or doc upload) ──────────────
    if request.content_type and request.content_type.startswith("multipart/form-data"):
        form = request.form
        files = request.files.getlist("documents")  # may be empty on autosave

        # ---------------------------------------------------------------
        #  **Ignore** completely empty heart-beats
        #  • no id      (brand-new draft)
        #  • no name
        #  • no description
        #  • no files / tools
        # ---------------------------------------------------------------
        nothing_to_save = (
            not form.get("id")
            and not form.get("name")
            and not form.get("description")
            and not form.get("systemPrompt")
            and not form.get("llmid")
            # and not form.get("embedding_llm_id")
            and not form.get("kb_description")
            and not files
            and all(not t for t in form.getlist("tools") or [])
        )
        if nothing_to_save:
            # Tell the caller it was an intentional no–op
            return ("", 204)

        # Pull existing row if it exists
        existing = store.get_agent(form.get("id")) if form.get("id") else {}

        # Check if this is an edit of an active agent
        is_editing_active = bool(existing and existing.get("published_version"))

        # Handle tools field - multipart forms have a special behavior:
        # 1. If no tools are selected, the field is not sent at all
        # 2. If tools are selected, they come as multiple values with the same key
        # 3. We need to check if ANY tool-related field was sent to know if it was intentional

        # Check for any field that indicates tools were part of the form submission
        # This includes checking for a hidden field or any tool-related field
        for key in form.keys():
            if key == "tools" or key == "tools[]" or key.startswith("tool_"):
                tools_explicitly_sent = True
                break

        # Also check if there's a special marker field (you may need to add this in frontend)
        if "tools_submitted" in form:
            tools_explicitly_sent = True

        if tools_explicitly_sent:
            # Tools were part of the form submission (even if empty)
            tool_ids = [t for t in form.getlist("tools") if t]
        else:
            # Tools were not part of this form submission - preserve existing
            tool_ids = existing.get("tools", []) if existing else []

        # Helper function to clean arrays - frontend sends [''] instead of [] for empty arrays
        def clean_array(arr):
            """Remove empty strings from array. If array contains only empty strings, return empty array."""
            cleaned = [item for item in arr if item and item.strip()]
            return cleaned

        # Handle sample questions - get list from form
        # If sample_questions field is present in form, use it; otherwise preserve existing
        sample_questions = []
        if "sample_questions" in form:
            sample_questions = clean_array(form.getlist("sample_questions"))
        elif existing:
            sample_questions = existing.get("sample_questions", [])

        agent = {
            "id": form.get("id") or str(uuid.uuid4()),
            "name": form.get("name", existing.get("name", "") if existing else ""),
            "description": form.get("description", existing.get("description", "") if existing else ""),
            "system_prompt": form.get("systemPrompt")
            or form.get("system_prompt", existing.get("system_prompt", "") if existing else ""),
            "tools": tool_ids,
            "llmid": form.get("llmid") or (existing.get("llmid", "") if existing else ""),
            "kb_description": form.get("kb_description", existing.get("kb_description", "") if existing else ""),
            "sample_questions": sample_questions,
            # "short_example_queries": short_example_queries,
        }
        # ----------------------------------------------------------------
        # Preserve / create the frozen snapshot that everybody else sees
        # ----------------------------------------------------------------
        if is_editing_active and not existing.get("published_version"):
            agent["published_version"] = {
                k: existing.get(k)
                for k in (
                    "name",
                    "description",
                    "system_prompt",
                    "tools",
                    "llmid",
                    # "embedding_llm_id",
                    "kb_description",
                    "sample_questions",
                )
            }
            agent["published_at"] = existing.get("published_at")
        elif existing and existing.get("published_version"):
            agent["published_version"] = existing["published_version"]
            agent["published_at"] = existing.get("published_at")

        # ── documents & indexing ─────────────────────────────────────────
        if files:
            # ── prevent concurrent indexing ───────────────────────────────
            existing_blob = agent.get("indexing") or {}
            if existing_blob.get("status") in ("pending", "running"):
                return jsonify({"error": "Indexing already in progress", "indexing": existing_blob}), 409

            # now upload / merge
            new_entries = [{"name": f.filename, "active": False, "size": get_file_size(f)} for f in files]
            docs = (existing.get("documents", []) if existing else []) + new_entries
            agent_with_doc = deepcopy(agent)
            agent_with_doc["documents"] = docs
            upload_documents(agent_with_doc, files)

            agent["documents"] = docs
            agent["indexing"] = None  # indexing_blob
        else:
            # autosave: keep existing docs / indexing untouched
            agent["documents"] = existing.get("documents", []) if existing else []
            agent["indexing"] = existing.get("indexing") if existing else None

    # ───────────── pure JSON (Edit Agent modal) ──────────────
    else:
        payload = request.get_json(force=True)
        logger.info("JSON POST /agents %s", payload.get("id"))
        missing = [f for f in ("id", "name") if not payload.get(f)]
        if missing:
            return jsonify({"error": f"Missing field(s): {', '.join(missing)}"}), 400

        existing = store.get_agent(payload["id"]) or {}
        agent = existing | payload  # shallow merge
        agent["documents"] = existing.get("documents", [])
        agent["indexing"] = existing.get("indexing")

    agent["owner"] = g.get("authIdentifier", "unknown")

    # Check if agent exists in DB
    agent_exists = store.agent_exists(agent["id"])

    saved = None
    try:
        if agent_exists:
            # Existing row → partial update keeps sharing intact
            logger.info(f"updating agent {agent['id']}")
            saved = update_draft_agent(agent, tools_explicitly_sent)

            if saved is None:
                # User doesn't own this agent
                return jsonify({"error": "Unauthorized to update this agent"}), 403
            if agent.get("name") != existing.get("name"):
                # Name changed - update project name too
                project = get_ua_project(agent)
                if project:
                    update_project_name(project, agent["name"])
                    update_agent_name_in_zone(project, DRAFT_ZONE, agent["name"])
        else:
            logger.info(f"creating agent {agent['id']}")
            # Create a project as soon as we have a new agent with name
            owner = g.get("run_as_login", "unknown")
            project = ensure_ua_project(agent, owner)
            if not project:
                return jsonify({"error": "Failed to create project for agent"}), 500
            agent["id"] = project.project_key
            saved = store.create_agent(agent)
            if saved is None:
                return jsonify({"error": "Failed to create agent"}), 500

    except ValueError as e:
        # Handle race condition where agent was created between our check and create attempt
        logger.warning(f"Race condition detected for agent {agent['id']}: {e}")
        # Try to update instead
        saved = update_draft_agent(agent, tools_explicitly_sent)
        if saved is None:
            return jsonify({"error": "Agent was created by another user"}), 403
    except Exception as e:
        logger.error(f"Unexpected error saving agent {agent['id']}: {e}")
        return jsonify({"error": "Failed to save agent"}), 500

    # Extra safety check
    if saved is None:
        logger.error(f"Saved is None for agent {agent['id']} - this shouldn't happen")
        return jsonify({"error": "Internal error: agent save failed"}), 500

    shares = store.get_agent_shares(saved["id"])
    share_count = len([s for s in shares if s["type"] == "user"])

    # build the "version rows" once
    version_rows = _agent_version_rows(saved, share_count, g.authIdentifier)

    # owner
    for row in version_rows:
        socketio.emit("agent:updated", row, room=f"user:{saved['owner']}")

    # every explicitly-shared user
    for s in shares:
        if s["type"] == "user":
            for row in version_rows:
                socketio.emit("agent:updated", {**row, "isShared": True}, room=f"user:{s['principal']}")
    return jsonify(saved), 200


@fetch_api.route("/agents/<agent_id>/publish", methods=["POST"])
@log_http_request
def publish_agent(agent_id):
    """Start the agent publishing process."""
    publishing_service = current_app.config.get("PUBLISHING_SERVICE")

    try:
        result = publishing_service.start_publish(agent_id, g.authIdentifier)
        return jsonify(result), 200
    except PermissionError as e:
        return jsonify({"error": str(e)}), 403
    except ValueError as e:
        return jsonify({"error": str(e)}), 404


@fetch_api.route("/agents/<agent_id>/publishing_status", methods=["GET"])
@log_http_request
def get_publishing_status(agent_id):
    """Get the current publishing status of an agent."""
    agent = get_store().get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    return jsonify(
        {
            "publishing_status": agent.get("publishing_status"),
            "published_at": agent.get("published_at"),
            "has_published_version": bool(agent.get("published_version")),
        }
    )


@fetch_api.route("/health/publishing", methods=["GET"])
@log_http_request
def publishing_health():
    """Health check endpoint for publishing service."""
    publishing_service = current_app.config.get("PUBLISHING_SERVICE")
    if not publishing_service:
        return jsonify({"status": "error", "message": "Service not initialized"}), 503

    stats = publishing_service.get_stats()
    status = "healthy" if stats["monitor_running"] else "unhealthy"

    return jsonify({"status": status, "stats": stats, "timestamp": datetime.utcnow().isoformat()})


@fetch_api.route("/agents/<agent_id>", methods=["GET"])
@log_http_request
def get_agent(agent_id):
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        logger.warning("Agent %s not found", agent_id)
        return jsonify({"error": "Agent not found"}), 404

    # Enrich with tool details
    enrich_agents_with_tools([agent])

    sc = store.get_share_counts([agent_id]).get(agent_id, 0)
    agent["shareCount"] = sc
    return jsonify(agent)


@fetch_api.route("/agents/<agent_id>", methods=["DELETE"])
@log_http_request
def delete_agent(agent_id):
    store = get_store()

    # Get agent details and shares before deletion
    agent = store.get_agent(agent_id)
    if not agent:
        logger.warning("DELETE missing agent %s", agent_id)
        return jsonify({"error": "Agent not found"}), 404

    agent_name = agent.get("name", agent_id)
    owner = agent.get("owner")

    # Get list of shared users before deletion
    shares = store.get_agent_shares(agent_id)
    shared_users = [s["principal"] for s in shares if s["type"] == "user"]

    # Perform deletion
    success = store.delete_agent(agent_id)
    if not success:
        logger.warning("DELETE missing agent %s", agent_id)
        return jsonify({"error": "Agent not found"}), 404

    delete_agent_project(agent_id)
    logger.info("Deleted agent %s and its project", agent_id)

    # Emit delete notifications
    # Notify owner
    socketio.emit(
        "agent:deleted",
        {"agentId": agent_id, "agentName": agent_name, "deletedBy": g.authIdentifier},
        room=f"user:{owner}",
    )

    # Notify all shared users
    for user in shared_users:
        socketio.emit(
            "agent:deleted",
            {"agentId": agent_id, "agentName": agent_name, "deletedBy": g.authIdentifier},
            room=f"user:{user}",
        )

    return jsonify({"message": "Agent deleted", "agentId": agent_id})


# ---------------------------------------------------------------------------#
# Document upload / delete (trigger indexing)
# ---------------------------------------------------------------------------#
@fetch_api.route("/agents/<agent_id>/documents", methods=["POST"])
@log_http_request
def upload_agent_document(agent_id):
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    files = request.files.getlist("documents")
    if not files:
        return jsonify({"error": "No files provided"}), 400

    new_entries = [{"name": f.filename, "active": False, "size": get_file_size(f)} for f in files]
    docs = agent.get("documents", []) + new_entries
    agent_with_doc = deepcopy(agent)
    agent_with_doc["documents"] = docs
    upload_documents(agent_with_doc, files)
    store.update_agent(agent_id, {"documents": docs, "indexing": None})
    return jsonify({"documents": docs}), 200
    # indexing_blob = {
    #     "status": "pending",
    #     "jobId": launch_index_job(agent),
    #     "updatedAt": datetime.utcnow().isoformat(),
    # }
    # wake_monitor()
    # return jsonify({"documents": docs, "indexing": indexing_blob}), 200


@fetch_api.route("/agents/<agent_id>/documents/index", methods=["POST"])
@log_http_request
def index_agent_document(agent_id):
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    docs = agent.get("documents", [])
    if not docs:
        return jsonify({"error": "No documents to index"}), 400
    indexing_blob = {
        "status": "pending",
        "jobId": launch_index_job(agent),
        "updatedAt": datetime.utcnow().isoformat(),
    }
    store.update_agent(agent_id, {"documents": docs, "indexing": indexing_blob})
    wake_monitor()
    return jsonify({"documents": docs, "indexing": indexing_blob}), 200


@fetch_api.route("/agents/<agent_id>/documents/<filename>", methods=["DELETE"])
@log_http_request
def delete_agent_document(agent_id, filename):
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    filename = unquote(filename)
    docs = agent.get("documents", [])
    if not any(d["name"].lower() == filename.lower() for d in docs):
        return jsonify({"error": "Document not found"}), 404

    # ── prevent concurrent indexing ───────────────────────────────
    existing_blob = agent.get("indexing") or {}
    if existing_blob.get("status") in ("pending", "running"):
        return jsonify({"error": "Indexing already in progress", "indexing": existing_blob}), 409

    # mark entry as pending deletion
    should_index = False
    indexing_blob = None
    for d in docs:
        if d["name"].lower() == filename.lower():
            d["deletePending"] = True
            if d.get("active"):
                should_index = True
            d["active"] = False

    # start job *after* marking
    remove_document(agent, filename)  # still delete physical file
    if should_index:
        logger.info("Documents deleted, launching re-indexing job for agent %s", agent_id)
        indexing_blob = {
            "status": "pending",
            "jobId": launch_index_job(agent),
            "updatedAt": datetime.utcnow().isoformat(),
        }
        store.update_agent(agent_id, {"documents": docs, "indexing": indexing_blob})
        wake_monitor()
        return jsonify({"documents": docs, "indexing": indexing_blob}), 200
    else:
        filtered_docs = [d for d in docs if d["name"].lower() != filename.lower() or not d.get("deletePending")]
        store.update_agent(agent_id, {"documents": filtered_docs, "indexing": indexing_blob})
        return jsonify({"documents": filtered_docs, "indexing": indexing_blob}), 200


@fetch_api.route("/agents/<agent_id>/documents", methods=["DELETE"])
@log_http_request
def delete_multiple_agent_documents(agent_id):
    """Bulk‐delete documents for one agent."""
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    payload = request.get_json(force=True)
    filenames = payload.get("filenames")
    if not filenames or not isinstance(filenames, list):
        return jsonify({"error": "Missing or invalid filenames"}), 400

    # Prevent concurrent indexing
    existing_blob = agent.get("indexing") or {}
    if existing_blob.get("status") in ("pending", "running"):
        return jsonify({"error": "Indexing already in progress", "indexing": existing_blob}), 409

    # Mark each for deletion
    docs = agent.get("documents", [])
    should_index = False
    for fn in filenames:
        for d in docs:
            if d["name"].lower() == fn.lower():
                d["deletePending"] = True
                if d.get("active"):
                    should_index = True
                d["active"] = False
                remove_document(agent, fn)
    indexing_blob = None
    if should_index:
        logger.info("Documents deleted, launching re-indexing job for agent %s", agent_id)
        # Kick off a single indexing job
        indexing_blob = {
            "status": "pending",
            "jobId": launch_index_job(agent),
            "updatedAt": datetime.utcnow().isoformat(),
        }
        store.update_agent(agent_id, {"documents": docs, "indexing": indexing_blob})
        wake_monitor()
        return jsonify({"documents": docs, "indexing": indexing_blob}), 200
    else:
        lower_filenames = [fn.lower() for fn in filenames]
        filtered_docs = [d for d in docs if d["name"].lower() not in lower_filenames or not d.get("deletePending")]
        store.update_agent(agent_id, {"documents": filtered_docs, "indexing": indexing_blob})
        return jsonify({"documents": filtered_docs, "indexing": indexing_blob}), 200


@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"],
                    "type": tool["type"],
                    "configured": tool["configured"],
                    "error": tool.get("error"),
                }
                for tool in tool_details
            ]
        }
    )


# --- Index status polling -----------------------------------------------


# ---------------------------------------------------------------------------#
# Index-status polling
# ---------------------------------------------------------------------------#
@fetch_api.route("/agents/<agent_id>/index_status", methods=["GET"])
@log_http_request
def get_index_status(agent_id):
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    blob = agent.get("indexing") or {}
    job_id = blob.get("jobId")
    if job_id:
        live = get_job_status(job_id)  # mock returns dict or None
        if live and live["status"] != blob["status"]:
            blob.update(live)
            blob["updatedAt"] = datetime.utcnow().isoformat()
            # If success → mark all docs active:true
            if blob["status"] == "success":
                docs = agent.get("documents", [])
                docs = [
                    {**d, "active": True}
                    for d in docs
                    if not d.get("deletePending")  # drop entries marked for deletion
                ]
                store.update_agent(agent_id, {"indexing": blob, "documents": docs})
                blob = blob.copy()  # refreshed for response
    return jsonify({"indexing": blob})


# backend/fetch_api.py
@fetch_api.route("/agents/index_status_bulk", methods=["POST"])
@log_http_request
def index_status_bulk():
    """
    Body: { "agentIds": ["id1", "id2", …] }
    Returns the *current* indexing blob (and, if finished, the documents list)
    straight from SQLite—no live DSS calls.
    """
    agent_ids = request.get_json(force=True).get("agentIds", [])
    store = get_store()
    results: dict[str, dict] = {}

    for aid in agent_ids:
        agent = store.get_agent(aid)
        if not agent:
            continue

        blob = agent.get("indexing") or {}
        result = {"indexing": blob}

        # only include documents if the job has completed (success or failure)
        if blob.get("status") in ("success", "failure"):
            result["documents"] = agent.get("documents", [])

        results[aid] = result

    return jsonify({"results": results})


# ---------------------------------------------------------------------------#
#  DISCARD CHANGES  – POST /api/agents/<id>/discard
# ---------------------------------------------------------------------------#


@fetch_api.route("/agents/<agent_id>/discard", methods=["POST"])
@log_http_request
def discard_agent_changes(agent_id: str):
    """Restore draft to the last published snapshot, documents included."""
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    pv = agent.get("published_version")
    if not pv:
        return jsonify({"error": "Agent has no published version"}), 400

    # ── 1 · abort any running indexing job ────────────────────────────────
    idx = agent.get("indexing") or {}
    if idx.get("status") in ("pending", "running") and idx.get("jobId"):
        try:
            pj_key = agent_id
            dataiku.api_client().get_project(pj_key).get_job(idx["jobId"]).abort()
        except Exception:
            logger.exception("Could not abort job %s for agent %s", idx.get("jobId"), agent_id)

    # ── 2 · delete draft-only files & those flagged deletePending ──────────
    draft_docs = agent.get("documents", [])
    pv_doc_names = {d["name"] for d in pv.get("documents", [])}
    files_changed = False

    for doc in draft_docs:
        name = doc.get("name")
        if name and (name not in pv_doc_names or doc.get("deletePending")):
            try:
                remove_document(agent, name)
                files_changed = True
            except Exception:
                logger.exception("Discard failed to delete %s for %s", name, agent_id)

    # ── 3 · restore any published file missing in draft ───────────────────
    try:
        project, draft_folder = get_ua_project_folder(agent)
        published_folder = get_managed_folder_in_zone(project, "published_documents", "published")

        if published_folder:
            draft_files = {it["path"] for it in draft_folder.list_contents()["items"] if it["type"] == "FILE"}

            for pv_doc in pv.get("documents", []):
                name = pv_doc["name"]
                if name not in draft_files:
                    try:
                        content = published_folder.get_file(name).content
                        draft_folder.put_file(name, content)
                        files_changed = True
                    except Exception:
                        logger.exception("Failed to restore %s for %s", name, agent_id)
    except Exception:
        logger.exception("Restore scan failed for agent %s", agent_id)

    # ── 4 · persist clean snapshot identical to published version ─────────
    now = datetime.utcnow().isoformat()
    restored_docs = [{"name": d["name"], "active": True} for d in pv.get("documents", [])]

    payload = {
        k: pv.get(k)
        for k in (
            "name",  # TODO what to do with name as this is related to project name ?
            "description",
            "sample_questions",
        )
    }
    payload.update({"documents": restored_docs, "indexing": None, "updatedAt": now})
    # Handle visual agent fields
    project = get_ua_project({"id": agent_id})
    draft_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
    prompt = pv.get("system_prompt", "")
    llm_id = pv.get("llmid", "")
    tools_ids = prepare_agent_tools(project, pv, DRAFT_ZONE)
    update_visual_agent(draft_agent, tools_ids=tools_ids, prompt=prompt, llm_id=llm_id)
    agent = store.update_agent(agent_id, payload)

    # ── 5 · (re)launch indexing if doc set changed ────────────────────────
    if files_changed:
        try:
            job_id = launch_index_job(agent, zone=DRAFT_ZONE)
            store.update_agent(
                agent_id,
                {"indexing": {"jobId": job_id, "status": "pending", "updatedAt": now}},
            )
            wake_monitor()
        except Exception:
            logger.exception("Failed to launch re-index job for %s", agent_id)

    # ── 6 · socket notifications ──────────────────────────────────────────
    ### TODO WHY EMIT HERE ??
    shares = store.get_agent_shares(agent_id)
    share_count = len([s for s in shares if s["type"] == "user"])

    # owner
    socketio.emit(
        "agent:updated",
        {
            "agentId": agent_id,
            "status": "active",
            "agentName": agent["name"],
            "isShared": False,
            "shareCount": share_count,
            "hasUnpublishedChanges": False,
        },
        room=f"user:{agent['owner']}",
    )
    # -- 7. Update agent name
    if agent.get("name") != pv.get("name"):
        update_project_name(project, pv["name"])
        update_agent_name_in_zone(project, DRAFT_ZONE, pv["name"])
    return jsonify(make_json_serializable(store.get_agent(agent_id))), 200


@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":
        return jsonify(store.get_preferences())

    payload = request.get_json(force=True) or {}
    updated = store.update_preferences(payload)
    return jsonify(updated)


@fetch_api.route("/agents/<agent_id>/favorite", methods=["POST"])
@log_http_request
def mark_agent_as_favorite(agent_id):
    """Mark or unmark an agent as a favorite."""
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return jsonify({"error": "Agent not found"}), 404

    payload = request.get_json(force=True)
    is_favorite = payload.get("favorite")
    if not isinstance(is_favorite, bool):
        return jsonify({"error": "Invalid 'favorite' value, must be true or false"}), 400

    # Update the favorite preference in the preferences table
    prefs = store.get_preferences()

    # Ensure the 'favorite' key exists, defaulting to an empty dict if not.
    favorite_prefs = prefs.setdefault("favorite", {})
    favorite_prefs[agent_id] = is_favorite

    store.update_preferences(prefs)

    return jsonify({"message": "Favorite status updated", "agentId": agent_id, "favorite": is_favorite})


@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):
        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:
        return jsonify({"error": f'Plugin "{plugin_id}" version information is unavailable.'}), 404


@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:
                return jsonify({"error": "Agent ID is required for draft trace retrieval"}), 400

            # Get draft conversation data
            convo = store.get_draft_conversation(aid) or {}

            # Find the message with matching ID and return its event log
            for m in 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
            message = store.get_message_trace(message_id)
            if not message:
                return jsonify({"error": "Trace data not found"}), 404
            return jsonify(message)

    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 = request.form.get("user_id")
        message_id = request.form.get("message_id")

        if not conv_id or not 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:
            return jsonify({"error": "No files provided"}), 400

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

        # Get managed folder
        folder_id = get_uploads_managedfolder_id()
        if not folder_id:
            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:
            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.warning(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
                attachment = {
                    "id": str(uuid.uuid4()),
                    "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:
                # 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)
                # Always upsert derived_docs for each file (associate with conversation and user)
                for attachment in attachments:
                    store.upsert_derived_document(
                        conv_id,
                        attachment["document_name"],
                        attachment["document_path"],
                        {"snapshots": [], "size": attachment["size"], "type": attachment["type"], "status": "uploaded"},
                    )
                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:
            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.warning(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
