import datetime
import io
import json
import logging
import mimetypes
import os
import shutil
import string
import tempfile
import threading
from typing import Any, Dict

import uuid
from collections import Counter
from urllib.parse import urljoin, urlsplit

import requests

from docsgpt.core.settings import settings
from docsgpt.events.publisher import publish_user_event
from docsgpt.parser.chunking_creator import ChunkerCreator
from docsgpt.parser.connectors.connector_creator import ConnectorCreator
from docsgpt.parser.embedding_pipeline import (
    assert_index_complete,
    embed_and_store_documents,
)
from docsgpt.parser.file.base_parser import NoTextLayerError
from docsgpt.parser.file.bulk import SimpleDirectoryReader, get_default_file_extractor
from docsgpt.parser.file.constants import SUPPORTED_SOURCE_EXTENSIONS
from docsgpt.parser.file.image_parser import (
    VISION_CONVERTIBLE_MIME_TYPES,
    convert_image_to_png,
)
from docsgpt.parser.remote.remote_creator import (
    RemoteCreator,
    normalize_remote_data,
)
from docsgpt.parser.schema.base import Document
from docsgpt.security.zip_archive import (
    extract_zip_safely,
    safe_zip_error_message,
    validate_zip_archive,
    ZipExtractionBudget,
    ZipExtractionError,
    ZipExtractionLimits,
)

from docsgpt.storage.db.base_repository import looks_like_uuid
from docsgpt.storage.db.repositories.agents import AgentsRepository
from docsgpt.storage.db.repositories.artifacts import ArtifactsRepository
from docsgpt.storage.db.repositories.attachments import AttachmentsRepository
from docsgpt.storage.db.repositories.ingest_chunk_progress import (
    IngestChunkProgressRepository,
)
from docsgpt.storage.db.repositories.sources import SourcesRepository
from docsgpt.storage.db.repositories.wiki_pages import (
    WikiPagesRepository,
    _content_hash,
    rebuild_wiki_directory_structure,
)
from docsgpt.storage.db.session import db_readonly, db_session
from docsgpt.storage.db.source_config import SourceConfig
from docsgpt.storage.storage_creator import StorageCreator
from docsgpt.upload_limits import (
    enforce_parseable_attachment,
    UnsupportedUploadTypeError,
)
from docsgpt.utils import (
    count_tokens_docs,
    get_encoding,
    num_tokens_from_string,
    safe_filename,
    truncate_to_line_boundary,
)

# Constants


MIN_TOKENS = 150
MAX_TOKENS = 1250
# Attachment content stored for prompting is capped in tokens — the same
# unit the gate is expressed in — so a stored row can never exceed the cap.
ATTACHMENT_MAX_TOKENS = 100_000
RECURSION_DEPTH = 2
INGEST_HEARTBEAT_INTERVAL_SECONDS = 30


def count_structure_files(node: dict) -> int:
    """Count leaf files in a nested ``directory_structure`` mapping.

    Directories are plain dicts of children; files are dicts carrying a
    ``token_count`` key. ``len()`` on the root only sees top-level entries,
    which undercounts any repo with subdirectories.

    Args:
        node: A ``directory_structure`` mapping (or any subtree of one).

    Returns:
        Number of file leaves beneath ``node``.
    """
    if not isinstance(node, dict):
        return 0
    total = 0
    for value in node.values():
        if isinstance(value, dict):
            if "token_count" in value and "size_bytes" in value:
                total += 1
            else:
                total += count_structure_files(value)
    return total


def add_file_to_structure(
    directory_structure: dict,
    file_path: str,
    file_type: str,
    *,
    size_bytes: int,
    token_count: int,
) -> None:
    """Insert one chunk's stats into a nested ``directory_structure``.

    Callers feed this *chunks*, so a file larger than one chunk arrives
    several times. Stats are accumulated rather than overwritten — the
    previous assignment kept only the final fragment, which made the
    per-file sizes shown in the UI wrong for every multi-chunk file
    (a 2.56M-token repo reported 733k).

    Args:
        directory_structure: Mapping mutated in place.
        file_path: Repo-relative path, e.g. ``"guides/setup.md"``.
        file_type: MIME type recorded on first insert.
        size_bytes: Byte length of this chunk.
        token_count: Token count of this chunk.

    Returns:
        None
    """
    path_parts = [p for p in file_path.split("/") if p]
    if not path_parts:
        return
    current_level = directory_structure
    for part in path_parts[:-1]:
        # Intermediate parts are directories
        child = current_level.get(part)
        if not isinstance(child, dict) or "token_count" in child:
            child = {}
            current_level[part] = child
        current_level = child

    leaf = path_parts[-1]
    existing = current_level.get(leaf)
    if isinstance(existing, dict) and "token_count" in existing:
        existing["size_bytes"] += size_bytes
        existing["token_count"] += token_count
    else:
        current_level[leaf] = {
            "type": file_type,
            "size_bytes": size_bytes,
            "token_count": token_count,
        }


def graph_extraction_key(source_id, updated_at) -> str:
    """Build the extract_graph idempotency key for a source's current state.

    The key embeds ``updated_at`` (the trigger-maintained column) so it
    changes whenever re-extraction is warranted — a re-ingest or re-enable
    bumps ``updated_at`` and yields a fresh key, bypassing the 24h completed
    cache so the worker re-runs. Two enqueues for the same state share a key
    and dedupe. ``updated_at`` is stringified deterministically.
    """
    return f"extract-graph:{source_id}:{updated_at}"


def _source_updated_at(source) -> str:
    """Return the source's ``updated_at`` (falling back to ``date``) as a string."""
    if not source:
        return ""
    stamp = source.get("updated_at") or source.get("date")
    return str(stamp) if stamp is not None else ""


def _reset_graph_for_source(source_id) -> None:
    """Drop a source's existing graph so a re-enable/re-ingest rebuilds from scratch.

    Clears nodes, edges, node→chunk links and the ingest checkpoint. Run at the
    enqueue site (not inside the worker) so a broker redelivery of an interrupted
    build still resumes from its checkpoint rather than restarting from zero.
    """
    from docsgpt.graphrag.store import GraphStore

    GraphStore().delete_by_source(str(source_id))


def _publish_graph_event(user, source_id, event_type, payload) -> None:
    """Publish a graph-extraction SSE event, scoped to the source. Never raises."""
    if not user:
        return
    try:
        publish_user_event(
            user,
            event_type,
            payload,
            scope={"kind": "source", "id": str(source_id)},
        )
    except Exception as e:
        logging.debug(f"Failed to publish graph event {event_type}: {e}")


def _maybe_enqueue_graph_extraction(cfg, source_id, user):
    """Reset and re-enqueue graph extraction after embed for a graphrag source.

    The graph lights up asynchronously once chunks are embedded, so ClassicRAG
    works immediately. A no-op for non-graphrag sources or when GraphRAG is
    unavailable. The prior graph is cleared first so a re-ingest rebuilds rather
    than accumulating stale nodes. The work is isolated so a broker hiccup can
    never fail an otherwise-successful ingest.
    """
    if cfg.kind != "graphrag":
        return
    from docsgpt.graphrag import graphrag_available

    if not graphrag_available():
        return

    source_id = str(source_id)
    try:
        from docsgpt.api.user.tasks import extract_graph

        with db_readonly() as conn:
            source = SourcesRepository(conn).get_any(source_id, user)
        _reset_graph_for_source(source_id)
        key = graph_extraction_key(source_id, _source_updated_at(source))
        extract_graph.delay(source_id, user, idempotency_key=key)
    except Exception as e:
        logging.warning(
            f"Failed to enqueue graph extraction for {source_id}: {e}",
            exc_info=True,
        )

# Re-exported here for backward-compatible imports
# (``from docsgpt.worker import _derive_source_id`` /
# ``DOCSGPT_INGEST_NAMESPACE``) from tests and any other in-tree callers.
# New code should import from ``docsgpt.storage.db.source_ids``
# directly to avoid pulling this Celery worker module into the API
# process at import time.
from docsgpt.storage.db.source_ids import (  # noqa: E402, F401
    DOCSGPT_INGEST_NAMESPACE,
    derive_source_id as _derive_source_id,
)


def _ingest_heartbeat_loop(source_id, stop_event, interval=INGEST_HEARTBEAT_INTERVAL_SECONDS):
    """Bump ``ingest_chunk_progress.last_updated`` until ``stop_event`` is set."""
    while not stop_event.wait(interval):
        try:
            with db_session() as conn:
                IngestChunkProgressRepository(conn).bump_heartbeat(source_id)
        except Exception as e:
            logging.warning(
                f"Heartbeat failed for {source_id}: {e}", exc_info=True
            )


def _start_ingest_heartbeat(source_id):
    """Spawn the heartbeat daemon and return ``(thread, stop_event)``."""
    stop_event = threading.Event()
    thread = threading.Thread(
        target=_ingest_heartbeat_loop,
        args=(str(source_id), stop_event),
        daemon=True,
        name=f"ingest-heartbeat-{source_id}",
    )
    thread.start()
    return thread, stop_event


def _stop_ingest_heartbeat(thread, stop_event):
    """Signal the heartbeat daemon to exit and wait briefly for it."""
    if stop_event is not None:
        stop_event.set()
    if thread is not None:
        thread.join(timeout=5)


def _make_parse_progress_callback(task, user, source_id, start_pct, end_pct):
    """Build a ``load_data`` callback mapping parse progress to
    ``[start_pct, end_pct]`` via ``update_state`` + a throttled
    ``stage='parsing'`` SSE event.
    """
    span = end_pct - start_pct
    source_id_str = str(source_id)
    state = {"last_pct": -1}

    def _callback(files_done, total_files):
        if not total_files:
            return
        pct = start_pct + int((files_done / total_files) * span)
        task.update_state(
            state="PROGRESS",
            meta={"current": pct, "status": "Parsing files"},
        )
        if user and pct > state["last_pct"]:
            publish_user_event(
                user,
                "source.ingest.progress",
                {
                    "current": pct,
                    "total": total_files,
                    "files_done": files_done,
                    "stage": "parsing",
                },
                scope={"kind": "source", "id": source_id_str},
            )
            state["last_pct"] = pct

    return _callback


# Define a function to extract metadata from a given filename.


def metadata_from_filename(title):
    return {"title": title}


def _normalize_file_name_map(file_name_map):
    if not file_name_map:
        return {}
    if isinstance(file_name_map, str):
        try:
            file_name_map = json.loads(file_name_map)
        except Exception:
            return {}
    return file_name_map if isinstance(file_name_map, dict) else {}


def _get_display_name(file_name_map, rel_path):
    if not file_name_map or not rel_path:
        return None
    if rel_path in file_name_map:
        return file_name_map[rel_path]
    base_name = os.path.basename(rel_path)
    return file_name_map.get(base_name)


def _apply_display_names_to_structure(structure, file_name_map, prefix=""):
    if not isinstance(structure, dict) or not file_name_map:
        return structure
    for name, node in structure.items():
        if isinstance(node, dict) and "type" in node and "size_bytes" in node:
            rel_path = f"{prefix}/{name}" if prefix else name
            display_name = _get_display_name(file_name_map, rel_path)
            if display_name:
                node["display_name"] = display_name
        elif isinstance(node, dict):
            next_prefix = f"{prefix}/{name}" if prefix else name
            _apply_display_names_to_structure(node, file_name_map, next_prefix)
    return structure


def _download_source_files_to_dir(storage, source_file_path, temp_dir):
    """Mirror a source's stored files into ``temp_dir``, preserving structure."""
    if not storage.is_directory(source_file_path):
        return
    for storage_file_path in storage.list_files(source_file_path):
        if storage.is_directory(storage_file_path):
            continue
        rel_path = os.path.relpath(storage_file_path, source_file_path)
        local_file_path = os.path.join(temp_dir, rel_path)
        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
        try:
            file_data = storage.get_file(storage_file_path)
            with open(local_file_path, "wb") as f:
                f.write(file_data.read())
        except Exception as e:
            logging.error(f"Error downloading file {storage_file_path}: {e}")
            continue


