from __future__ import annotations

from datetime import datetime
from urllib.parse import unquote

from backend.config import get_my_agents_max_upload_size_bytes
from backend.constants import DRAFT_ZONE, CloudConnections
from backend.indexing_monitor import wake_monitor
from backend.services.agent_assets import (
    compute_indexing_status,
    ensure_merged_documents_assets,
    ensure_sharepoint_assets,
    get_ua_project_folder,
    launch_index_job,
    prepare_agent_tools,
    remove_connection,
    remove_document,
    run_cloud_sync_scenarios,
    upload_documents,
)
from backend.utils.docs_utils import (
    SHAREPOINT_PATH_TYPE_FILE,
    SHAREPOINT_PATH_TYPE_FOLDER,
    doc_name_equals,
    is_connection_folder_doc,
    is_delete_pending,
    is_local_doc,
    make_sharepoint_doc_entry,
    normalize_sharepoint_items,
)
from backend.utils.logging_utils import get_logger
from backend.utils.project_utils import (
    build_sharepoint_managed_folder_name,
    create_sharepoint_managed_folder,
    get_kb_tool_in_zone,
    get_managed_folder_in_zone,
    get_managed_folder_size,
    get_outputs_from_embed_docs_recipe,
    get_ua_project,
    get_visual_agent_in_zone,
    update_visual_agent,
)
from backend.utils.utils import get_doc_name, get_file_size, get_store

logger = get_logger(__name__)


def attach_connection_folder_sizes(agent_id: str, docs: list | None) -> list:
    docs = docs or []
    if not docs:
        return docs

    enriched_docs = []
    for doc in docs:
        if isinstance(doc, dict) and is_connection_folder_doc(doc):
            if doc.get("pathType") == SHAREPOINT_PATH_TYPE_FILE:
                enriched_docs.append(doc)
                continue
            size = get_managed_folder_size(agent_id, doc.get("folderId", ""))
            enriched_docs.append({**doc, "size": size})
        else:
            enriched_docs.append(doc)
    return enriched_docs


def build_indexing_blob(
    agent_id: str,
    zone: str = DRAFT_ZONE,
    include_local: bool = True,
    include_cloud: bool = True,
    start_job: bool = True,
) -> dict:
    job_id = (
        launch_index_job(agent_id, zone=zone, include_local=include_local, include_cloud=include_cloud)
        if start_job
        else None
    )
    return {
        "status": "pending",
        "jobId": job_id,
        "updatedAt": datetime.utcnow().isoformat(),
    }


def _mark_document_for_deletion(doc) -> tuple[bool, str | None]:
    if is_connection_folder_doc(doc):
        was_active = bool(doc.get("active"))
        doc["deletePending"] = True
        doc["active"] = False
        return was_active, doc.get("connectionType", "")
    if isinstance(doc, dict):
        should_index = bool(doc.get("active"))
        doc["deletePending"] = True
        doc["active"] = False
        return should_index, None
    if isinstance(doc, str):
        return True, None
    return False, None


def _collect_documents_for_deletion(docs: list, filenames: list) -> tuple[bool, list[tuple[str, str]], set[str]]:
    should_index = False
    connection_folder_delete_targets = []
    local_filenames_to_delete = set()

    for fn in filenames:
        for doc in docs:
            if not doc_name_equals(doc, fn):
                continue
            doc_should_index, connection_type = _mark_document_for_deletion(doc)
            should_index = should_index or doc_should_index
            if connection_type is not None:
                connection_folder_delete_targets.append((connection_type, fn))
            else:
                local_filenames_to_delete.add(fn)

    return should_index, connection_folder_delete_targets, local_filenames_to_delete


def _build_deleted_documents_indexing_blob(
    agent_id: str, filtered_docs: list
) -> tuple[dict | None, dict | None, int | None]:
    scenario_runs = []
    cloud_docs = [d for d in filtered_docs if is_connection_folder_doc(d)]
    if cloud_docs:
        project = get_ua_project(agent_id)
        scenario_runs = run_cloud_sync_scenarios(project, cloud_docs, DRAFT_ZONE)
        if not scenario_runs:
            return None, {"error": "Failed to start cloud sync scenario"}, 500

    indexing_blob = build_indexing_blob(
        agent_id,
        include_local=not bool(scenario_runs),
        include_cloud=False,
        start_job=not bool(scenario_runs),
    )
    if scenario_runs:
        indexing_blob["scenarioRuns"] = scenario_runs
    return indexing_blob, None, None


