{%- if cookiecutter.enable_rag and (cookiecutter.use_celery or cookiecutter.use_taskiq or cookiecutter.use_arq or cookiecutter.use_prefect) %}
# ruff: noqa: I001 - Imports structured for Jinja2 template conditionals
"""RAG ingestion & sync tasks — processes documents asynchronously."""

import asyncio
import hashlib
import json
import logging
import tempfile
from pathlib import Path
from typing import Any

import redis.asyncio as aioredis
{%- if cookiecutter.use_celery %}
from celery import shared_task
{%- elif cookiecutter.use_taskiq %}
from app.worker.taskiq_app import broker
{%- elif cookiecutter.use_prefect %}
from prefect import flow
{%- endif %}

from app.core.config import settings
from app.db.session import get_worker_db_context
from app.repositories import sync_source as sync_source_repo
from app.services.rag.config import DocumentExtensions
from app.services.rag.connectors import CONNECTOR_REGISTRY
from app.services.rag.documents import DocumentProcessor
from app.services.rag.embeddings import EmbeddingService
from app.services.rag.ingestion import IngestionService
from app.services.sync_source import SyncSourceService
{%- if cookiecutter.use_milvus %}
from app.services.rag.vectorstore import MilvusVectorStore as VectorStore
{%- elif cookiecutter.use_qdrant %}
from app.services.rag.vectorstore import QdrantVectorStore as VectorStore
{%- elif cookiecutter.use_chromadb %}
from app.services.rag.vectorstore import ChromaVectorStore as VectorStore
{%- elif cookiecutter.use_pgvector %}
from app.services.rag.vectorstore import PgVectorStore as VectorStore
{%- endif %}

logger = logging.getLogger(__name__)


def _build_ingestion_service() -> IngestionService:
    rag_settings = settings.rag
    embed_service = EmbeddingService(settings=rag_settings)
    vector_store = VectorStore(settings=rag_settings, embedding_service=embed_service)
    processor = DocumentProcessor(settings=rag_settings)
    return IngestionService(processor=processor, vector_store=vector_store)


{%- if cookiecutter.use_celery %}


@shared_task(bind=True, max_retries=2, soft_time_limit=300, time_limit=360)  # type: ignore
def ingest_document_task(self: Any, rag_document_id: str, collection_name: str, filepath: str, source_path: str, replace: bool = False) -> dict[str, Any]:
    """Process a document: parse, chunk, embed, store in vector DB."""
    logger.info("Starting ingestion: %s -> %s", source_path, collection_name)
    try:
        return asyncio.run(_run_ingestion(rag_document_id, collection_name, filepath, source_path, replace))
    except Exception as exc:
        logger.error("Ingestion failed: %s", exc)
        asyncio.run(_update_status(rag_document_id, "error", error_message=str(exc)))
        raise self.retry(exc=exc, countdown=30) from exc


@shared_task(bind=True, max_retries=1, soft_time_limit=600, time_limit=720)  # type: ignore
def sync_collection_task(self: Any, sync_log_id: str, source: str, collection_name: str, mode: str, path: str) -> dict[str, Any]:
    """Sync a collection from a local directory."""
    logger.info("Starting sync: %s -> %s (mode=%s)", source, collection_name, mode)
    try:
        return asyncio.run(_run_sync(sync_log_id, source, collection_name, mode, path))
    except Exception as exc:
        logger.error("Sync failed: %s", exc)
        asyncio.run(_update_sync_log(sync_log_id, "error", error_message=str(exc)))
        raise self.retry(exc=exc, countdown=60) from exc
{%- elif cookiecutter.use_taskiq %}


@broker.task
async def ingest_document_task(rag_document_id: str, collection_name: str, filepath: str, source_path: str, replace: bool = False) -> dict[str, Any]:
    """Process a document: parse, chunk, embed, store in vector DB."""
    logger.info("Starting ingestion: %s -> %s", source_path, collection_name)
    try:
        return await _run_ingestion(rag_document_id, collection_name, filepath, source_path, replace)
    except Exception as exc:
        logger.error("Ingestion failed: %s", exc)
        await _update_status(rag_document_id, "error", error_message=str(exc))
        raise


@broker.task
async def sync_collection_task(sync_log_id: str, source: str, collection_name: str, mode: str, path: str) -> dict[str, Any]:
    """Sync a collection from a local directory."""
    logger.info("Starting sync: %s -> %s (mode=%s)", source, collection_name, mode)
    try:
        return await _run_sync(sync_log_id, source, collection_name, mode, path)
    except Exception as exc:
        logger.error("Sync failed: %s", exc)
        await _update_sync_log(sync_log_id, "error", error_message=str(exc))
        raise
{%- elif cookiecutter.use_arq %}