# Define a function to generate a random string of a given length.


def generate_random_string(length):
    return "".join([string.ascii_letters[i % 52] for i in range(length)])


# Zip extraction security limits. Kept as module constants for backward
# compatibility with worker callers/tests; values come from operator settings.
MAX_UNCOMPRESSED_SIZE = settings.UPLOAD_MAX_ARCHIVE_BYTES
MAX_FILE_COUNT = settings.UPLOAD_MAX_ARCHIVE_FILES
MAX_COMPRESSION_RATIO = settings.UPLOAD_MAX_ARCHIVE_RATIO


def _is_path_safe(base_path: str, target_path: str) -> bool:
    """
    Check if target_path is safely within base_path (prevents zip slip attacks).

    Args:
        base_path: The base directory where extraction should occur.
        target_path: The full path where a file would be extracted.

    Returns:
        True if the path is safe, False otherwise.
    """
    # Resolve to absolute paths and check containment
    base_resolved = os.path.realpath(base_path)
    target_resolved = os.path.realpath(target_path)
    return target_resolved.startswith(base_resolved + os.sep) or target_resolved == base_resolved


def _validate_zip_safety(zip_path: str, extract_to: str) -> None:
    """
    Validate a zip file for security issues before extraction.

    Checks for:
    - Zip bombs (excessive compression ratio or uncompressed size)
    - Too many files
    - Path traversal attacks (zip slip)

    Args:
        zip_path: Path to the zip file.
        extract_to: Destination directory.

    Raises:
        ZipExtractionError: If the zip file fails security validation.
    """
    del extract_to  # Path checks are platform-independent inside the validator.
    validate_zip_archive(
        zip_path,
        ZipExtractionLimits(
            max_uncompressed_bytes=MAX_UNCOMPRESSED_SIZE,
            max_files=MAX_FILE_COUNT,
            max_compression_ratio=MAX_COMPRESSION_RATIO,
            max_member_bytes=settings.UPLOAD_MAX_FILE_BYTES,
            max_depth=RECURSION_DEPTH,
        ),
    )


def extract_zip_recursive(
    zip_path: str,
    extract_to: str,
    current_depth: int = 0,
    max_depth: int = 5,
    _budget: ZipExtractionBudget | None = None,
) -> None:
    """
    Recursively extract zip files with security protections.

    Security measures:
    - Limits recursion depth to prevent infinite loops
    - Validates uncompressed size to prevent zip bombs
    - Limits number of files to prevent resource exhaustion
    - Checks compression ratio to detect zip bombs
    - Validates paths to prevent zip slip attacks

    Args:
        zip_path (str): Path to the zip file to be extracted.
        extract_to (str): Destination path for extracted files.
        current_depth (int): Current depth of recursion.
        max_depth (int): Maximum allowed depth of recursion to prevent infinite loops.

    Raises:
        ZipExtractionError: If the archive fails safety validation, so ingestion
            fails loudly instead of indexing an empty directory.
    """
    if current_depth > max_depth:
        logging.warning(f"Reached maximum recursion depth of {max_depth}")
        return

    try:
        budget = _budget or ZipExtractionBudget()
        extract_zip_safely(
            zip_path,
            extract_to,
            ZipExtractionLimits(
                max_uncompressed_bytes=MAX_UNCOMPRESSED_SIZE,
                max_files=MAX_FILE_COUNT,
                max_compression_ratio=MAX_COMPRESSION_RATIO,
                max_member_bytes=settings.UPLOAD_MAX_FILE_BYTES,
                max_depth=max_depth - current_depth,
            ),
            budget,
        )
        os.remove(zip_path)  # Remove the zip file after extracting

    except ZipExtractionError as e:
        logging.error(
            "Zip security validation failed for %s: %s",
            safe_zip_error_message(zip_path),
            safe_zip_error_message(e),
        )
        # Remove the potentially malicious zip file
        try:
            os.remove(zip_path)
        except OSError:
            pass
        raise
    except Exception as e:
        logging.error(
            "Error extracting zip file %s: %s",
            safe_zip_error_message(zip_path),
            safe_zip_error_message(e),
            exc_info=True,
        )
        raise


def download_file(url, params, dest_path):
    try:
        response = requests.get(url, params=params, timeout=100)
        response.raise_for_status()
        with open(dest_path, "wb") as f:
            f.write(response.content)
    except requests.RequestException as e:
        logging.error(f"Error downloading file: {e}")
        raise


def upload_index(full_path, file_data):
    files = None
    try:
        headers = {}
        if settings.INTERNAL_KEY:
            headers["X-Internal-Key"] = settings.INTERNAL_KEY

        if settings.VECTOR_STORE == "faiss":
            faiss_path = full_path + "/index.faiss"
            pkl_path = full_path + "/index.pkl"

            if not os.path.exists(faiss_path):
                logging.error(f"FAISS index file not found: {faiss_path}")
                raise FileNotFoundError(f"FAISS index file not found: {faiss_path}")

            if not os.path.exists(pkl_path):
                logging.error(f"FAISS pickle file not found: {pkl_path}")
                raise FileNotFoundError(f"FAISS pickle file not found: {pkl_path}")

            files = {
                "file_faiss": open(faiss_path, "rb"),
                "file_pkl": open(pkl_path, "rb"),
            }
            response = requests.post(
                urljoin(settings.API_URL, "/api/upload_index"),
                files=files,
                data=file_data,
                headers=headers,
                timeout=100,
            )
        else:
            response = requests.post(
                urljoin(settings.API_URL, "/api/upload_index"),
                data=file_data,
                headers=headers,
                timeout=100,
            )
        response.raise_for_status()
    except (requests.RequestException, FileNotFoundError) as e:
        logging.error(f"Error uploading index: {e}")
        raise
    finally:
        if settings.VECTOR_STORE == "faiss" and files is not None:
            for file in files.values():
                file.close()


# Define the main function for ingesting and processing documents.


def ingest_worker(
    self,
    directory,
    formats,
    job_name,
    file_path,
    filename,
    user,
    retriever="classic",
    file_name_map=None,
    config=None,
    idempotency_key=None,
    source_id=None,
):
    """
    Ingest and process documents.

    Args:
        self: Reference to the instance of the task.
        directory (str): Specifies the directory for ingesting ('inputs' or 'temp').
        formats (list of str): List of file extensions to consider for ingestion (e.g., [".rst", ".md"]).
        job_name (str): Name of the job for this ingestion task (original, unsanitized).
        file_path (str): Complete file path to use consistently throughout the pipeline.
        filename (str): Original unsanitized filename provided by the user.
        user (str): Identifier for the user initiating the ingestion (original, unsanitized).
        retriever (str): Type of retriever to use for processing the documents.
        file_name_map (dict|str|None): Optional mapping of safe relative paths to original filenames.
        config (dict|None): Per-source ``SourceConfig`` dict. ``None``/``{}`` →
            classic defaults (byte-identical to prior behavior).
        idempotency_key (str|None): When provided, the ``source_id`` is derived
            deterministically from the key so a retried task reuses the same
            source row instead of duplicating it.
        source_id (str|None): UUID minted by the HTTP route and returned in
            its response. When supplied, the worker uses it verbatim so SSE
            envelopes carry the same id the frontend already has — required
            for non-idempotent uploads where the route can't predict
            ``_derive_source_id(idempotency_key)``.

    Returns:
        dict: Information about the completed ingestion task, including input parameters and a "limited" flag.
    """
    input_files = None
    recursive = True
    limit = None
    exclude = True
    sample = False

    storage = StorageCreator.get_storage()

    logging.info(f"Ingest path: {file_path}", extra={"user": user, "job": job_name})

    # Source id resolution order:
    #   1. Caller-supplied ``source_id`` (HTTP route minted + returned to
    #      the frontend) — keeps the route response and the SSE event
    #      payloads in lockstep on the non-idempotent path.
    #   2. Deterministic uuid5 from ``idempotency_key`` — retried tasks
    #      reuse the original source row instead of duplicating it.
    #   3. Fresh uuid4 (caller has neither) — opaque, single-shot only.
    if source_id:
        source_uuid = uuid.UUID(source_id)
    else:
        source_uuid = _derive_source_id(idempotency_key)
    source_id_for_events = str(source_uuid)
    # Only emit ``queued`` on the original attempt. Celery retries re-run
    # the body, and re-publishing here would oscillate the toast through
    # ``queued`` again between ``failed`` and ``completed``.
    if self.request.retries == 0:
        publish_user_event(
            user,
            "source.ingest.queued",
            {
                "job_name": job_name,
                "filename": filename,
                "source_id": source_id_for_events,
                "operation": "upload",
            },
            scope={"kind": "source", "id": source_id_for_events},
        )

    # Wrap the entire body in try/except so a failure between the
    # ``queued`` publish above and the inner work (e.g. tempdir
    # creation, OS-level resource exhaustion) still emits a terminal
    # ``failed`` event rather than leaving the toast wedged on
    # 'training' until the polling fallback rescues it 30s later.
    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            os.makedirs(temp_dir, exist_ok=True)

            if storage.is_directory(file_path):
                # Handle directory case
                logging.info(f"Processing directory: {file_path}")
                files_list = storage.list_files(file_path)

                for storage_file_path in files_list:
                    if storage.is_directory(storage_file_path):
                        continue

                    # Create relative path structure in temp directory
                    rel_path = os.path.relpath(storage_file_path, file_path)
                    local_file_path = os.path.join(temp_dir, rel_path)

                    os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

                    # Download file
                    try:
                        file_data = storage.get_file(storage_file_path)
                        with open(local_file_path, "wb") as f:
                            f.write(file_data.read())
                    except Exception as e:
                        logging.error(
                            f"Error downloading file {storage_file_path}: {e}"
                        )
                        continue
            else:
                # Handle single file case
                temp_filename = os.path.basename(file_path)
                temp_file_path = os.path.join(temp_dir, temp_filename)

                file_data = storage.get_file(file_path)
                with open(temp_file_path, "wb") as f:
                    f.write(file_data.read())

                # Handle zip files
                if temp_filename.lower().endswith(".zip"):
                    logging.info(f"Extracting zip file: {temp_filename}")
                    extract_zip_recursive(
                        temp_file_path,
                        temp_dir,
                        current_depth=0,
                        max_depth=RECURSION_DEPTH,
                    )

            self.update_state(state="PROGRESS", meta={"current": 1})
            if sample:
                logging.info(f"Sample mode enabled. Using {limit} documents.")
            reader = SimpleDirectoryReader(
                input_dir=temp_dir,
                input_files=input_files,
                recursive=recursive,
                required_exts=formats,
                exclude_hidden=exclude,
                file_metadata=metadata_from_filename,
            )
            # Parsing/OCR owns 1-50% of the bar; embedding takes 50-100%.
            raw_docs = reader.load_data(
                progress_callback=_make_parse_progress_callback(
                    self, user, source_uuid, start_pct=1, end_pct=50,
                )
            )

            directory_structure = getattr(reader, "directory_structure", {})
            logging.info(f"Directory structure from reader: {directory_structure}")
            file_name_map = _normalize_file_name_map(file_name_map)
            if file_name_map:
                for doc in raw_docs:
                    extra_info = getattr(doc, "extra_info", None)
                    if not isinstance(extra_info, dict):
                        continue
                    rel_path = extra_info.get("source") or extra_info.get("file_path")
                    display_name = _get_display_name(file_name_map, rel_path)
                    if display_name:
                        display_name = str(display_name)
                        extra_info["filename"] = display_name
                        extra_info["file_name"] = display_name
                        extra_info["title"] = display_name
                directory_structure = _apply_display_names_to_structure(
                    directory_structure, file_name_map
                )

            cfg = SourceConfig.parse(config)
            chunker = ChunkerCreator.create_chunker(
                cfg.chunking.strategy,
                chunking_strategy=cfg.chunking.strategy,
                max_tokens=cfg.chunking.max_tokens,
                min_tokens=cfg.chunking.min_tokens,
                duplicate_headers=cfg.chunking.duplicate_headers,
            )
            raw_docs = chunker.chunk(documents=raw_docs)

            docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs]

            vector_store_path = os.path.join(temp_dir, "vector_store")
            os.makedirs(vector_store_path, exist_ok=True)

            heartbeat_thread, heartbeat_stop = _start_ingest_heartbeat(source_uuid)
            try:
                embed_and_store_documents(
                    docs, vector_store_path, source_uuid, self,
                    attempt_id=getattr(self.request, "id", None),
                    user_id=user,
                    progress_start=50, progress_end=100,
                )
            finally:
                _stop_ingest_heartbeat(heartbeat_thread, heartbeat_stop)
            # Defense-in-depth: chunk-progress is the authoritative
            # record of how many chunks landed; mismatch raises so the
            # task fails loud rather than caching a partial index.
            assert_index_complete(source_uuid)

            tokens = count_tokens_docs(docs)

            self.update_state(state="PROGRESS", meta={"current": 100})

            if sample:
                for i in range(min(5, len(raw_docs))):
                    logging.info(f"Sample document {i}: {raw_docs[i]}")
            file_data = {
                "name": job_name,
                "file": filename,
                "user": user,
                "tokens": tokens,
                "retriever": retriever,
                "id": source_id_for_events,
                "type": "local",
                "file_path": file_path,
                "directory_structure": json.dumps(directory_structure),
            }
            if file_name_map:
                file_data["file_name_map"] = json.dumps(file_name_map)
            if config:
                file_data["config"] = json.dumps(config)

            upload_index(vector_store_path, file_data)
            publish_user_event(
                user,
                "source.ingest.completed",
                {
                    "source_id": source_id_for_events,
                    "filename": filename,
                    "tokens": tokens,
                    "operation": "upload",
                    # Forward-looking contract: ``limited`` is always
                    # ``False`` today but is carried on the wire so a
                    # future token-cap detection path can flip it and
                    # the frontend slice / UploadToast already react.
                    "limited": False,
                },
                scope={"kind": "source", "id": source_id_for_events},
            )
            _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)
    except Exception as e:
        logging.error(f"Error in ingest_worker: {e}", exc_info=True)
        publish_user_event(
            user,
            "source.ingest.failed",
            {
                "source_id": source_id_for_events,
                "filename": filename,
                "operation": "upload",
                "error": str(e)[:1024],
            },
            scope={"kind": "source", "id": source_id_for_events},
        )
        raise
    return {
        "directory": directory,
        "formats": formats,
        "name_job": job_name,  # Use original job_name
        "filename": filename,
        "user": user,  # Use original user
        "limited": False,
    }