def _persist_deleted_documents(
    store, agent_id: str, docs: list, filtered_docs: list, should_index: bool
) -> tuple[dict, int]:
    indexing_blob = None
    if should_index:
        indexing_blob, error_payload, status_code = _build_deleted_documents_indexing_blob(agent_id, filtered_docs)
        if error_payload:
            return error_payload, status_code or 500
        store.update_agent(agent_id, {"documents": docs, "indexing": indexing_blob})
        wake_monitor()
        return {"documents": docs, "indexing": indexing_blob}, 200

    store.update_agent(agent_id, {"documents": filtered_docs, "indexing": indexing_blob})
    return {"documents": filtered_docs, "indexing": indexing_blob}, 200


def activate_indexed_documents(docs: list | None) -> list:
    activated_docs = []
    for doc in docs or []:
        if is_delete_pending(doc):
            continue
        if isinstance(doc, str):
            activated_docs.append({"name": doc, "active": True})
        elif isinstance(doc, dict):
            activated_docs.append({**doc, "active": True})
        else:
            activated_docs.append(doc)
    return activated_docs


def refresh_indexing_blob(agent) -> tuple[dict, list]:
    store = get_store()
    blob = (agent.indexing or {}).copy()
    job_id = blob.get("jobId")
    scenario_runs = blob.get("scenarioRuns") or []
    docs = agent.documents or []

    if not job_id and not scenario_runs:
        return blob, docs

    live = compute_indexing_status(agent.id, job_id, scenario_runs)
    if not live or live.get("status") == blob.get("status"):
        return blob, docs

    blob.update(live)
    blob["updatedAt"] = datetime.utcnow().isoformat()
    payload = {"indexing": blob}

    if blob["status"] == "success":
        docs = activate_indexed_documents(docs)
        payload["documents"] = docs

    store.update_agent(agent.id, payload)
    return blob, docs


def get_documents_total_size_bytes(agent_id: str) -> int:
    try:
        store = get_store()
        agent = store.get_agent(agent_id)
        if not agent:
            return 0

        project = get_ua_project(agent_id)
        folder = get_managed_folder_in_zone(project, f"documents_{DRAFT_ZONE}", DRAFT_ZONE)
        local_total_size = get_managed_folder_size(agent_id, folder.id) if folder else 0
        connection_folder_ids = {
            doc.get("folderId")
            for doc in agent.documents or []
            if is_connection_folder_doc(doc) and not is_delete_pending(doc) and doc.get("folderId")
        }
        connection_total_size = sum(
            get_managed_folder_size(agent_id, folder_id) or 0 for folder_id in connection_folder_ids
        )
        return int((local_total_size or 0) + connection_total_size)
    except Exception as e:
        logger.warning("Could not compute total documents size for %s: %s", agent_id, e)
        return 0


def upload_agent_document(agent_id: str, files) -> tuple[dict, int]:
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return {"error": "Agent not found"}, 404

    if not files:
        return {"error": "No files provided"}, 400

    file_sizes = [get_file_size(f) for f in files]
    new_files_total_size = sum(file_sizes)
    current_docs_total_size = get_documents_total_size_bytes(agent_id)
    max_upload_size_bytes = get_my_agents_max_upload_size_bytes()

    if max_upload_size_bytes is not None and current_docs_total_size + new_files_total_size > max_upload_size_bytes:
        return (
            {
                "error": "max_upload_size_exceeded",
                "maxSizeBytes": max_upload_size_bytes,
                "currentSizeBytes": current_docs_total_size,
                "newFilesSizeBytes": new_files_total_size,
            },
            413,
        )

    new_entries = [{"name": f.filename, "active": False, "size": size} for f, size in zip(files, file_sizes)]
    docs = (agent.documents or []) + new_entries
    agent_with_doc = agent.model_dump()
    agent_with_doc["documents"] = docs
    upload_documents(agent_with_doc, files)
    store.update_agent(agent_id, {"documents": docs, "indexing": None})
    return {"documents": docs}, 200