async def ingest_document_task(ctx: dict, rag_document_id: str, collection_name: str, filepath: str, source_path: str, replace: bool = False) -> dict[str, Any]:
    """Process a document: parse, chunk, embed, store in vector DB."""
    logger.info("Starting ingestion: %s -> %s", source_path, collection_name)
    try:
        return await _run_ingestion(rag_document_id, collection_name, filepath, source_path, replace)
    except Exception as exc:
        logger.error("Ingestion failed: %s", exc)
        await _update_status(rag_document_id, "error", error_message=str(exc))
        raise


async def sync_collection_task(ctx: dict, sync_log_id: str, source: str, collection_name: str, mode: str, path: str) -> dict[str, Any]:
    """Sync a collection from a local directory."""
    logger.info("Starting sync: %s -> %s (mode=%s)", source, collection_name, mode)
    try:
        return await _run_sync(sync_log_id, source, collection_name, mode, path)
    except Exception as exc:
        logger.error("Sync failed: %s", exc)
        await _update_sync_log(sync_log_id, "error", error_message=str(exc))
        raise
{%- endif %}



{%- if cookiecutter.use_celery %}


@shared_task(bind=True, max_retries=2, soft_time_limit=600, time_limit=720)  # type: ignore
def sync_single_source_task(self: Any, source_id: str, sync_log_id: str | None = None) -> dict[str, Any]:
    """Sync a single connector source. If sync_log_id provided, use existing log."""
    logger.info("Starting source sync: %s", source_id)
    try:
        return asyncio.run(_run_source_sync(source_id, sync_log_id=sync_log_id))
    except Exception as exc:
        logger.error("Source sync failed: %s", exc)
        raise self.retry(exc=exc, countdown=60) from exc


@shared_task  # type: ignore
def check_scheduled_syncs() -> None:
    """Periodic task: find sources due for sync and dispatch individual tasks."""
    async def _check() -> None:
        async with get_worker_db_context() as db:
            sources = await sync_source_repo.get_due_for_sync(db)
            for source in sources:
                sync_single_source_task.delay(str(source.id))
            logger.info("Scheduled sync check: dispatched %d source(s)", len(sources))

    asyncio.run(_check())
{%- elif cookiecutter.use_taskiq %}


@broker.task
async def sync_single_source_task(source_id: str, sync_log_id: str | None = None) -> dict[str, Any]:
    """Sync a single connector source. If sync_log_id provided, use existing log."""
    logger.info("Starting source sync: %s", source_id)
    return await _run_source_sync(source_id, sync_log_id=sync_log_id)


@broker.task
async def check_scheduled_syncs() -> None:
    """Periodic task: find sources due for sync and dispatch individual tasks."""
    async with get_worker_db_context() as db:
        sources = await sync_source_repo.get_due_for_sync(db)
        for source in sources:
            await sync_single_source_task.kiq(str(source.id))
        logger.info("Scheduled sync check: dispatched %d source(s)", len(sources))
{%- elif cookiecutter.use_arq %}


async def sync_single_source_task(ctx: dict, source_id: str, sync_log_id: str | None = None) -> dict[str, Any]:
    """Sync a single connector source. If sync_log_id provided, use existing log."""
    logger.info("Starting source sync: %s", source_id)
    return await _run_source_sync(source_id, sync_log_id=sync_log_id)


async def check_scheduled_syncs(ctx: dict) -> None:
    """Periodic task: find sources due for sync and dispatch individual tasks."""
    async with get_worker_db_context() as db:
        sources = await sync_source_repo.get_due_for_sync(db)
        pool = ctx["redis"]
        for source in sources:
            await pool.enqueue_job("sync_single_source_task", str(source.id))
        logger.info("Scheduled sync check: dispatched %d source(s)", len(sources))

{%- elif cookiecutter.use_prefect %}


@flow(name="ingest-document", log_prints=True)
async def ingest_document_flow(rag_document_id: str, collection_name: str, filepath: str, source_path: str, replace: bool = False) -> dict[str, Any]:
    """Process a document: parse, chunk, embed, store in vector DB."""
    logger.info("Starting ingestion: %s -> %s", source_path, collection_name)
    try:
        return await _run_ingestion(rag_document_id, collection_name, filepath, source_path, replace)
    except Exception as exc:
        logger.error("Ingestion failed: %s", exc)
        await _update_status(rag_document_id, "error", error_message=str(exc))
        raise


@flow(name="sync-collection", log_prints=True)
async def sync_collection_flow(sync_log_id: str, source: str, collection_name: str, mode: str, path: str) -> dict[str, Any]:
    """Sync a collection from a local directory."""
    logger.info("Starting sync: %s -> %s (mode=%s)", source, collection_name, mode)
    try:
        return await _run_sync(sync_log_id, source, collection_name, mode, path)
    except Exception as exc:
        logger.error("Sync failed: %s", exc)
        await _update_sync_log(sync_log_id, "error", error_message=str(exc))
        raise