def reingest_source_worker(self, source_id, user):
    """
    Re-ingestion worker that handles incremental updates by:
    1. Adding chunks from newly added files
    2. Removing chunks from deleted files

    Args:
        self: Task instance
        source_id: ID of the source to re-ingest
        user: User identifier

    Returns:
        dict: Information about the re-ingestion task

    Note:
        Reingest does its own ``vector_store.add_chunk`` work rather
        than going through ``embed_and_store_documents`` so it does
        *not* emit per-percent SSE progress events — only ``queued``,
        ``completed`` (carrying ``chunks_added`` / ``chunks_deleted``),
        or ``failed``. v1 limitation; revisit if reingest gains a
        progress-driven UI.
    """
    # Declared at the function scope so the outer except can include
    # ``name`` in the failed event payload when the failure happens
    # after the source lookup. Empty string until the lookup succeeds.
    source_name = ""
    # Tracks inner-block failures so a ``completed`` event reflects
    # partial-success accurately rather than masking it.
    inner_warnings: list[str] = []

    try:
        from docsgpt.vectorstore.vector_creator import VectorCreator

        self.update_state(
            state="PROGRESS",
            meta={"current": 10, "status": "Initializing re-ingestion scan"},
        )

        with db_readonly() as conn:
            source = SourcesRepository(conn).get_any(source_id, user)
        if not source:
            raise ValueError(f"Source {source_id} not found or access denied")
        source_id = str(source["id"])
        source_name = source.get("name") or ""
        # Re-chunk with the source's persisted config (classic defaults
        # when absent/empty), keeping reingest consistent with ingest.
        cfg = SourceConfig.parse(source.get("config"))

        # Publish ``queued`` *after* canonicalising ``source_id`` so the
        # event references the same id as the source row. Trade-off
        # documented: a Celery-backend or PG-lookup hiccup before this
        # publish means the toast may see only a ``failed`` event with
        # no preceding ``queued`` — acceptable for v1 since both
        # conditions also imply broader system trouble. Gate on first
        # attempt only so Celery retries don't re-emit ``queued`` after
        # a prior attempt already published ``failed``.
        if self.request.retries == 0:
            publish_user_event(
                user,
                "source.ingest.queued",
                {
                    "source_id": source_id,
                    "name": source_name,
                    # ``filename`` labels the upload toast on auto-create.
                    "filename": source_name,
                    "operation": "reingest",
                },
                scope={"kind": "source", "id": source_id},
            )

        storage = StorageCreator.get_storage()
        source_file_path = source.get("file_path", "")
        file_name_map = _normalize_file_name_map(source.get("file_name_map"))

        self.update_state(
            state="PROGRESS", meta={"current": 20, "status": "Scanning current files"}
        )

        with tempfile.TemporaryDirectory() as temp_dir:
            _download_source_files_to_dir(storage, source_file_path, temp_dir)

            reader = SimpleDirectoryReader(
                input_dir=temp_dir,
                recursive=True,
                required_exts=list(SUPPORTED_SOURCE_EXTENSIONS),
                exclude_hidden=True,
                file_metadata=metadata_from_filename,
            )
            reader.load_data()
            directory_structure = reader.directory_structure
            logging.info(
                f"Directory structure built with token counts: {directory_structure}"
            )

            try:
                old_directory_structure = source.get("directory_structure") or {}
                if isinstance(old_directory_structure, str):
                    try:
                        old_directory_structure = json.loads(old_directory_structure)
                    except Exception:
                        old_directory_structure = {}

                def _flatten_directory_structure(struct, prefix=""):
                    files = set()
                    if isinstance(struct, dict):
                        for name, meta in struct.items():
                            current_path = (
                                os.path.join(prefix, name) if prefix else name
                            )
                            if isinstance(meta, dict) and (
                                "type" in meta and "size_bytes" in meta
                            ):
                                files.add(current_path)
                            elif isinstance(meta, dict):
                                files |= _flatten_directory_structure(
                                    meta, current_path
                                )
                    return files

                old_files = _flatten_directory_structure(old_directory_structure)
                new_files = _flatten_directory_structure(directory_structure)

                added_files = sorted(new_files - old_files)
                removed_files = sorted(old_files - new_files)

                if added_files:
                    logging.info(f"Files added since last ingest: {added_files}")
                else:
                    logging.info("No files added since last ingest.")

                if removed_files:
                    logging.info(f"Files removed since last ingest: {removed_files}")
                else:
                    logging.info("No files removed since last ingest.")

            except Exception as e:
                logging.error(
                    f"Error comparing directory structures: {e}", exc_info=True
                )
                added_files = []
                removed_files = []
            try:
                if not added_files and not removed_files:
                    logging.info("No changes detected.")
                    publish_user_event(
                        user,
                        "source.ingest.completed",
                        {
                            "source_id": source_id,
                            "name": source_name,
                            "filename": source_name,
                            "operation": "reingest",
                            "no_changes": True,
                            "chunks_added": 0,
                            "chunks_deleted": 0,
                        },
                        scope={"kind": "source", "id": source_id},
                    )
                    return {
                        "source_id": source_id,
                        "user": user,
                        "status": "no_changes",
                        "added_files": [],
                        "removed_files": [],
                    }

                vector_store = VectorCreator.create_vectorstore(
                    settings.VECTOR_STORE,
                    source_id,
                    settings.EMBEDDINGS_KEY,
                )

                self.update_state(
                    state="PROGRESS",
                    meta={"current": 40, "status": "Processing file changes"},
                )

                # 1) Delete chunks from removed files
                deleted = 0
                if removed_files:
                    try:
                        for ch in vector_store.get_chunks() or []:
                            metadata = (
                                ch.get("metadata", {})
                                if isinstance(ch, dict)
                                else getattr(ch, "metadata", {})
                            )
                            raw_source = metadata.get("source")

                            source_file = str(raw_source) if raw_source else ""

                            if source_file in removed_files:
                                cid = ch.get("doc_id")
                                if cid:
                                    try:
                                        vector_store.delete_chunk(cid)
                                        deleted += 1
                                    except Exception as de:
                                        logging.error(
                                            f"Failed deleting chunk {cid}: {de}"
                                        )
                        logging.info(
                            f"Deleted {deleted} chunks from {len(removed_files)} removed files"
                        )
                    except Exception as e:
                        logging.error(
                            f"Error during deletion of removed file chunks: {e}",
                            exc_info=True,
                        )
                        inner_warnings.append(
                            f"deletion failed: {str(e)[:200]}"
                        )

                # 2) Add chunks from new files
                added = 0
                if added_files:
                    try:
                        # Build list of local files for added files only
                        added_local_files = []
                        for rel_path in added_files:
                            local_path = os.path.join(temp_dir, rel_path)
                            if os.path.isfile(local_path):
                                added_local_files.append(local_path)

                        if added_local_files:
                            reader_new = SimpleDirectoryReader(
                                input_files=added_local_files,
                                exclude_hidden=True,
                                errors="ignore",
                                file_metadata=metadata_from_filename,
                            )
                            raw_docs_new = reader_new.load_data()
                            chunker_new = ChunkerCreator.create_chunker(
                                cfg.chunking.strategy,
                                chunking_strategy=cfg.chunking.strategy,
                                max_tokens=cfg.chunking.max_tokens,
                                min_tokens=cfg.chunking.min_tokens,
                                duplicate_headers=cfg.chunking.duplicate_headers,
                            )
                            chunked_new = chunker_new.chunk(documents=raw_docs_new)

                            for (
                                file_path,
                                token_count,
                            ) in reader_new.file_token_counts.items():
                                try:
                                    rel_path = os.path.relpath(
                                        file_path, start=temp_dir
                                    )
                                    path_parts = rel_path.split(os.sep)
                                    current_dir = directory_structure

                                    for part in path_parts[:-1]:
                                        if part in current_dir and isinstance(
                                            current_dir[part], dict
                                        ):
                                            current_dir = current_dir[part]
                                        else:
                                            break

                                    filename = path_parts[-1]
                                    if filename in current_dir and isinstance(
                                        current_dir[filename], dict
                                    ):
                                        current_dir[filename][
                                            "token_count"
                                        ] = token_count
                                        logging.info(
                                            f"Updated token count for {rel_path}: {token_count}"
                                        )
                                except Exception as e:
                                    logging.warning(
                                        f"Could not update token count for {file_path}: {e}"
                                    )

                            for d in chunked_new:
                                meta = dict(d.extra_info or {})
                                try:
                                    raw_src = meta.get("source")
                                    if isinstance(raw_src, str) and os.path.isabs(
                                        raw_src
                                    ):
                                        meta["source"] = os.path.relpath(
                                            raw_src, start=temp_dir
                                        )
                                except Exception:
                                    pass
                                display_name = _get_display_name(
                                    file_name_map, meta.get("source")
                                )
                                if display_name:
                                    display_name = str(display_name)
                                    meta["filename"] = display_name
                                    meta["file_name"] = display_name
                                    meta["title"] = display_name

                                vector_store.add_chunk(d.text, metadata=meta)
                                added += 1
                            logging.info(
                                f"Added {added} chunks from {len(added_files)} new files"
                            )
                    except Exception as e:
                        logging.error(
                            f"Error during ingestion of new files: {e}", exc_info=True
                        )
                        inner_warnings.append(
                            f"add failed: {str(e)[:200]}"
                        )

                # 3) Update source directory structure timestamp
                try:
                    total_tokens = sum(reader.file_token_counts.values())
                    directory_structure = _apply_display_names_to_structure(
                        directory_structure, file_name_map
                    )

                    now = datetime.datetime.now(datetime.timezone.utc)
                    with db_session() as conn:
                        SourcesRepository(conn).update(
                            source_id, user,
                            {
                                "directory_structure": directory_structure,
                                "date": now,
                                "tokens": total_tokens,
                            },
                        )
                except Exception as e:
                    logging.error(
                        f"Error updating directory_structure in DB: {e}", exc_info=True
                    )

                self.update_state(
                    state="PROGRESS",
                    meta={"current": 100, "status": "Re-ingestion completed"},
                )

                completed_payload: dict = {
                    "source_id": source_id,
                    "name": source_name,
                    "filename": source_name,
                    "operation": "reingest",
                    "chunks_added": added,
                    "chunks_deleted": deleted,
                    "tokens": int(total_tokens) if "total_tokens" in locals() else 0,
                }
                if inner_warnings:
                    # Surface the per-block failures so the toast can warn
                    # rather than claim a clean success.
                    completed_payload["warnings"] = inner_warnings
                publish_user_event(
                    user,
                    "source.ingest.completed",
                    completed_payload,
                    scope={"kind": "source", "id": source_id},
                )
                _maybe_enqueue_graph_extraction(cfg, source_id, user)

                return {
                    "source_id": source_id,
                    "user": user,
                    "status": "completed",
                    "added_files": added_files,
                    "removed_files": removed_files,
                    "chunks_added": added,
                    "chunks_deleted": deleted,
                }
            except Exception as e:
                logging.error(
                    f"Error while processing file changes: {e}", exc_info=True
                )
                raise

    except Exception as e:
        logging.error(f"Error in reingest_source_worker: {e}", exc_info=True)
        publish_user_event(
            user,
            "source.ingest.failed",
            {
                "source_id": str(source_id),
                "name": source_name,
                "filename": source_name,
                "operation": "reingest",
                "error": str(e)[:1024],
            },
            scope={"kind": "source", "id": str(source_id)},
        )
        raise