def index_agent_documents(agent_id: str) -> tuple[dict, int]:
    from backend.services.guardrails_service import enforce_indexing_guardrails

    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return {"error": "Agent not found"}, 404

    project = get_ua_project(agent_id)
    ensure_merged_documents_assets(project, agent.model_dump(mode="json"), DRAFT_ZONE)

    docs = agent.documents or []
    local_docs = [d if isinstance(d, dict) else {"name": d, "active": True} for d in docs if is_local_doc(d)]
    cloud_docs = [d for d in docs if is_connection_folder_doc(d)]
    if not docs:
        return {"error": "No documents to index"}, 400

    if local_docs:
        error_dict, status_code = enforce_indexing_guardrails(agent_id, local_docs, store, remove_document)
        if error_dict:
            return error_dict, status_code

    scenario_runs = []
    if cloud_docs:
        scenario_runs = run_cloud_sync_scenarios(project, cloud_docs, DRAFT_ZONE)
        if not scenario_runs:
            return {"error": "Failed to start cloud sync scenario"}, 500

    indexing_blob = build_indexing_blob(
        agent_id,
        include_local=bool(local_docs) and not bool(scenario_runs),
        include_cloud=False,
        start_job=not bool(scenario_runs),
    )
    if scenario_runs:
        indexing_blob["scenarioRuns"] = scenario_runs
    store.update_agent(agent_id, {"documents": docs, "indexing": indexing_blob})
    wake_monitor()
    return {"documents": docs, "indexing": indexing_blob}, 200


def _build_sharepoint_documents(
    agent_id: str,
    project,
    connection: str,
    site: str,
    site_id: str,
    drive: str,
    drive_id: str,
    selected_items: list[dict],
) -> tuple[list[str], list[str], list[dict]]:
    folder_ids = []
    folder_names = []
    sharepoint_docs = []
    source_folders_by_path = {}

    for selected_item in selected_items:
        selected_path = selected_item["path"]
        path_type = selected_item.get("pathType") or SHAREPOINT_PATH_TYPE_FOLDER
        folder_path = selected_item["sourcePath"]

        if folder_path in source_folders_by_path:
            folder, folder_name = source_folders_by_path[folder_path]
        else:
            folder_name = build_sharepoint_managed_folder_name(
                DRAFT_ZONE,
                connection,
                site_id or site,
                drive_id or drive,
                folder_path,
            )
            folder = create_sharepoint_managed_folder(
                project, folder_name, connection, site, site_id, drive, drive_id, folder_path, DRAFT_ZONE
            )
            source_folders_by_path[folder_path] = (folder, folder_name)
            folder_ids.append(folder.id)
            folder_names.append(folder_name)

        if path_type == SHAREPOINT_PATH_TYPE_FILE:
            sharepoint_docs.append(
                make_sharepoint_doc_entry(
                    None,
                    connection,
                    site,
                    drive,
                    site_id,
                    drive_id,
                    selected_path,
                    0,
                    path_type=SHAREPOINT_PATH_TYPE_FILE,
                    source_folder_id=folder.id,
                    source_folder_path=folder_path,
                    relative_path=selected_item.get("relativePath"),
                )
            )
            continue

        sharepoint_docs.append(
            make_sharepoint_doc_entry(
                folder.id,
                connection,
                site,
                drive,
                site_id,
                drive_id,
                selected_path,
                get_managed_folder_size(agent_id, folder.id),
                path_type=SHAREPOINT_PATH_TYPE_FOLDER,
            )
        )

    return folder_ids, folder_names, sharepoint_docs


def _update_sharepoint_visual_agent(project, updated_agent: dict) -> None:
    if get_kb_tool_in_zone(project, DRAFT_ZONE):
        return

    vis_agent = get_visual_agent_in_zone(project, DRAFT_ZONE)
    if not vis_agent:
        return

    tools_ids = prepare_agent_tools(project, updated_agent, DRAFT_ZONE)
    update_visual_agent(
        agent=vis_agent,
        tools_ids=tools_ids,
        llm_id=updated_agent.get("llmid", ""),
        prompt=updated_agent.get("system_prompt", ""),
    )