@flow(name="sync-single-source", log_prints=True)
async def sync_single_source_flow(source_id: str, sync_log_id: str | None = None) -> dict[str, Any]:
    """Sync a single connector source. If sync_log_id provided, use existing log."""
    logger.info("Starting source sync: %s", source_id)
    return await _run_source_sync(source_id, sync_log_id=sync_log_id)


@flow(name="rag-sync-check", log_prints=True)
async def check_scheduled_syncs_flow() -> None:
    """Scheduled flow: find sources due for sync and dispatch individual flows."""
    import asyncio
    async with get_worker_db_context() as db:
        sources = await sync_source_repo.get_due_for_sync(db)
        tasks = [asyncio.create_task(sync_single_source_flow(str(source.id))) for source in sources]
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        logger.info("Scheduled sync check: dispatched %d source(s)", len(sources))
{%- endif %}




async def _run_ingestion(rag_document_id: str, collection_name: str, filepath: str, source_path: str, replace: bool) -> dict[str, Any]:
    from app.services.rag_document import RAGDocumentService
    ingestion_service = _build_ingestion_service()

    file_path = Path(filepath)
    try:
        result = await ingestion_service.ingest_file(filepath=file_path, collection_name=collection_name, replace=replace, source_path=source_path)
        async with get_worker_db_context() as db:
            await RAGDocumentService(db).complete_ingestion(rag_document_id, vector_document_id=result.document_id)
        await _notify_ws(rag_document_id, "done", source_path)
        logger.info("Ingestion complete: %s", source_path)
        return {"status": "done", "document_id": result.document_id, "filename": source_path}
    except Exception as e:
        await _update_status(rag_document_id, "error", error_message=str(e))
        raise


async def _run_sync(sync_log_id: str, source: str, collection_name: str, mode: str, path: str) -> dict[str, Any]:
    from app.services.rag_document import RAGDocumentService
    from app.services.rag_sync import RAGSyncService
    ingestion_service = _build_ingestion_service()

    target_path = Path(path).resolve()
    if not target_path.exists():
        await _update_sync_log(sync_log_id, "error", error_message=f"Path not found: {path}")
        return {"status": "error", "message": f"Path not found: {path}"}

    if target_path.is_file():
        files = [target_path]
    else:
        files = [f for f in target_path.rglob("*") if f.is_file() and not f.name.startswith(".")]

    allowed = {ext.value for ext in DocumentExtensions}
    files = [f for f in files if f.suffix.lower() in allowed]
    ingested = updated = skipped = failed = 0

    for filepath in files:
        async with get_worker_db_context() as db:
            sync_log_check = await RAGSyncService(db).get_sync_log(sync_log_id)
            if sync_log_check.status == "cancelled":
                logger.info("Sync %s cancelled by user", sync_log_id)
                return {"status": "cancelled", "ingested": ingested, "updated": updated, "skipped": skipped, "failed": failed}

        source_path = str(filepath.resolve())
        if mode in ("new_only", "update_only"):
            existing_id = await ingestion_service.find_existing(collection_name, source_path)

            if mode == "new_only":
                if existing_id:
                    # File exists — check if content changed via hash
                    file_hash = hashlib.sha256(filepath.read_bytes()).hexdigest()
                    existing_hash = await ingestion_service.get_existing_hash(collection_name, source_path)
                    if existing_hash and file_hash == existing_hash:
                        skipped += 1
                        continue
                    # Hash changed — will re-ingest below

            elif mode == "update_only":
                if not existing_id:
                    skipped += 1
                    continue
                file_hash = hashlib.sha256(filepath.read_bytes()).hexdigest()
                existing_hash = await ingestion_service.get_existing_hash(collection_name, source_path)
                if existing_hash and file_hash == existing_hash:
                    skipped += 1
                    continue
        try:
            result = await ingestion_service.ingest_file(filepath=filepath, collection_name=collection_name, replace=True)
            if result.status.value == "done":
                if result.message and "replaced" in result.message:
                    updated += 1
                else:
                    ingested += 1
                async with get_worker_db_context() as db:
                    doc = await RAGDocumentService(db).create_document(
                        collection_name=collection_name,
                        filename=filepath.name,
                        filesize=filepath.stat().st_size,
                        filetype=filepath.suffix.lstrip(".").lower(),
                    )
                    await RAGDocumentService(db).complete_ingestion(
                        str(doc.id), vector_document_id=result.document_id
                    )
            else:
                failed += 1
        except Exception as e:
            logger.warning("Sync file error %s: %s", filepath.name, e)
            failed += 1

    async with get_worker_db_context() as db:
        await RAGSyncService(db).complete_sync(
            sync_log_id,
            status="done" if failed == 0 else "error",
            total_files=len(files),
            ingested=ingested,
            updated=updated,
            skipped=skipped,
            failed=failed,
        )

    return {"status": "done", "ingested": ingested, "updated": updated, "skipped": skipped, "failed": failed}