def remote_worker(
    self,
    source_data,
    name_job,
    user,
    loader,
    directory="temp",
    retriever="classic",
    sync_frequency="never",
    operation_mode="upload",
    doc_id=None,
    config=None,
    idempotency_key=None,
    source_id=None,
):
    safe_user = safe_filename(user)
    full_path = os.path.join(directory, safe_user, uuid.uuid4().hex)
    os.makedirs(full_path, exist_ok=True)

    # Source id resolution order matches ``ingest_worker``:
    #   1. ``operation_mode == "sync"`` reuses the existing source's ``doc_id``.
    #   2. Caller-supplied ``source_id`` (the HTTP route minted it and
    #      already returned it to the frontend) — keeps the route
    #      response and the SSE event payloads in lockstep on the
    #      no-idempotency-key path.
    #   3. Deterministic uuid5 from ``idempotency_key`` — retried tasks
    #      reuse the original source row instead of duplicating it.
    #   4. Fresh uuid4 — opaque, single-shot only.
    if operation_mode == "sync" and doc_id:
        source_uuid = str(doc_id)
    elif source_id:
        source_uuid = uuid.UUID(source_id)
    else:
        source_uuid = _derive_source_id(idempotency_key)
    source_id_for_events = str(source_uuid)

    # Emit the queued event before any work that could fail (including
    # ``update_state``) so the toast UI always sees a queued envelope
    # before any subsequent failed event. Gated on first attempt so
    # Celery retries don't re-emit ``queued`` after a prior ``failed``.
    if self.request.retries == 0:
        publish_user_event(
            user,
            "source.ingest.queued",
            {
                "source_id": source_id_for_events,
                "job_name": name_job,
                "loader": loader,
                "operation": operation_mode,
            },
            scope={"kind": "source", "id": source_id_for_events},
        )

    # Wrap ``update_state`` plus the entire body so any pre-loader
    # failure (Celery backend down, OS resource issue) still emits a
    # terminal ``failed`` event rather than wedging the toast.
    try:
        self.update_state(state="PROGRESS", meta={"current": 1})
        logging.info("Initializing remote loader with type: %s", loader)
        remote_loader = RemoteCreator.create_loader(loader)
        raw_docs = remote_loader.load_data(source_data)

        cfg = SourceConfig.parse(config)
        chunker = ChunkerCreator.create_chunker(
            cfg.chunking.strategy,
            chunking_strategy=cfg.chunking.strategy,
            max_tokens=cfg.chunking.max_tokens,
            min_tokens=cfg.chunking.min_tokens,
            duplicate_headers=cfg.chunking.duplicate_headers,
        )
        raw_docs = chunker.chunk(documents=raw_docs)
        docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs]
        tokens = count_tokens_docs(docs)
        logging.info("Total tokens calculated: %d", tokens)

        # Build directory structure from loaded documents
        # Format matches local file uploads: nested structure with type, size_bytes, token_count
        directory_structure = {}
        for doc in raw_docs:
            # Get the file path from extra_info
            # For crawlers: file_path is a virtual path like "guides/setup.md"
            # For other remotes: use key or title as fallback
            file_path = ""
            if doc.extra_info:
                file_path = (
                    doc.extra_info.get("file_path", "")
                    or doc.extra_info.get("key", "")
                    or doc.extra_info.get("title", "")
                )
            if not file_path:
                file_path = doc.doc_id or ""

            if file_path:
                # Calculate token count
                token_count = num_tokens_from_string(doc.text) if doc.text else 0

                # Estimate size in bytes from text content
                size_bytes = len(doc.text.encode("utf-8")) if doc.text else 0

                # Guess mime type from extension
                file_name = (
                    file_path.split("/")[-1] if "/" in file_path else file_path
                )
                ext = os.path.splitext(file_name)[1].lower()
                mime_types = {
                    ".txt": "text/plain",
                    ".md": "text/markdown",
                    ".pdf": "application/pdf",
                    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                    ".doc": "application/msword",
                    ".html": "text/html",
                    ".json": "application/json",
                    ".csv": "text/csv",
                    ".xml": "application/xml",
                    ".py": "text/x-python",
                    ".js": "text/javascript",
                    ".ts": "text/typescript",
                    ".jsx": "text/jsx",
                    ".tsx": "text/tsx",
                }
                file_type = mime_types.get(ext, "application/octet-stream")

                # Build nested directory structure from path
                # e.g., "guides/setup.md" -> {"guides": {"setup.md": {...}}}
                add_file_to_structure(
                    directory_structure, file_path, file_type,
                    size_bytes=size_bytes, token_count=token_count,
                )

        # ``len(directory_structure)`` counts only top-level entries, so a
        # 1,474-file repo logged "44 files". Count the leaves instead — this
        # line is the operational signal that a remote ingest succeeded.
        logging.info(
            f"Built directory structure with "
            f"{count_structure_files(directory_structure)} files across "
            f"{len(directory_structure)} top-level entries: "
            f"{list(directory_structure.keys())}"
        )

        if operation_mode == "upload":
            embed_and_store_documents(
                docs, full_path, source_uuid, self,
                attempt_id=getattr(self.request, "id", None),
                user_id=user,
            )
            assert_index_complete(source_uuid)
        elif operation_mode == "sync":
            if not doc_id:
                logging.error("Invalid doc_id provided for sync operation: %s", doc_id)
                raise ValueError("doc_id must be provided for sync operation.")
            embed_and_store_documents(
                docs, full_path, source_uuid, self,
                attempt_id=getattr(self.request, "id", None),
                user_id=user,
            )
            assert_index_complete(source_uuid)
        self.update_state(state="PROGRESS", meta={"current": 100})

        # Serialize remote_data as JSON if it's a dict (for S3, Reddit, etc.)
        remote_data_serialized = (
            json.dumps(source_data) if isinstance(source_data, dict) else source_data
        )
        file_data = {
            "name": name_job,
            "user": user,
            "tokens": tokens,
            "retriever": retriever,
            "id": source_id_for_events,
            "type": loader,
            "remote_data": remote_data_serialized,
            "sync_frequency": sync_frequency,
            "directory_structure": json.dumps(directory_structure),
        }
        if config:
            file_data["config"] = json.dumps(config)

        if operation_mode == "sync":
            last_sync_now = datetime.datetime.now(datetime.timezone.utc)
            file_data["last_sync"] = last_sync_now

            try:
                with db_session() as conn:
                    repo = SourcesRepository(conn)
                    src = repo.get_any(source_id_for_events, user)
                    if src is not None:
                        repo.update(str(src["id"]), user, {"date": last_sync_now})
            except Exception as upd_err:
                logging.warning(
                    f"Failed to update last_sync for source {source_id_for_events}: {upd_err}"
                )
        upload_index(full_path, file_data)
        publish_user_event(
            user,
            "source.ingest.completed",
            {
                "source_id": source_id_for_events,
                "job_name": name_job,
                "loader": loader,
                "operation": operation_mode,
                "tokens": tokens,
                # Forward-looking contract: see ingest_worker.
                "limited": False,
            },
            scope={"kind": "source", "id": source_id_for_events},
        )
        _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)
    except Exception as e:
        logging.error("Error in remote_worker task: %s", str(e), exc_info=True)
        publish_user_event(
            user,
            "source.ingest.failed",
            {
                "source_id": source_id_for_events,
                "job_name": name_job,
                "loader": loader,
                "operation": operation_mode,
                "error": str(e)[:1024],
            },
            scope={"kind": "source", "id": source_id_for_events},
        )
        raise
    finally:
        if os.path.exists(full_path):
            shutil.rmtree(full_path)
    logging.info("remote_worker task completed successfully")
    return {
        "id": source_id_for_events,
        "urls": source_data,
        "name_job": name_job,
        "user": user,
        "limited": False,
    }


def sync(
    self,
    source_data,
    name_job,
    user,
    loader,
    sync_frequency,
    retriever,
    doc_id=None,
    directory="temp",
):
    try:
        remote_worker(
            self,
            source_data,
            name_job,
            user,
            loader,
            directory,
            retriever,
            sync_frequency,
            "sync",
            doc_id,
        )
    except Exception as e:
        logging.error(f"Error during sync: {e}", exc_info=True)
        return {"status": "error", "error": str(e)}
    return {"status": "success"}


def sync_worker(self, frequency):
    from sqlalchemy import text as sql_text

    sync_counts = Counter()
    with db_readonly() as conn:
        result = conn.execute(
            sql_text(
                "SELECT id, name, user_id, type, remote_data, retriever "
                "FROM sources WHERE sync_frequency = :freq"
            ),
            {"freq": frequency},
        )
        rows = result.fetchall()

    for row in rows:
        doc = dict(row._mapping)
        name = doc.get("name")
        user = doc.get("user_id")
        source_type = doc.get("type")
        retriever = doc.get("retriever")
        doc_id = str(doc.get("id"))

        sync_counts["total_sync_count"] += 1

        # Connector sources have no RemoteCreator loader and need an OAuth
        # token to sync, which a scheduled task lacks — skip them.
        if source_type and source_type.startswith("connector"):
            sync_counts["sync_skipped"] += 1
            continue

        source_data = normalize_remote_data(source_type, doc.get("remote_data"))
        if not source_data:
            # No syncable URL/config — skip instead of dispatching a sync
            # that can only fail (and emit a spurious failed event).
            sync_counts["sync_skipped"] += 1
            continue

        resp = sync(
            self, source_data, name, user, source_type, frequency, retriever, doc_id
        )
        sync_counts[
            "sync_success" if resp["status"] == "success" else "sync_failure"
        ] += 1
    return {
        key: sync_counts[key]
        for key in [
            "total_sync_count", "sync_success", "sync_failure", "sync_skipped",
        ]
    }