def _parse_sharepoint_documents_request(payload: dict) -> tuple[dict | None, int | None, dict]:
    connection = (payload.get("connection") or "").strip()
    site = (payload.get("site") or "").strip()
    drive = (payload.get("drive") or "").strip()
    site_id = (payload.get("siteId") or "").strip()
    drive_id = (payload.get("driveId") or "").strip()
    items = payload.get("items") or []

    if not connection or not site or not drive:
        return {"error": "connection, site and drive are required"}, 400, {}
    if not isinstance(items, list):
        return {"error": "items must be a list"}, 400, {}

    selected_items = normalize_sharepoint_items(items)
    if not selected_items:
        return {"error": "At least one item is required"}, 400, {}

    return (
        None,
        None,
        {
            "connection": connection,
            "site": site,
            "drive": drive,
            "site_id": site_id,
            "drive_id": drive_id,
            "selected_items": selected_items,
        },
    )


def _load_agent_and_non_sharepoint_docs(agent_id: str) -> tuple[dict, list]:
    store = get_store()
    db_agent = store.get_agent(agent_id)
    if not db_agent:
        raise ValueError(f"Agent {agent_id} not found")

    existing_docs = db_agent.documents or []
    non_sharepoint_docs = [
        d
        for d in existing_docs
        if not (isinstance(d, dict) and d.get("connectionType") == CloudConnections.SHAREPOINT_ONLINE.value)
    ]
    return db_agent, non_sharepoint_docs


def _finalize_sharepoint_documents_configuration(
    agent_id: str,
    store,
    db_agent,
    project,
    non_sharepoint_docs: list,
    sharepoint_docs: list[dict],
) -> list:
    updated_docs = [*non_sharepoint_docs, *sharepoint_docs]
    updated_agent = {**db_agent.model_dump(mode="json"), "documents": updated_docs}
    recipe = ensure_sharepoint_assets(project, updated_agent, DRAFT_ZONE)
    _update_sharepoint_visual_agent(project, updated_agent)
    kb_id, imgs_id = get_outputs_from_embed_docs_recipe(recipe) if recipe else (None, None)
    for sharepoint_doc in sharepoint_docs:
        sharepoint_doc["kbId"] = kb_id
        sharepoint_doc["imgsFolderId"] = imgs_id
        sharepoint_doc["recipeId"] = recipe.id if recipe else None
    updated_docs = [*non_sharepoint_docs, *sharepoint_docs]
    store.update_agent(agent_id, {"documents": updated_docs})
    return updated_docs


def configure_sharepoint_documents_folder(agent_id: str, payload: dict) -> tuple[dict, int]:
    error_payload, status_code, request = _parse_sharepoint_documents_request(payload)
    if error_payload:
        return error_payload, status_code or 400

    try:
        project = get_ua_project(agent_id)
        store = get_store()
        db_agent, non_sharepoint_docs = _load_agent_and_non_sharepoint_docs(agent_id)

        folder_ids, folder_names, sharepoint_docs = _build_sharepoint_documents(
            agent_id,
            project,
            request["connection"],
            request["site"],
            request["site_id"],
            request["drive"],
            request["drive_id"],
            request["selected_items"],
        )
        updated_docs = _finalize_sharepoint_documents_configuration(
            agent_id,
            store,
            db_agent,
            project,
            non_sharepoint_docs,
            sharepoint_docs,
        )

        return (
            {
                "ok": True,
                "folderId": folder_ids[0],
                "folderIds": folder_ids,
                "folderName": folder_names[0],
                "folderNames": folder_names,
                "type": CloudConnections.SHAREPOINT_ONLINE.value,
                "documents": updated_docs,
            },
            200,
        )
    except Exception as e:
        logger.error(f"Failed to configure SharePoint folder for {agent_id}: {e}", exc_info=True)
        return {"error": str(e)}, 500