async def _update_status(rag_document_id: str, status: str, error_message: str | None = None) -> None:
    from app.services.rag_document import RAGDocumentService
    try:
        async with get_worker_db_context() as db:
            doc_svc = RAGDocumentService(db)
            if status == "error":
                await doc_svc.fail_ingestion(rag_document_id, error_message=error_message or "Unknown error")
            elif status == "done":
                # vector_document_id required for complete_ingestion; callers
                # set status="done" directly in _run_ingestion with the ID.
                pass
    except Exception as e:
        logger.warning("Failed to update RAGDocument status: %s", e)


async def _notify_ws(rag_document_id: str, status: str, filename: str) -> None:
    try:
        r = aioredis.from_url(f"redis://{settings.REDIS_HOST}:{settings.REDIS_PORT}/{settings.REDIS_DB}")  # type: ignore[no-untyped-call]
        await r.publish("rag_status", json.dumps({
            "document_id": rag_document_id, "status": status, "filename": filename,
        }))
        await r.aclose()
    except Exception as e:
        logger.warning("Failed to send WS notification: %s", e)


async def _update_sync_log(sync_log_id: str, status: str, error_message: str | None = None) -> None:
    from app.services.rag_sync import RAGSyncService
    try:
        async with get_worker_db_context() as db:
            await RAGSyncService(db).complete_sync(sync_log_id, status=status, error_message=error_message)
    except Exception as e:
        logger.warning("Failed to update SyncLog: %s", e)


async def _run_source_sync(source_id: str, sync_log_id: str | None = None) -> dict[str, Any]:
    """Core sync logic for connector-based sources (shared between all task frameworks).

    Fetches files from a remote connector (e.g. Google Drive, S3), downloads them
    to a temporary directory, and ingests each into the vector store.
    """
    from app.services.rag_sync import RAGSyncService
    async with get_worker_db_context() as db:
        source_svc = SyncSourceService(db)

        source = await source_svc.get_source(source_id)
        connector_cls = CONNECTOR_REGISTRY.get(source.connector_type)
        if not connector_cls:
            await source_svc.update_after_sync(
                source_id, "error", f"Unknown connector: {source.connector_type}"
            )
            return {"status": "error", "message": f"Unknown connector: {source.connector_type}"}

        raw_config = source.config if isinstance(source.config, dict) else json.loads(source.config)
        config = SyncSourceService.decrypt_config_dict(raw_config)
        collection_name = source.collection_name
        sync_mode = source.sync_mode

        # Use existing SyncLog (from API trigger) or create new one (from scheduler)
        if sync_log_id:
            log_id = sync_log_id
        else:
            log = await source_svc.trigger_sync(source_id)
            log_id = str(log.id)

    connector = connector_cls()
    ingestion_svc = _build_ingestion_service()

    ingested = skipped = failed = total = 0

    try:
        files = await connector.list_files(config)
        total = len(files)

        with tempfile.TemporaryDirectory() as tmp_dir:
            for remote_file in files:
                try:
                    local_path = await connector.download_file(remote_file, Path(tmp_dir), config=config)
                    await ingestion_svc.ingest_file(
                        filepath=local_path,
                        collection_name=collection_name,
                        replace=(sync_mode == "full"),
                        source_path=remote_file.source_path,
                    )
                    ingested += 1
                except Exception as e:
                    logger.warning("Failed to sync %s: %s", remote_file.name, e)
                    failed += 1
    except Exception as e:
        logger.error("Source sync failed for %s: %s", source_id, e)
        failed = max(failed, 1)

    async with get_worker_db_context() as db:
        sync_svc = RAGSyncService(db)
        source_svc = SyncSourceService(db)
        try:
            await sync_svc.complete_sync(
                log_id,
                status="done" if not failed else "error",
                total_files=total,
                ingested=ingested,
                skipped=skipped,
                failed=failed,
            )
            await source_svc.update_after_sync(
                source_id,
                status="done" if not failed else "error",
                error=f"{failed} files failed" if failed else None,
            )
        except Exception:
            logger.error("Failed to update sync status for source %s", source_id)

    logger.info(
        "Source sync complete: %s — total=%d, ingested=%d, skipped=%d, failed=%d",
        source_id, total, ingested, skipped, failed,
    )
    return {
        "status": "done" if not failed else "error",
        "total": total,
        "ingested": ingested,
        "skipped": skipped,
        "failed": failed,
    }
{%- endif %}