# Line-oriented text formats that stay parseable after a head-truncation.
# Structured/binary formats (pdf, docx, xlsx, json, html, ...) are excluded —
# cutting them mid-stream would break their parsers entirely.
_TRUNCATABLE_ATTACHMENT_SUFFIXES = {
    ".csv",
    ".tsv",
    ".txt",
    ".log",
    ".md",
    ".mdx",
    ".rst",
    ".jsonl",
}


class AttachmentRejectedError(Exception):
    """A deterministic reason an attachment must not be parsed (e.g. zip bomb).

    Raised before parsing and marked non-retryable on the Celery task so a
    poison upload fails once instead of retrying identically.
    """


def _reject_unparseable_attachment(
    local_path: str, filename: str, parser_extensions
) -> None:
    """Reject an attachment with no parser whose contents are binary.

    Defence in depth behind the route's gate, and stricter than it: this runs
    against the parser table actually loaded, so a suffix the route trusted
    (.webp with docling absent, say) is still refused rather than opened as
    plain text by ``SimpleDirectoryReader`` and stored as garbage.

    Args:
        local_path: Filesystem path of the attachment about to be parsed.
        filename: The upload's original filename, which carries the suffix.
        parser_extensions: Suffixes the live ``file_extractor`` can parse.

    Raises:
        AttachmentRejectedError: If the file has no parser and is not text.
    """
    try:
        enforce_parseable_attachment(local_path, filename, parser_extensions)
    except UnsupportedUploadTypeError as exc:
        raise AttachmentRejectedError(str(exc)) from exc


def _reject_attachment_zip_bomb(local_path: str) -> None:
    """Reject a zip-container attachment that decompresses to too much.

    Applies the shared guard (``document_reader.reject_zip_bomb_path``) on the
    attachment path, which previously had no such check.

    Args:
        local_path: Filesystem path of the attachment about to be parsed.

    Raises:
        AttachmentRejectedError: If the archive exceeds the entry-count or
            inner-uncompressed-size caps.
    """
    from docsgpt.parser.document_reader import reject_zip_bomb_path

    reason = reject_zip_bomb_path(local_path)
    if reason is not None:
        raise AttachmentRejectedError(reason)


def _readable_without_text(filename: str) -> bool:
    """Whether a model can read an attachment from its original file alone.

    PDFs and images are sent to models as the file itself (a PDF as page
    images on image-only models); every other format needs extracted text.

    Args:
        filename: The upload's original filename.

    Returns:
        bool: True for PDFs and images.
    """
    mime_type = mimetypes.guess_type(filename)[0] or ""
    return mime_type == "application/pdf" or mime_type.startswith("image/")


def _store_png_copy(storage, relative_path: str, mime_type: str) -> tuple[str, dict]:
    """Store a PNG copy, beside the original, of an image the providers reject.

    Args:
        storage: The storage backend holding the upload.
        relative_path: Storage path of the original image.
        mime_type: The original's mime type, e.g. ``image/tiff``.

    Returns:
        tuple[str, dict]: The PNG's storage path, and the conversion record
        kept in the attachment's metadata.

    Raises:
        DocumentParseError: If the image cannot be decoded.
    """
    with storage.get_file(relative_path) as source:
        png_bytes, frames = convert_image_to_png(source)
    png_path = f"{os.path.splitext(relative_path)[0]}.png"
    storage.save_file(io.BytesIO(png_bytes), png_path)
    return png_path, {"from": mime_type, "to": "image/png", "frames": frames}


def _bounded_attachment_copy(local_path: str) -> tuple[str, bool]:
    """Bound how much of a text attachment reaches the parser.

    Attachment content is capped at ~250k chars after parsing, so bytes past
    ``ATTACHMENT_TEXT_MAX_BYTES`` only cost parse time and memory. Oversized
    line-oriented text files are head-truncated on a line boundary into a
    temp copy; the stored original is never modified (local storage hands the
    canonical file path to ``process_file`` callbacks).

    Args:
        local_path: Filesystem path handed to the parse callback.

    Returns:
        Tuple of (path to parse, whether it is a temp copy the caller must
        delete).
    """
    max_bytes = settings.ATTACHMENT_TEXT_MAX_BYTES
    if max_bytes <= 0:
        return local_path, False
    suffix = os.path.splitext(local_path)[1].lower()
    if suffix not in _TRUNCATABLE_ATTACHMENT_SUFFIXES:
        return local_path, False
    try:
        if os.path.getsize(local_path) <= max_bytes:
            return local_path, False
        with open(local_path, "rb") as src:
            head = src.read(max_bytes)
    except OSError:
        return local_path, False
    head = truncate_to_line_boundary(head)
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(head)
    logging.warning(
        f"Attachment {os.path.basename(local_path)} exceeds "
        f"ATTACHMENT_TEXT_MAX_BYTES ({max_bytes}); parsing first {len(head)} bytes"
    )
    return tmp.name, True


def _upsert_attachment_row(
    user, filename, relative_path, *, mime_type, content, token_count, metadata, attachment_id
):
    """Create or update the attachment row for one upload handle.

    The upload route mints a UUID-shaped ``attachment_id`` (stored in the
    storage path); the PG ``attachments.id`` is DB-generated, so the handle
    lives in ``legacy_mongo_id``. Task retries share the handle — an earlier
    attempt's failure row is updated in place, never duplicated.
    """
    with db_session() as conn:
        repo = AttachmentsRepository(conn)
        existing = repo.get_by_legacy_id(str(attachment_id), user)
        if existing:
            repo.update(
                existing["id"],
                user,
                {
                    "filename": filename,
                    "upload_path": relative_path,
                    "mime_type": mime_type,
                    "content": content,
                    "token_count": token_count,
                    "metadata": metadata,
                },
            )
        else:
            repo.create(
                user,
                filename,
                relative_path,
                mime_type=mime_type,
                content=content,
                token_count=token_count,
                metadata=metadata,
                legacy_mongo_id=str(attachment_id),
            )


def record_attachment_failure(user, file_info, error, parser=None):
    """Persist a failure row so a broken parse is visible to a DB scan.

    Called on the worker's error path and by the poison guard; never raises —
    the original exception must keep propagating, and the parser's error text
    goes into ``metadata.extraction``, never into ``content``.
    """
    attachment_id = file_info.get("attachment_id")
    filename = file_info.get("filename") or ""
    if not attachment_id:
        return
    try:
        metadata = {
            **(file_info.get("metadata") or {}),
            "extraction": {
                "status": "failed",
                "parser": parser,
                "truncated": False,
                "error": str(error)[:1024],
            },
        }
        with db_session() as conn:
            repo = AttachmentsRepository(conn)
            # The conditional update carries the no-clobber guard in its own
            # WHERE clause: a success row committed by a concurrent duplicate
            # execution (broker redelivery, the poison guard racing a live
            # attempt) or earlier in this attempt is never overwritten with a
            # NULL-content failure.
            if repo.update_metadata_if_content_null(str(attachment_id), user, metadata):
                return
            if repo.get_by_legacy_id(str(attachment_id), user) is not None:
                return
            repo.create(
                user,
                filename,
                file_info.get("path") or "",
                mime_type=mimetypes.guess_type(filename)[0] or "application/octet-stream",
                content=None,
                token_count=None,
                metadata=metadata,
                legacy_mongo_id=str(attachment_id),
            )
    except Exception:
        logging.error(
            f"Failed to record failure row for attachment {attachment_id}",
            extra={"user": user},
            exc_info=True,
        )


def attachment_worker(self, file_info, user):
    """
    Process and store a single attachment without vectorization.
    """

    filename = file_info["filename"]
    attachment_id = file_info["attachment_id"]
    relative_path = file_info["path"]
    metadata = file_info.get("metadata", {})
    parser_name = None

    publish_user_event(
        user,
        "attachment.queued",
        {"attachment_id": str(attachment_id), "filename": filename},
        scope={"kind": "attachment", "id": str(attachment_id)},
    )

    try:
        self.update_state(state="PROGRESS", meta={"current": 10})
        storage = StorageCreator.get_storage()

        self.update_state(
            state="PROGRESS", meta={"current": 30, "status": "Processing content"}
        )
        publish_user_event(
            user,
            "attachment.progress",
            {
                "attachment_id": str(attachment_id),
                "filename": filename,
                "current": 30,
                "stage": "processing",
            },
            scope={"kind": "attachment", "id": str(attachment_id)},
        )

        # Attachments only: under the docling engine PDFs are read via their
        # text layer where one exists (``ATTACHMENT_PDF_TEXT_FAST_PATH``; a
        # no-op under anydoc, which is already that fast). Source ingestion
        # calls SimpleDirectoryReader without a ``file_extractor`` and so keeps
        # the ``DOC_PARSER_ENGINE`` default map, which retrieval quality depends on.
        file_extractor = get_default_file_extractor(
            ocr_enabled=settings.OCR_ATTACHMENTS_ENABLED,
            pdf_text_fast_path=settings.ATTACHMENT_PDF_TEXT_FAST_PATH,
        )
        _parser = file_extractor.get(os.path.splitext(filename)[1].lower())
        parser_name = type(_parser).__name__ if _parser is not None else "SimpleDirectoryReader"

        def _parse_local_file(local_path: str, **kwargs) -> Document:
            _reject_unparseable_attachment(local_path, filename, set(file_extractor))
            _reject_attachment_zip_bomb(local_path)
            parse_path, is_temp_copy = _bounded_attachment_copy(local_path)
            try:
                return SimpleDirectoryReader(
                    input_files=[parse_path],
                    exclude_hidden=True,
                    errors="ignore",
                    file_extractor=file_extractor,
                    file_metadata=metadata_from_filename,
                ).load_data()[0]
            finally:
                if is_temp_copy:
                    try:
                        os.unlink(parse_path)
                    except OSError:
                        pass

        extraction_status = "ok"
        no_text_reason = None
        try:
            attachment_document = storage.process_file(relative_path, _parse_local_file)
        except NoTextLayerError as exc:
            # A PDF with no text layer is still a usable attachment: models
            # that read PDFs natively, or as page images, are sent the stored
            # file rather than its text. Keep it with nothing to inline; the
            # LLM handler names it to any model that cannot read it.
            if not _readable_without_text(filename):
                raise
            logging.info(
                f"Attachment {filename} has no text layer; keeping it for native reading",
                extra={"user": user},
            )
            attachment_document = Document(text="", extra_info={})
            extraction_status = "no_text"
            no_text_reason = str(exc)[:1024]
        content = attachment_document.text
        # A fast-path parser may have delegated to its fallback for this file,
        # so record the engine that actually ran rather than the one selected.
        parser_name = getattr(_parser, "last_engine", None) or parser_name
        parser_metadata = {
            key: value
            for key, value in (attachment_document.extra_info or {}).items()
            if key.startswith("transcript_") or key == "parse_warnings"
        }
        if parser_metadata:
            metadata = {**metadata, **parser_metadata}

        # Gate and cut in the same unit. The old form gated on tokens but cut
        # at 250k *chars*, which for dense scripts (CJK ~1.4 tokens/char)
        # stored 300k+ tokens while looking like a clean extraction.
        encoding = get_encoding()
        tokens = encoding.encode_ordinary(content)
        original_tokens = len(tokens)
        truncated = original_tokens > ATTACHMENT_MAX_TOKENS
        if truncated:
            content = encoding.decode(tokens[:ATTACHMENT_MAX_TOKENS])
            token_count = ATTACHMENT_MAX_TOKENS
        else:
            token_count = original_tokens

        metadata = {
            **metadata,
            "extraction": {
                "status": extraction_status,
                "parser": parser_name,
                "truncated": truncated,
                "original_tokens": original_tokens,
                "stored_tokens": token_count,
                **({"reason": no_text_reason} if no_text_reason else {}),
            },
        }

        self.update_state(
            state="PROGRESS", meta={"current": 80, "status": "Storing in database"}
        )
        publish_user_event(
            user,
            "attachment.progress",
            {
                "attachment_id": str(attachment_id),
                "filename": filename,
                "current": 80,
                "stage": "storing",
            },
            scope={"kind": "attachment", "id": str(attachment_id)},
        )

        mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
        if mime_type in VISION_CONVERTIBLE_MIME_TYPES:
            relative_path, conversion = _store_png_copy(storage, relative_path, mime_type)
            mime_type = conversion["to"]
            metadata = {**metadata, "image_conversion": conversion}

        _upsert_attachment_row(
            user,
            filename,
            relative_path,
            mime_type=mime_type,
            content=content,
            token_count=token_count,
            metadata=metadata,
            attachment_id=attachment_id,
        )

        logging.info(
            f"Stored attachment with ID: {attachment_id}", extra={"user": user}
        )

        self.update_state(state="PROGRESS", meta={"current": 100, "status": "Complete"})

        publish_user_event(
            user,
            "attachment.completed",
            {
                "attachment_id": str(attachment_id),
                "filename": filename,
                "token_count": token_count,
                "mime_type": mime_type,
                "extraction_status": extraction_status,
            },
            scope={"kind": "attachment", "id": str(attachment_id)},
        )

        return {
            "filename": filename,
            "path": relative_path,
            "token_count": token_count,
            "attachment_id": attachment_id,
            "mime_type": mime_type,
            "metadata": metadata,
        }
    except Exception as e:
        logging.error(
            f"Error processing file {filename}: {e}",
            extra={"user": user},
            exc_info=True,
        )
        record_attachment_failure(user, file_info, e, parser=parser_name)
        publish_user_event(
            user,
            "attachment.failed",
            {
                "attachment_id": str(attachment_id),
                "filename": filename,
                "error": str(e)[:1024],
            },
            scope={"kind": "attachment", "id": str(attachment_id)},
        )
        raise