def delete_agent_document(agent_id: str, filename: str) -> tuple[dict, int]:
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return {"error": "Agent not found"}, 404

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

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

    if not any(doc_name_equals(d, filename) for d in docs):
        return {"error": "Document not found"}, 404

    should_index, connection_folder_delete_targets, _ = _collect_documents_for_deletion(docs, [filename])
    connection_folder_types_to_delete = {connection_type for connection_type, _ in connection_folder_delete_targets}

    if connection_folder_types_to_delete:
        for connection_type in connection_folder_types_to_delete:
            remove_connection(agent, connection_type, doc_name=filename)
    else:
        remove_document(agent, filename)

    filtered_docs = [d for d in docs if not doc_name_equals(d, filename) or not is_delete_pending(d)]
    return _persist_deleted_documents(store, agent_id, docs, filtered_docs, should_index)


def delete_multiple_agent_documents(agent_id: str, filenames: list) -> tuple[dict, int]:
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return {"error": "Agent not found"}, 404

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

    docs = agent.documents or []
    should_index, connection_folder_delete_targets, local_filenames_to_delete = _collect_documents_for_deletion(
        docs, filenames
    )

    for connection_type, fn in connection_folder_delete_targets:
        remove_connection(agent, connection_type, doc_name=fn)
    for local_fn in local_filenames_to_delete:
        remove_document(agent, local_fn)

    lower_filenames = [filename.lower() for filename in filenames]
    filtered_docs = [d for d in docs if get_doc_name(d).lower() not in lower_filenames or not is_delete_pending(d)]
    return _persist_deleted_documents(store, agent_id, docs, filtered_docs, should_index)


def get_index_status(agent_id: str) -> tuple[dict, int]:
    store = get_store()
    agent = store.get_agent(agent_id)
    if not agent:
        return {"error": "Agent not found"}, 404

    blob, _ = refresh_indexing_blob(agent)
    return {"indexing": blob}, 200


def index_status_bulk(agent_ids: list) -> dict:
    store = get_store()
    results: dict[str, dict] = {}

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

        blob, docs = refresh_indexing_blob(agent)
        result = {"indexing": blob}
        if blob.get("status") in ("success", "failure"):
            result["documents"] = attach_connection_folder_sizes(aid, docs)
        results[aid] = result

    return {"results": results}


def _remove_discarded_local_documents(agent, draft_docs: list, published_doc_names: set[str]) -> bool:
    files_changed = False

    for doc in draft_docs:
        name = get_doc_name(doc)
        if name and is_local_doc(doc) and (name not in published_doc_names or is_delete_pending(doc)):
            try:
                remove_document(agent, name)
                files_changed = True
            except Exception:
                logger.exception("Discard failed to delete %s for %s", name, agent.id)

    return files_changed


def _restore_missing_published_documents(agent, published_docs: list, draft_folder, published_folder) -> bool:
    files_changed = False

    draft_files = {item["path"] for item in draft_folder.list_contents()["items"] if item["type"] == "FILE"}
    for published_doc in published_docs:
        if not is_local_doc(published_doc):
            continue

        name = get_doc_name(published_doc)
        if name in draft_files:
            continue

        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)

    return files_changed


def _normalize_restored_documents(published_docs: list | None) -> list:
    restored_docs = []
    for doc in published_docs or []:
        if isinstance(doc, str):
            restored_docs.append({"name": doc, "active": True})
        elif is_local_doc(doc):
            restored_docs.append({**doc, "active": True})
        else:
            restored_docs.append(doc)
    return restored_docs


def restore_discarded_documents(agent, published_version) -> tuple[list, bool]:
    draft_docs = agent.documents or []
    published_docs = published_version.documents or []
    published_doc_names = {get_doc_name(doc) for doc in published_docs if is_local_doc(doc)}

    files_changed = _remove_discarded_local_documents(agent, draft_docs, published_doc_names)

    try:
        project, draft_folder = get_ua_project_folder(agent.id)
        published_folder = get_managed_folder_in_zone(project, "published_documents", "published")

        if published_folder:
            files_changed = (
                _restore_missing_published_documents(agent, published_docs, draft_folder, published_folder)
                or files_changed
            )
    except Exception:
        logger.exception("Restore scan failed for agent %s", agent.id)

    return _normalize_restored_documents(published_docs), files_changed