def parse_document_worker(self, artifact_id, parent, user_id, options):
    """Thin Celery-task wrapper; delegates to the process-agnostic ``run_parse_document``."""
    return run_parse_document(artifact_id, parent, user_id, options)


def run_parse_document(artifact_id, parent, user_id, options):
    """Parse an input artifact's bytes to a shaped result; runnable inline OR on the parsing queue.

    Security: the artifact is re-resolved through the run-scoped gate here (never trusting a
    raw storage path) so authz is enforced independently, in addition to the pre-enqueue check
    in the tool. ``read_document`` calls this directly (in-process) when it already runs inside a
    Celery worker; the web process dispatches ``parse_document`` to the parsing queue, which lands
    here via ``parse_document_worker``.
    """
    from docsgpt.agents.tools.artifact_ref import resolve_artifact_id
    from docsgpt.parser.document_reader import bound_parse_payload, parse_document_bytes

    options = options or {}
    parent = parent or {}
    conversation_id = parent.get("conversation_id")
    workflow_run_id = parent.get("workflow_run_id")
    if conversation_id is None and workflow_run_id is None:
        return {"status": "error", "error": "parse_document requires a conversation_id or workflow_run_id."}

    # Re-resolve through the parent-scoped gate so a forged/cross-run id is rejected
    # here too; resolve a short ref to an id within this parent only.
    try:
        with db_readonly() as conn:
            repo = ArtifactsRepository(conn)
            resolved_id = resolve_artifact_id(
                repo, artifact_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
            )
            artifact = (
                repo.get_artifact_in_parent(
                    resolved_id, conversation_id=conversation_id, workflow_run_id=workflow_run_id
                )
                if resolved_id is not None
                else None
            )
            if artifact is None:
                return {"status": "error", "error": f"input artifact {artifact_id} not found in this conversation/run."}
            version = repo.get_version(resolved_id, artifact["current_version"])
    except Exception:
        logging.error("run_parse_document: failed to resolve input artifact", exc_info=True)
        return {"status": "error", "error": f"failed to load input artifact {artifact_id}."}

    if not version or not version.get("storage_path"):
        return {"status": "error", "error": f"input artifact {artifact_id} has no stored content."}

    display_name = version.get("filename") or artifact.get("title") or str(resolved_id)
    filename = safe_filename(display_name)
    try:
        data = StorageCreator.get_storage().get_file(version["storage_path"]).read()
    except Exception:
        logging.error("run_parse_document: failed to read input artifact bytes", exc_info=True)
        return {"status": "error", "error": f"failed to read input artifact {artifact_id}."}

    # Parse returns the FULL content (no max_chars/window here) so the persisted artifact is the
    # complete parse; all view-bounding is applied by ``bound_parse_payload`` below.
    result = parse_document_bytes(
        data,
        filename,
        output=options.get("output", "markdown"),
        ocr=options.get("ocr", "auto"),
        pages=options.get("pages"),
        engine=options.get("engine", "auto"),
        include_tables=bool(options.get("include_tables", True)),
    )
    if result.get("error"):
        return {"status": "error", "error": result["error"]}

    payload = {"status": "ok", **result}
    if options.get("persist"):
        # Persist the FULL shaped result by reference (bytes live in the artifact); only the
        # bounded view computed below rides back through the Redis result backend.
        artifact_ref = _persist_parse_result(result, display_name, user_id, parent, options)
        if isinstance(artifact_ref, dict) and artifact_ref.get("error"):
            payload["artifact_error"] = artifact_ref["error"]
        elif artifact_ref is not None:
            payload["artifact"] = artifact_ref
    # Bound the Redis-backed VIEW across all shapes: content is capped by max_chars (else the
    # default head+tail window), chunks are count/length-capped, and structured (needed for
    # json_schema validation) rides back as-is. The FULL result already lives in the artifact.
    payload = bound_parse_payload(payload, max_chars=options.get("max_chars"))
    return payload


def _persist_parse_result(result, title, user_id, parent, options):
    """Persist the full shaped parse result as an owner/parent-scoped ``data`` artifact; return its ref."""
    from docsgpt.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact

    try:
        data = json.dumps(result).encode("utf-8")
    except (TypeError, ValueError):
        logging.error("parse_document_worker: parse result is not JSON-serializable", exc_info=True)
        return {"error": "parse result is not JSON-serializable."}
    base = safe_filename(title) or "document"
    filename = f"{base}.parsed.json"
    try:
        return persist_new_artifact(
            user_id=user_id,
            kind="data",
            data=data,
            filename=filename,
            mime_type="application/json",
            title=f"{title} (parsed)",
            conversation_id=parent.get("conversation_id"),
            workflow_run_id=parent.get("workflow_run_id"),
            message_id=parent.get("message_id"),
            produced_by={"tool": "read_document", "action": "read_document", "tool_id": options.get("tool_id")},
        )
    except QuotaExceeded as exc:
        return {"error": str(exc)}


def agent_webhook_worker(self, agent_id, payload):
    """Process the webhook payload for an agent.

    Raises on failure: Celery treats a returned dict as success and
    would skip retries, leaving the caller with a stale 200.
    """
    self.update_state(state="PROGRESS", meta={"current": 1})
    try:
        with db_readonly() as conn:
            repo = AgentsRepository(conn)
            agent_config = None
            if looks_like_uuid(str(agent_id)):
                # Access without user scoping — webhooks authenticate via
                # the incoming token, not a user context.
                from sqlalchemy import text as sql_text
                from docsgpt.storage.db.base_repository import row_to_dict
                result = conn.execute(
                    sql_text("SELECT * FROM agents WHERE id = CAST(:id AS uuid)"),
                    {"id": str(agent_id)},
                )
                row = result.fetchone()
                if row is not None:
                    agent_config = row_to_dict(row)
            if agent_config is None:
                agent_config = repo.get_by_legacy_id(str(agent_id))
        if not agent_config:
            raise ValueError(f"Agent with ID {agent_id} not found.")
        input_data = json.dumps(payload)
    except Exception as e:
        logging.error(f"Error processing agent webhook: {e}", exc_info=True)
        raise
    self.update_state(state="PROGRESS", meta={"current": 50})
    try:
        # Shared headless path with the scheduler; approval-gated tools auto-deny.
        from docsgpt.agents.headless_runner import run_agent_headless
        from docsgpt.quotas.service import QuotaExceededError

        outcome = run_agent_headless(
            agent_config,
            input_data,
            tool_allowlist=_webhook_tool_allowlist(agent_config),
            endpoint="webhook",
        )
        result = {
            "answer": outcome.get("answer", ""),
            "sources": outcome.get("sources", []),
            "tool_calls": outcome.get("tool_calls", []),
            "thought": outcome.get("thought", ""),
        }
    except QuotaExceededError as e:
        # Returned, not raised: retrying cannot succeed before the quota resets.
        logging.warning(
            f"Webhook skipped for agent {agent_id}: {e}", extra={"agent_id": agent_id}
        )
        return {"status": "quota_exceeded", "error": str(e)}
    except Exception as e:
        logging.error(f"Error running agent logic: {e}", exc_info=True)
        raise
    else:
        logging.info(
            f"Webhook processed for agent {agent_id}", extra={"agent_id": agent_id}
        )
        return {"status": "success", "result": result}
    finally:
        self.update_state(state="PROGRESS", meta={"current": 100})


def _webhook_tool_allowlist(agent_config):
    """Deny-all on approval-gated tools for webhooks (per-agent opt-in is TBD)."""
    return []


def ingest_connector(
    self,
    job_name: str,
    user: str,
    source_type: str,
    session_token=None,
    file_ids=None,
    folder_ids=None,
    recursive=True,
    retriever: str = "classic",
    operation_mode: str = "upload",
    doc_id=None,
    sync_frequency: str = "never",
    config=None,
    idempotency_key=None,
    source_id=None,
) -> Dict[str, Any]:
    """
    Ingestion for internal knowledge bases (GoogleDrive, etc.).

    Args:
        job_name: Name of the ingestion job
        user: User identifier
        source_type: Type of remote source ("google_drive", "dropbox", etc.)
        session_token: Authentication token for the service
        file_ids: List of file IDs to download
        folder_ids: List of folder IDs to download
        recursive: Whether to recursively download folders
        retriever: Type of retriever to use
        operation_mode: "upload" for initial ingestion, "sync" for incremental sync
        doc_id: Document ID for sync operations (required when operation_mode="sync")
        sync_frequency: How often to sync ("never", "daily", "weekly", "monthly")
        config: Per-source ``SourceConfig`` dict. ``None``/``{}`` → classic
            defaults (byte-identical to prior behavior).
        idempotency_key: When provided, the ``source_id`` is derived
            deterministically so a retried upload reuses the same source row.
        source_id: When supplied, the worker uses it verbatim so SSE envelopes
            carry the same id the HTTP route already returned to the frontend
            — required for non-idempotent uploads where the route can't
            predict ``_derive_source_id(idempotency_key)``.
    """
    logging.info(
        f"Starting remote ingestion from {source_type} for user: {user}, job: {job_name}"
    )

    # Source id resolution mirrors ``ingest_worker`` / ``remote_worker``:
    # sync mode reuses ``doc_id``; otherwise the caller-supplied
    # ``source_id`` (minted by the HTTP route and already echoed to the
    # client) wins; fall back to ``_derive_source_id`` only when neither
    # is supplied. Without rule (2) the no-idempotency-key path would
    # mint a fresh uuid4 here that the frontend has no way to correlate
    # SSE envelopes to.
    if operation_mode == "sync" and doc_id:
        source_uuid = str(doc_id)
    elif source_id:
        source_uuid = uuid.UUID(source_id)
    else:
        source_uuid = _derive_source_id(idempotency_key)
    source_id_for_events = str(source_uuid)

    # First-attempt gate: Celery retries re-run the body, and a
    # repeated ``queued`` here would oscillate the toast through
    # ``queued`` again between ``failed`` and ``completed``.
    if self.request.retries == 0:
        publish_user_event(
            user,
            "source.ingest.queued",
            {
                "source_id": source_id_for_events,
                "job_name": job_name,
                "loader": source_type,
                "operation": operation_mode,
            },
            scope={"kind": "source", "id": source_id_for_events},
        )

    self.update_state(state="PROGRESS", meta={"current": 1})

    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            # Step 1: Initialize the appropriate loader
            self.update_state(
                state="PROGRESS",
                meta={"current": 10, "status": "Initializing connector"},
            )

            if not session_token:
                raise ValueError(f"{source_type} connector requires session_token")

            if not ConnectorCreator.is_supported(source_type):
                raise ValueError(
                    f"Unsupported connector type: {source_type}. Supported types: {ConnectorCreator.get_supported_connectors()}"
                )

            remote_loader = ConnectorCreator.create_connector(
                source_type, session_token
            )

            # Create a clean config for storage
            api_source_config = {
                "file_ids": file_ids or [],
                "folder_ids": folder_ids or [],
                "recursive": recursive,
            }

            # Step 2: Download files to temp directory
            self.update_state(
                state="PROGRESS", meta={"current": 20, "status": "Downloading files"}
            )
            download_info = remote_loader.download_to_directory(
                temp_dir, api_source_config
            )

            if download_info.get("empty_result", False) or not download_info.get(
                "files_downloaded", 0
            ):
                logging.warning(f"No files were downloaded from {source_type}")
                # Connector returned no files — surface as a benign
                # ``completed`` event with zero tokens so the toast
                # closes out cleanly instead of waiting on polling.
                publish_user_event(
                    user,
                    "source.ingest.completed",
                    {
                        "source_id": source_id_for_events,
                        "job_name": job_name,
                        "loader": source_type,
                        "operation": operation_mode,
                        "tokens": 0,
                        "no_changes": True,
                    },
                    scope={"kind": "source", "id": source_id_for_events},
                )
                # Create empty result directly instead of calling a separate method
                return {
                    "name": job_name,
                    "user": user,
                    "tokens": 0,
                    "type": source_type,
                    "source_config": api_source_config,
                    "directory_structure": "{}",
                }

            # Step 3: Use SimpleDirectoryReader to process downloaded files
            self.update_state(
                state="PROGRESS", meta={"current": 40, "status": "Processing files"}
            )
            reader = SimpleDirectoryReader(
                input_dir=temp_dir,
                recursive=True,
                required_exts=list(SUPPORTED_SOURCE_EXTENSIONS),
                exclude_hidden=True,
                file_metadata=metadata_from_filename,
            )
            # Parsing/OCR fills 40-60% of the bar; embedding takes 60-100%.
            raw_docs = reader.load_data(
                progress_callback=_make_parse_progress_callback(
                    self, user, source_uuid, start_pct=40, end_pct=60,
                )
            )
            directory_structure = getattr(reader, "directory_structure", {})

            # Step 4: Process documents (chunking, embedding, etc.)
            cfg = SourceConfig.parse(config)
            chunker = ChunkerCreator.create_chunker(
                cfg.chunking.strategy,
                chunking_strategy=cfg.chunking.strategy,
                max_tokens=cfg.chunking.max_tokens,
                min_tokens=cfg.chunking.min_tokens,
                duplicate_headers=cfg.chunking.duplicate_headers,
            )
            raw_docs = chunker.chunk(documents=raw_docs)

            # Preserve source information in document metadata
            for doc in raw_docs:
                if hasattr(doc, "extra_info") and doc.extra_info:
                    source = doc.extra_info.get("source")
                    if source and os.path.isabs(source):
                        # Convert absolute path to relative path
                        doc.extra_info["source"] = os.path.relpath(
                            source, start=temp_dir
                        )

            docs = [Document.to_vector_format(raw_doc) for raw_doc in raw_docs]

            # Validate operation_mode here too (the source_uuid path
            # at the top of the function only branches on the
            # sync+doc_id combination; surfacing the wrong-mode error
            # this far in matches the legacy behaviour).
            if operation_mode == "sync" and not doc_id:
                logging.error(
                    "Invalid doc_id provided for sync operation: %s", doc_id
                )
                raise ValueError("doc_id must be provided for sync operation.")
            if operation_mode not in ("upload", "sync"):
                raise ValueError(f"Invalid operation_mode: {operation_mode}")

            vector_store_path = os.path.join(temp_dir, "vector_store")
            os.makedirs(vector_store_path, exist_ok=True)

            self.update_state(
                state="PROGRESS", meta={"current": 60, "status": "Storing documents"}
            )
            embed_and_store_documents(
                docs, vector_store_path, source_uuid, self,
                attempt_id=getattr(self.request, "id", None),
                user_id=user,
                progress_start=60, progress_end=100,
            )
            assert_index_complete(source_uuid)

            tokens = count_tokens_docs(docs)

            # Step 6: Upload index files
            file_data = {
                "user": user,
                "name": job_name,
                "tokens": tokens,
                "retriever": retriever,
                "id": source_id_for_events,
                "type": "connector:file",
                "remote_data": json.dumps(
                    {"provider": source_type, **api_source_config}
                ),
                "directory_structure": json.dumps(directory_structure),
                "sync_frequency": sync_frequency,
            }
            if config:
                file_data["config"] = json.dumps(config)

            file_data["last_sync"] = datetime.datetime.now(datetime.timezone.utc)

            if operation_mode == "sync":
                try:
                    with db_session() as conn:
                        repo = SourcesRepository(conn)
                        src = repo.get_any(source_id_for_events, user)
                        if src is not None:
                            repo.update(
                                str(src["id"]), user,
                                {"date": file_data["last_sync"]},
                            )
                except Exception as upd_err:
                    logging.warning(
                        "Failed to update last_sync for source %s: %s",
                        source_id_for_events,
                        upd_err,
                    )

            upload_index(vector_store_path, file_data)

            # Ensure we mark the task as complete
            self.update_state(
                state="PROGRESS", meta={"current": 100, "status": "Complete"}
            )

            logging.info(f"Remote ingestion completed: {job_name}")

            publish_user_event(
                user,
                "source.ingest.completed",
                {
                    "source_id": source_id_for_events,
                    "job_name": job_name,
                    "loader": source_type,
                    "operation": operation_mode,
                    "tokens": tokens,
                },
                scope={"kind": "source", "id": source_id_for_events},
            )
            _maybe_enqueue_graph_extraction(cfg, source_id_for_events, user)

            return {
                "user": user,
                "name": job_name,
                "tokens": tokens,
                "type": source_type,
                "id": source_id_for_events,
                "status": "complete",
            }
    except Exception as e:
        logging.error(f"Error during remote ingestion: {e}", exc_info=True)
        publish_user_event(
            user,
            "source.ingest.failed",
            {
                "source_id": source_id_for_events,
                "job_name": job_name,
                "loader": source_type,
                "operation": operation_mode,
                "error": str(e)[:1024],
            },
            scope={"kind": "source", "id": source_id_for_events},
        )
        raise


def mcp_oauth(self, config: Dict[str, Any], user_id: str = None) -> Dict[str, Any]:
    """Worker to handle MCP OAuth flow asynchronously.

    Publishes SSE events at each phase boundary so the frontend can
    drive the OAuth popup directly from the push channel. The
    ``mcp.oauth.awaiting_redirect`` envelope carries the
    ``authorization_url`` once the upstream OAuth client surfaces it,
    eliminating the prior polling-only path for that URL.
    """

    # Bind ``task_id`` and the publish helpers OUTSIDE the outer try so
    # the ``except`` handler at the bottom can reach them even when an
    # early statement raises. Without this, ``publish_oauth`` would
    # UnboundLocalError on top of the original failure.
    task_id = self.request.id if getattr(self, "request", None) else None

    def publish_oauth(event_type: str, payload: Dict[str, Any]) -> None:
        # MCP OAuth can be invoked without a route-bound user_id by
        # legacy paths. Skip the SSE publish in that case \u2014 the caller
        # has no per-user channel to subscribe to, and the status is
        # surfaced via the task's return value.
        if not user_id or task_id is None:
            return
        publish_user_event(
            user_id,
            event_type,
            {"task_id": task_id, **payload},
            scope={"kind": "mcp_oauth", "id": task_id},
        )

    def publish_awaiting_redirect(authorization_url: str) -> None:
        """Callback invoked by ``DocsGPTOAuth.redirect_handler`` once
        the OAuth client has minted the authorization URL.

        Carrying the URL on the SSE envelope lets the frontend open the
        popup directly from the event \u2014 the prior polling-only path
        for the URL is gone.
        """
        publish_oauth(
            "mcp.oauth.awaiting_redirect",
            {
                "message": "Awaiting OAuth redirect...",
                "authorization_url": authorization_url,
            },
        )

    try:
        import asyncio

        from docsgpt.agents.tools.mcp_tool import MCPTool

        publish_oauth("mcp.oauth.in_progress", {"message": "Starting OAuth..."})

        tool_config = config.copy()
        tool_config["oauth_task_id"] = task_id
        # Inject the awaiting-redirect publish callback. ``MCPTool`` pops
        # it out of the config and threads it into ``DocsGPTOAuth`` so
        # the publish fires synchronously from inside
        # ``redirect_handler`` \u2014 the only point where the URL is known.
        tool_config["oauth_redirect_publish"] = publish_awaiting_redirect
        mcp_tool = MCPTool(tool_config, user_id)

        async def run_oauth_discovery():
            if not mcp_tool._client:
                mcp_tool._setup_client()
            return await mcp_tool._execute_with_client("list_tools")

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            loop.run_until_complete(run_oauth_discovery())
            tools = mcp_tool.get_actions_metadata()

            publish_oauth(
                "mcp.oauth.completed",
                {"tools": tools, "tools_count": len(tools)},
            )

            return {"success": True, "tools": tools, "tools_count": len(tools)}
        except Exception as e:
            error_msg = f"OAuth failed: {str(e)}"
            logging.error("MCP OAuth discovery failed: %s", error_msg, exc_info=True)
            publish_oauth("mcp.oauth.failed", {"error": error_msg[:1024]})
            return {"success": False, "error": error_msg}
        finally:
            loop.close()
    except Exception as e:
        error_msg = f"OAuth init failed: {str(e)}"
        logging.error("MCP OAuth init failed: %s", error_msg, exc_info=True)
        publish_oauth("mcp.oauth.failed", {"error": error_msg[:1024]})
        return {"success": False, "error": error_msg}


def reembed_wiki_page_worker(self, source_id, path, content_hash, user):
    """Re-embed one wiki page after an edit, or purge its chunks on delete.

    Targeted delete of the page's existing chunks runs first, so a deleted
    page becomes a pure purge and an edited page is re-embedded cleanly. The
    chunk metadata shape matches the reingest path (``source``/``title``/
    ``filename``) so retrieval and future targeted deletes line up.

    Args:
        self: Celery task instance.
        source_id: Source the page belongs to.
        path: Page path; doubles as the chunk ``source`` metadata key.
        content_hash: The edit's content hash (idempotency anchor).
        user: Owner identifier used to load the source row.

    Returns:
        A status dict: ``{"status", "added", "deleted"}`` on re-embed, or
        ``{"status": "deleted", "deleted": n}`` when the page is gone.
    """
    from docsgpt.vectorstore.vector_creator import VectorCreator

    source_id = str(source_id)

    with db_readonly() as conn:
        source = SourcesRepository(conn).get_any(source_id, user)
    if not source:
        raise ValueError(f"Source {source_id} not found or access denied")
    source_id = str(source["id"])
    cfg = SourceConfig.parse(source.get("config"))

    store = VectorCreator.create_vectorstore(
        settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
    )
    deleted = store.delete_chunks_by_source_path(path)

    with db_readonly() as conn:
        page = WikiPagesRepository(conn).get_by_path(source_id, path)

    if page is None:
        return {"status": "deleted", "deleted": deleted}

    try:
        title = page.get("title") or path
        chunker = ChunkerCreator.create_chunker(
            cfg.chunking.strategy,
            chunking_strategy=cfg.chunking.strategy,
            max_tokens=cfg.chunking.max_tokens,
            min_tokens=cfg.chunking.min_tokens,
            duplicate_headers=cfg.chunking.duplicate_headers,
        )
        chunks = chunker.chunk(
            documents=[
                Document(
                    text=page["content"],
                    extra_info={
                        "source": path,
                        "title": title,
                        "filename": path,
                    },
                )
            ]
        )

        added = 0
        for chunk in chunks:
            store.add_chunk(
                chunk.text,
                metadata={"source": path, "title": title, "filename": path},
            )
            added += 1

        with db_session() as conn:
            WikiPagesRepository(conn).set_embed_status(source_id, path, "embedded")
            # These chunks were just embedded with the configured model, so the
            # source now names it. Wiki sources created before this was recorded
            # carry NULL, which the boot mismatch check reads as the legacy
            # model and reports as stale; stamping here heals them on the next
            # page edit.
            SourcesRepository(conn).update(
                source_id, user, {"model": settings.EMBEDDINGS_NAME}
            )
    except Exception:
        with db_session() as conn:
            WikiPagesRepository(conn).set_embed_status(source_id, path, "failed")
        raise

    return {"status": "embedded", "added": added, "deleted": deleted}


def _wiki_page_path_from_rel(rel_path):
    """Build a validated leading-slash wiki page path, or None if invalid.

    Runs the derived path through ``validate_tool_path`` so a stored filename
    carrying traversal (``..``) never lands as a ``wiki_pages.path``.
    """
    from docsgpt.agents.tools.path_utils import validate_tool_path

    raw = "/" + rel_path.replace(os.sep, "/").lstrip("/")
    return validate_tool_path(raw)


def _chunk_text(chunk) -> str:
    if isinstance(chunk, dict):
        return chunk.get("text") or chunk.get("page_content") or ""
    return getattr(chunk, "text", None) or getattr(chunk, "page_content", "") or ""


def _chunk_metadata(chunk) -> dict:
    if isinstance(chunk, dict):
        meta = chunk.get("metadata")
    else:
        meta = getattr(chunk, "metadata", None)
    return meta if isinstance(meta, dict) else {}


def _chunk_doc_id(chunk):
    if isinstance(chunk, dict):
        return chunk.get("doc_id")
    return getattr(chunk, "doc_id", None)


def _url_to_virtual_path(value: str) -> str:
    parts = urlsplit(value)
    path = parts.path.strip("/")
    if path:
        return path
    if parts.netloc:
        return parts.netloc
    return value.replace("://", "/")


def _chunk_page_path(metadata) -> str:
    for key in ("file_path", "file_name", "filename", "title", "source"):
        value = metadata.get(key)
        if value:
            value = str(value)
            if "://" in value:
                value = _url_to_virtual_path(value)
            return value
    return ""


def _chunk_order_hint(metadata):
    for key in ("chunk", "chunk_index", "index", "start", "start_index", "offset"):
        value = metadata.get(key)
        if isinstance(value, (int, float)):
            return value
        if isinstance(value, str):
            try:
                return float(value)
            except ValueError:
                continue
    return None


MIN_CHUNK_OVERLAP_TRIM = 32


def _join_chunk_texts(texts: list[str]) -> str:
    """Concatenate chunk texts, trimming an exact overlapping boundary.

    Only an exact suffix/prefix overlap of at least ``MIN_CHUNK_OVERLAP_TRIM``
    characters between consecutive chunks is removed (the common
    ``chunk_overlap`` case); shorter, possibly legitimate repeats are kept.
    """
    parts: list[str] = []
    for text in texts:
        if parts:
            prev = parts[-1]
            limit = min(len(prev), len(text))
            overlap = 0
            for size in range(limit, MIN_CHUNK_OVERLAP_TRIM - 1, -1):
                if prev.endswith(text[:size]):
                    overlap = size
                    break
            if overlap:
                text = text[overlap:]
        if text:
            parts.append(text)
    return "\n\n".join(parts)


def convert_source_to_wiki_worker(self, source_id, user):
    """Convert an ingested source into a wiki by reassembling pages from chunks.

    Groups the source's existing vector-store chunks by their page path
    (``metadata.source``), concatenating each group into one wiki page. This
    works uniformly for every source type, including crawler/remote/connector
    sources that hold no stored originals. Chunks with a missing or invalid
    path are skipped and reported. Embedding is enqueued per page (async); the
    source's ``config.kind`` flips to ``wiki`` only after at least one page is
    materialized — a source with no usable chunks keeps its original kind and
    directory structure.

    Args:
        self: Celery task instance.
        source_id: ID of the source to convert.
        user: Owner identifier used to load the source row and author pages.

    Returns:
        A summary dict ``{"status": "converted", "pages_created", "skipped"}``,
        ``{"status": "no_pages", ...}`` when nothing was reassembled, or
        ``{"status": "already_wiki"}`` when the source is already a wiki.
    """
    from docsgpt.api.user.tasks import reembed_wiki_page
    from docsgpt.vectorstore.vector_creator import VectorCreator

    source_id = str(source_id)

    with db_readonly() as conn:
        source = SourcesRepository(conn).get_any(source_id, user)
    if not source:
        raise ValueError(f"Source {source_id} not found or access denied")
    source_id = str(source["id"])

    cfg = SourceConfig.parse(source.get("config"))
    if cfg.kind == "wiki":
        return {"status": "already_wiki", "pages_created": 0, "skipped": []}

    file_name_map = _normalize_file_name_map(source.get("file_name_map"))

    created_pages: list[tuple[str, str]] = []
    skipped: list[dict] = []

    store = VectorCreator.create_vectorstore(
        settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
    )

    grouped: dict[str, list[tuple[object, str]]] = {}
    original_doc_ids: list[object] = []
    for chunk in store.get_chunks() or []:
        # Track every original chunk so a skipped one (empty/no-path/invalid)
        # is purged on convert too, not left orphaned in the vector store after
        # the source flips to wiki. Only deleted once pages are created below.
        doc_id = _chunk_doc_id(chunk)
        if doc_id is not None:
            original_doc_ids.append(doc_id)
        metadata = _chunk_metadata(chunk)
        rel_path = _chunk_page_path(metadata)
        text = _chunk_text(chunk)
        if not text.strip():
            continue
        if not rel_path:
            skipped.append({"file": "", "reason": "missing path"})
            continue
        if _wiki_page_path_from_rel(rel_path) is None:
            skipped.append({"file": rel_path, "reason": "invalid path"})
            continue
        grouped.setdefault(rel_path, []).append((_chunk_order_hint(metadata), text))

    with db_session() as conn:
        repo = WikiPagesRepository(conn)
        for rel_path in sorted(grouped):
            page_path = _wiki_page_path_from_rel(rel_path)
            entries = grouped[rel_path]
            if any(hint is not None for hint, _ in entries):
                entries = sorted(
                    enumerate(entries),
                    key=lambda item: (
                        item[1][0] if item[1][0] is not None else float("inf"),
                        item[0],
                    ),
                )
                entries = [entry for _, entry in entries]
            content = _join_chunk_texts([text for _, text in entries])
            title = _get_display_name(file_name_map, rel_path) or os.path.basename(
                rel_path
            )
            repo.upsert(
                source_id,
                page_path,
                content,
                title=title,
                updated_by=user,
                updated_via="agent",
            )
            created_pages.append((page_path, _content_hash(content)))

    if not created_pages:
        return {
            "status": "no_pages",
            "pages_created": 0,
            "skipped": skipped,
        }

    with db_session() as conn:
        rebuild_wiki_directory_structure(conn, source_id, user)

    for doc_id in original_doc_ids:
        try:
            store.delete_chunk(doc_id)
        except Exception as exc:
            logging.error(f"Failed deleting original chunk {doc_id}: {exc}")

    for page_path, content_hash in created_pages:
        reembed_wiki_page.delay(
            source_id,
            page_path,
            content_hash,
            user=user,
            idempotency_key=f"reembed-wiki:{source_id}:{page_path}:{content_hash}",
        )

    with db_session() as conn:
        SourcesRepository(conn).update(
            source_id, user, {"config": cfg.wiki_enabled()}
        )

    return {
        "status": "converted",
        "pages_created": len(created_pages),
        "skipped": skipped,
    }


def extract_graph_worker(self, source_id, user):
    """Build a graphrag source's knowledge graph from its embedded chunks.

    Loads the source, fetches its chunks from the vector store, and runs the
    per-chunk LLM extraction pipeline. The chunks carry ``doc_id`` + ``text``,
    matching the retrievable ids the extraction pipeline links against. No-ops
    cleanly when GraphRAG is unavailable or the source has no chunks yet.

    Streams ``graph.extract.progress`` SSE events as chunks are processed and a
    terminal ``graph.extract.completed``/``graph.extract.failed`` on exit.

    Args:
        self: Celery task instance.
        source_id: Source whose graph is being built.
        user: Owner identifier used to load the source and attribute token usage.

    Returns:
        ``{"status": "unavailable"}`` when GraphRAG is off, otherwise the
        extraction summary ``{nodes, edges, chunks_processed, ...}``.
    """
    from docsgpt.graphrag import graphrag_available
    from docsgpt.graphrag.extraction import extract_graph_for_source
    from docsgpt.vectorstore.vector_creator import VectorCreator

    source_id = str(source_id)

    if not graphrag_available():
        return {"status": "unavailable"}

    with db_readonly() as conn:
        source = SourcesRepository(conn).get_any(source_id, user)
    if not source:
        raise ValueError(f"Source {source_id} not found or access denied")
    source_id = str(source["id"])
    cfg = SourceConfig.parse(source.get("config"))

    store = VectorCreator.create_vectorstore(
        settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
    )
    chunks = store.get_chunks() or []

    total = len(chunks)
    # Throttle: at most ~20 progress events regardless of chunk count.
    step = max(1, total // 20)

    def _progress(info):
        current = int(info.get("current", 0))
        if current and current % step != 0 and current != info.get("total"):
            return
        _publish_graph_event(
            user,
            source_id,
            "graph.extract.progress",
            {
                "source_id": source_id,
                "current": current,
                "total": int(info.get("total", total)),
                "nodes": int(info.get("nodes", 0)),
                "edges": int(info.get("edges", 0)),
            },
        )

    _publish_graph_event(
        user,
        source_id,
        "graph.extract.progress",
        {
            "source_id": source_id,
            "current": 0,
            "total": total,
            "nodes": 0,
            "edges": 0,
        },
    )

    try:
        summary = extract_graph_for_source(
            source_id,
            user,
            chunks,
            config=cfg,
            request_id=getattr(self.request, "id", None),
            progress_cb=_progress,
        )
    except Exception as e:
        _publish_graph_event(
            user,
            source_id,
            "graph.extract.failed",
            {"source_id": source_id, "error": str(e)[:1024]},
        )
        raise

    _publish_graph_event(
        user,
        source_id,
        "graph.extract.completed",
        {"source_id": source_id, **summary},
    )
    return summary