"""
Training Data Routes
Endpoints for managing ChromaDB training data (DDL, documentation, question-SQL pairs)
and API Training Data (training table CRUD operations)
"""
import asyncio
import json
import re
import time
import os
from functools import wraps
import pandas as pd
from io import BytesIO
from fastapi import APIRouter, HTTPException, Depends, Body, UploadFile, File as FastAPIFile, Form, Request
from fastapi.responses import StreamingResponse, FileResponse, Response
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any, Tuple
from datetime import datetime
from sqlalchemy.orm import Session
from sqlalchemy.exc import OperationalError
from fastapi.concurrency import run_in_threadpool

from app.services.askme_service import get_askme_service
from app.core.database import get_db, SessionLocal
from app.models.agent import Training
from app.services.agent_config import get_agent_config, agent_has_file_datasource
from app.core.config import settings, get_base_url
from app.utils.folder_scan_scheduler import (
    find_scheduled_job_for_file_path,
    get_scheduled_job_status,
    is_file_path_scheduler_running,
    is_folder_scan_stop_pending,
    release_folder_scan_job_after_stop,
    schedule_folder_scan_job,
    signal_immediate_folder_scan_stop,
    try_resume_background_scheduler,
)
from app.services.file_schedule_service import FileScheduleService
from app.services.folder_scan_job_prompts import (
    FOLDER_SCAN_FILE_PATH_ALREADY_EXISTS,
    FOLDER_SCAN_JOB_STOP_ALREADY_CANCELLED,
    FOLDER_SCAN_JOB_STOP_ALREADY_COMPLETED,
    FOLDER_SCAN_JOB_STOP_ALREADY_FAILED,
    FOLDER_SCAN_JOB_STOP_NOT_RUNNING,
    FOLDER_SCAN_JOB_STOPPED,
    FOLDER_SCAN_JOB_STOP_PENDING,
    FOLDER_SCAN_SCHEDULER_REST_UPLOAD_BLOCKED,
)
from app.services.document_service import FILE_SYS_METADATA_DB
from app.services.document_db_service import DocumentDbService
from app.utils.file_schedule_utils import (
    active_file_path_schedule,
    build_folder_scan_step_user_message,
    dedupe_schedule_list_by_file_path,
    effective_schedule_status,
    folder_scan_schedule_seconds,
    get_folder_scan_job_status,
    global_scheduler_rest_phase,
    schedule_details_payload,
    scheduler_message_from_schedule_row,
    scheduler_status_message,
)

_file_schedule_svc = FileScheduleService()
from app.core.dependencies import get_current_user
from app.utils.document_preview import (
    DEFAULT_PREVIEW_PDF_CACHE_DIR,
    apply_preview_contract_to_mapping,
    build_document_preview_url,
    document_media_type,
    extract_and_cache_page_pdf,
    render_page_png_bytes,
    resolve_preview_file_response,
)

# Import DocumentService for file operations (new modular structure)
# Supports both SQLAlchemy Session and db_file modes for backward compatibility
try:
    from app.services.document_service import (
        DocumentService,
        FILE_SYS_METADATA_DB,
        FILE_VECTORS_DB_DIR,
        FILE_UPLOAD_DISK_ROOT,
        FILE_IMAGE_DISK_ROOT,
        FILE_EXTRACT_DISK_ROOT,
    )
    # Alias for backward compatibility with existing code
    DocumentAssistant = DocumentService
    DOCUMENT_ASSISTANT_AVAILABLE = True
except ImportError:
    DOCUMENT_ASSISTANT_AVAILABLE = False
    DocumentAssistant = None
    DocumentService = None


# Helper function to convert image paths to URLs
def convert_image_path_to_url(image_path: str, agent_id: str, base_url: str = None) -> str:
    """
    Convert a local file path to a backend URL for serving images.
    
    Args:
        image_path: Local file path (e.g., "extracted_images_/file_data/.../image.png")
        agent_id: Agent ID (kept for signature compatibility)
        base_url: Base URL of the backend server
    
    Returns:
        URL to access the image through the backend API
    """
    if not image_path:
        return ""
    
    # Check if already a URL (prevent double conversion)
    if image_path.startswith("http://") or image_path.startswith("https://"):
        return image_path
    
    # Use provided base_url or default from initial_config.txt (BASE_URL + SCHEME)
    if base_url is None:
        base_url = get_base_url()
    
    # Normalize path separators to forward slashes
    image_path = image_path.replace("\\", "/")
    
    # Construct the URL - simplified for the new API endpoint structure
    # Served by static_files router: /api/images/{file_path} (disk path unchanged)
    url = f"{base_url}/api/images/{image_path}"
    return url


# Schemas
class CreateTrainingDataRequest(BaseModel):
    """Create training data - question is required, others optional"""
    agent_id: str
    question: Optional[str] = None
    sql: Optional[str] = None
    ddl: Optional[str] = None
    documentation: Optional[str] = None


class RemoveRequest(BaseModel):
    agent_id: str
    id: str


class GetAllDBTrainingDataRequest(BaseModel):
    user_id: str
    agent_id: str


# ===== API Training Data Schemas (for training table) ===== #
class CreateAPITrainingRequest(BaseModel):
    agent_id: str
    name: str
    endpoint: str
    base_params: str
    response_schema: str
    description: str
    is_sap: bool = False  # Maps to extra1: "Y" or "N"


class GetAPITrainingRequest(BaseModel):
    user_id: str
    agent_id: str


class UpdateAPITrainingRequest(BaseModel):
    id: int
    agent_id: str
    name: str
    endpoint: str
    base_params: str
    response_schema: str
    description: str
    is_sap: bool = False  # Maps to extra1: "Y" or "N"


class DeleteAPITrainingRequest(BaseModel):
    id: int
    agent_id: str


# ===== Document Training Center Schemas ===== #
class GetFilesFoldersRequest(BaseModel):
    user_id: str
    agent_id: str


class GetFileExtractedImagesRequest(BaseModel):
    user_id: str
    agent_id: str
    file_id: str


class UploadFolderRequest(BaseModel):
    user_id: str
    agent_id: str
    folder_path: str


class UploadFileRequest(BaseModel):
    user_id: str
    agent_id: str
    folder_path: str = ""  # Optional, defaults to "Root" if empty


class CreateFolderRequest(BaseModel):
    user_id: str
    agent_id: str
    folder_name: str  # Name of the folder to create


class DeleteFolderRequest(BaseModel):
    user_id: str
    agent_id: str
    folder_path: str  # Path of the folder to delete


class SearchFilesRequest(BaseModel):
    user_id: str
    agent_id: str
    keyword: str  # Search term to match against filename, folder, or content preview
    folder_path: Optional[str] = None  # Optional: restrict search to a specific folder # Cap result size to bound DB work


class ScheduleFolderRequest(BaseModel):
    agent_id: str
    file_path: str
    free_time: Optional[int] = None


class StopScheduleFolderRequest(BaseModel):
    """JSON body for POST /document/schedule-stop.

    - Provide ``job_id`` to cancel a **single** specific schedule.
    - Omit ``job_id`` (or pass null) to cancel **all** active schedules for the agent.
    """

    agent_id: str = Field(..., min_length=1, description="Agent ID")
    job_id: Optional[str] = Field(None, description="Folder scan job ID from schedule-folder. Omit to stop all active schedules for the agent.")
    file_path: Optional[str] = Field(None, description="File/Folder path (optional)")

    model_config = {
        "json_schema_extra": {
            "example": {
                "agent_id": "1",
                "job_id": "0bf9f19d-23ac-4cd7-8ed3-12e8b42e5136",
                "file_path": "/path/to/folder"
            }
        }
    }


class ScheduleStopStatusRequest(BaseModel):
    """Poll whether a schedule stop has fully completed (for UI spinners)."""

    agent_id: str = Field(..., min_length=1)
    job_id: str = Field(..., min_length=1)


class ListSchedulesRequest(BaseModel):
    user_id: str
    agent_id: str
    file_path: Optional[str] = None


router = APIRouter(prefix="/training", tags=["Training Data (ChromaDB & API Training)"])

# Large folder ingest can run for many files; allow longer than default proxy/client limits.
FOLDER_UPLOAD_REQUEST_TIMEOUT_SECONDS = 30 * 60


def folder_upload_request_timeout(func):
    """Apply a 30-minute wall-clock limit to folder upload handlers only."""

    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any):
        try:
            return await asyncio.wait_for(
                func(*args, **kwargs),
                timeout=FOLDER_UPLOAD_REQUEST_TIMEOUT_SECONDS,
            )
        except asyncio.TimeoutError:
            minutes = FOLDER_UPLOAD_REQUEST_TIMEOUT_SECONDS // 60
            return {
                "success": False,
                "message": f"Request timed out after {minutes} minutes.",
                "details": {},
            }

    return wrapper


@router.post("/create-db-training-data", status_code=201)
async def create_db_training_data(
    request: CreateTrainingDataRequest,
    current_user: Any = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Create Database training data.
    Question is required, sql, ddl, and documentation are optional.
    """
    try:
        service = get_askme_service(agent_id=request.agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{request.agent_id}' not found",
                "details": {}
            }
        
        doc_ids = []
        
        # Add SQL if provided
        if request.sql:
            await run_in_threadpool(service.add_traning_data, type="sql", query=request.sql, question=request.question)
            doc_ids.append({"type": "sql"})
        
        # Add DDL if provided
        if request.ddl:
            await run_in_threadpool(service.add_traning_data, type="ddl", query=request.ddl, question=None)
            doc_ids.append({"type": "ddl"})
        
        # Add documentation if provided
        if request.documentation:
            await run_in_threadpool(service.add_traning_data, type="documentation", query=request.documentation, question=None)
            doc_ids.append({"type": "documentation"})
        
        if not doc_ids:
            return {
                "success": False,
                "message": "At least one of sql, ddl, or documentation must be provided",
                "details": {}
            }
        
        return {
            "success": True,
            "message": "Training data added successfully",
            "details": {
                "question": request.question,
                "created": doc_ids
            }
        }   
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error creating database training data: {str(e)}",
            "details": {"question": request.question}
        }


@router.post("/get-all-db-training-data")
async def get_all_db_training_data(
    request: GetAllDBTrainingDataRequest,
    current_user: Any = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """
    Get all Database training data (SQL, DDL, Documentation).
    Returns all training data organized by type.
    """
    try:
        service = get_askme_service(agent_id=request.agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{request.agent_id}' not found",
                "details": {}
            }
        
        # Get training data as DataFrame
        training_df = await run_in_threadpool(service.get_traning_data)
        
        if training_df is None or training_df.empty:
            return {
                "success": True,
                "message": "No training data found",
                "details": {
                    "sql": [],
                    "ddl": [],
                    "documentation": []
                }
            }
            
        # Convert DataFrame to list of dicts
        dataframe = training_df.to_dict(orient="records")
        
        # Organize data by type
        sql_data = []
        ddl_data = []
        documentation_data = []
        
        for item_data in dataframe:
            item_id = item_data.get("id", "")
            question = item_data.get("question", "")
            content = item_data.get("content", "")
            training_data_type = item_data.get("training_data_type", "")
            
            if training_data_type == "sql":
                sql_data.append({
                    "id": item_id,
                    "question": question,
                    "sql": content
                })
            elif training_data_type == "ddl":
                ddl_data.append({
                    "id": item_id,
                    "ddl": content
                })
            elif training_data_type == "documentation":
                documentation_data.append({
                    "id": item_id,
                    "documentation": content
                })
        
        return {
            "success": True,
            "message": "Training data retrieved successfully",
            "details": {
                "sql": sql_data,
                "ddl": ddl_data,
                "documentation": documentation_data
            }
        }   
    except Exception as e:
        return {
            "success": False,
            "message": f"Error getting all database training data: {str(e)}",
            "details": {}
        }


@router.post("/upload-db-training-data")
async def upload_db_training_data(
    file: UploadFile = FastAPIFile(...),
    agent_id: str = Form(...)
):
    """
    Upload Excel file to bulk import database training data
    Excel format: columns must be ["question", "content", "training_data_type"]
    training_data_type must be "sql", "ddl", or "documentation"
    Similar to add_bulk_data_dialog in trainer_dash.py
    """
    try:
        service = get_askme_service(agent_id=agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{agent_id}' not found",
                "details": {}
            }
        
        # Validate file type
        if not file.filename.endswith(('.xlsx', '.xls')):
            return {
                "success": False,
                "message": "Invalid file type. Please upload an Excel file (.xlsx or .xls)",
                "details": {}
            }
        
        # Read Excel file
        try:
            file_content = await file.read()
            df = pd.read_excel(BytesIO(file_content), engine='openpyxl')
        except Exception as e:
            return {
                "success": False,
                "message": f"Failed to read Excel file: {str(e)}",
                "details": {}
            }
        
        # Normalize columns (strip whitespace, handle non-string headers)
        df.columns = df.columns.astype(str).str.strip()

        # Validate that only allowed columns are present
        allowed_columns = {"question", "content", "training_data_type"}
        extra_columns = [col for col in df.columns if col not in allowed_columns]

        if extra_columns:
            return {
                "success": False,
                "message": f"Excel format error: Invalid columns found: {extra_columns}. Only 'question', 'content', and 'training_data_type' are allowed.",
                "details": {"columns_found": list(df.columns)}
            }

        # Check if question column is missing and add it with default text
        if "question" not in df.columns:
            df["question"] = "None"

        # Validate required columns (excluding question since we just added it if missing)
        required_cols = ["content", "training_data_type"]
        missing_cols = [col for col in required_cols if col not in df.columns]
        
        if missing_cols:
            return {
                "success": False,
                "message": f"Excel format error: Missing required columns: {missing_cols}",
                "details": {"columns_found": list(df.columns)}
            }
        
        
        # Handle empty questions: fill NaNs with default text
        df["question"] = df["question"].fillna("None")
        
        # Check for empty fields (skip question column as it is filled)
        temp_df = df.drop(columns=["question"])
        if temp_df.isnull().values.any().any():
            return {
                "success": False,
                "message": "Excel contains empty fields in content or training_data_type.",
                "details": {}
            }
        
        # Validate training_data_type values
        valid_types = ["sql", "ddl", "documentation"]
        if not df["training_data_type"].isin(valid_types).all():
            invalid_types = df[~df["training_data_type"].isin(valid_types)]["training_data_type"].unique().tolist()
            return {
                "success": False,
                "message": f"training_data_type must be one of {valid_types}. Found invalid types: {invalid_types}",
                "details": {}
            }
        
        # Process and add training data
        success_count = 0
        failed_count = 0
        errors = []
        
        for idx, row in df.iterrows():
            try:
                training_data_type = str(row["training_data_type"]).lower()
                question = str(row["question"]) if pd.notna(row["question"]) else ""
                content = str(row["content"])
                
                # Ensure question is set if empty (double check)
                if not question or question.strip() == "":
                    question = "None"
                
                # Add training data using askme_service
                await run_in_threadpool(
                    service.add_traning_data,
                    type=training_data_type,
                    # Pass the processed question for all types, or None if logic demands, 
                    # but user requested filling the column.
                    question=question, 
                    query=content
                )
                
                success_count += 1
            except Exception as e:
                failed_count += 1
                errors.append(f"Row {idx + 2}: {str(e)}")
        
        return {
            "success": True,
            "message": f"Bulk upload completed. {success_count} records added, {failed_count} failed.",
            "details": {
                "success_count": success_count,
                "failed_count": failed_count,
                "total_rows": len(df),
                "errors": errors if errors else None
            }
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Error uploading file: {str(e)}",
            "details": {}
        }


@router.post("/download-db-training-data")
async def download_db_training_data(request: GetAllDBTrainingDataRequest):
    """
    Download all database training data as Excel file
    Similar to load_data in trainer_dash.py
    """
    try:
        service = get_askme_service(agent_id=request.agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{request.agent_id}' not found",
                "details": {}
            }
        
        # Get training data as DataFrame (similar to trainer_dash.py line 26)
        training_df = await run_in_threadpool(service.get_traning_data)
        
        if training_df is None or training_df.empty:
            return {
                "success": False,
                "message": "No training data found to download",
                "details": {}
            }
        
        # Convert DataFrame to dict records (similar to trainer_dash.py line 26)
        dataframe = training_df.to_dict(orient="records")
        
        # Prepare data for Excel - format: ["id", "question", "content", "training_data_type"]
        excel_data = []
        for data in dataframe:
            excel_data.append({
                "id": data.get("id", ""),
                "question": data.get("question", ""),
                "content": data.get("content", ""),
                "training_data_type": data.get("training_data_type", "")
            })
        
        # Create DataFrame with exact column order (similar to trainer_dash.py line 124-135)
        df = pd.DataFrame(excel_data, columns=["id", "question", "content", "training_data_type"])
        
        # Convert to Excel bytes - matching format from trainer_dash.py line 132-136
        output = BytesIO()
        def write_excel():
            with pd.ExcelWriter(output, engine='openpyxl') as writer:
                df.to_excel(writer, index=False, sheet_name='Traning Data')
            output.seek(0)
            return output
        
        output = await run_in_threadpool(write_excel)
        
        # Return as streaming response
        return StreamingResponse(
            BytesIO(output.read()),
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers={"Content-Disposition": "attachment; filename=traning_data.xlsx"}
        )
    except Exception as e:
        return {
            "success": False,
            "message": f"Error downloading training data: {str(e)}",
            "details": {}
        }


@router.post("/delete-db-training-data")
async def delete_db_training_data(request: RemoveRequest):
    """
    Delete Database training data by ID
    Similar to trainer_dash.py line 103
    """
    try:
        service = get_askme_service(agent_id=request.agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{request.agent_id}' not found",
                "details": {}
            }
        
        # Delete training data using askme_service (similar to trainer_dash.py line 103)
        result = await run_in_threadpool(service.delete_traning_data, id=request.id)
        
        return {
            "success": True,
            "message": "Database training data deleted successfully",
            "details": {"id": request.id}
        }
    except Exception as e:
        return {
            "success": False,
            "message": str(e),
            "details": {}
        }


# ===== API Training Data Endpoints (training table CRUD) ===== #

@router.post("/api-training/create", status_code=201)
async def create_api_training_data(request: CreateAPITrainingRequest, db: Session = Depends(get_db)):
    """
    Create API training data entry.
    Uses retry logic for database locks and ensures database connection is closed after response.
    Similar to create_training_data in training_center.py
    """
    max_retries = 5
    retry_delay = 0.1  # Start with 100ms delay
    
    # Use the provided session for the first attempt
    current_db = db
    manually_created_session = False
    
    try:
        for attempt in range(max_retries):
            try:
                # Create a new session for retries (after first attempt)
                if attempt > 0:
                    current_db = SessionLocal()
                    manually_created_session = True
                
                try:
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    
                    
                    # Check if name already exists (unique constraint)
                    def check_existing():
                        return current_db.query(Training).filter(Training.name == request.name).first()

                    existing = await run_in_threadpool(check_existing)
                    if existing:
                        if manually_created_session:
                            await run_in_threadpool(current_db.close)
                        return {
                            "success": False,
                            "message": f"Training data with name '{request.name}' already exists",
                            "details": {}
                        }
                    
                    # Create new training data entry
                    training_data = Training(
                        agent_id=request.agent_id,
                        name=request.name,
                        endpoint=request.endpoint,
                        base_params=request.base_params,
                        response_schema=request.response_schema,
                        description=request.description,
                        extra1="Y" if request.is_sap else "N",
                        extra2="",
                        created_at=current_time,
                        updated_at=current_time
                    )
                    
                    await run_in_threadpool(current_db.add, training_data)
                    await run_in_threadpool(current_db.commit)
                    await run_in_threadpool(current_db.refresh, training_data)
                    
                    result = {
                        "success": True,
                        "message": "API training data created successfully",
                        "details": {
                            "id": training_data.id,
                            "agent_id": training_data.agent_id,
                            "name": training_data.name,
                            "endpoint": training_data.endpoint,
                            "base_params": training_data.base_params,
                            "response_schema": training_data.response_schema,
                            "description": training_data.description,
                            "extra1": training_data.extra1,
                            "extra2": training_data.extra2,
                            "created_at": training_data.created_at
                        }
                    }
                    
                    # Close manually created session before returning
                    if manually_created_session:
                        current_db.close()
                    
                    return result
                    
                except OperationalError as e:
                    current_db.rollback()
                    # Check if it's a database locked error
                    if "database is locked" in str(e).lower() and attempt < max_retries - 1:
                        # Close the current session before retry
                        if manually_created_session:
                            current_db.close()
                            manually_created_session = False
                        # Exponential backoff: wait longer for each retry
                        time.sleep(retry_delay * (2 ** attempt))
                        continue
                    else:
                        if manually_created_session:
                            current_db.close()
                        return {
                            "success": False,
                            "message": f"Error creating API training data: Database is locked. Please try again later. {str(e)}",
                            "details": {}
                        }
                except HTTPException:
                    if manually_created_session:
                        current_db.close()
                    raise
                except Exception as e:
                    current_db.rollback()
                    if manually_created_session:
                        current_db.close()
                    return {
                        "success": False,
                        "message": f"Error creating API training data: {str(e)}",
                        "details": {}
                    }
            except HTTPException:
                raise
            except Exception as e:
                if manually_created_session and current_db:
                    current_db.close()
                return {
                    "success": False,
                    "message": f"Error creating API training data: {str(e)}",
                    "details": {}
                }
        
        # If we've exhausted all retries
        return {
            "success": False,
            "message": "Error creating API training data: Database is locked after multiple retries. Please try again later.",
            "details": {}
        }
    finally:
        # Ensure any manually created session is closed
        if manually_created_session and current_db:
            try:
                current_db.close()
            except Exception:
                pass


@router.post("/api-training/all")
async def get_all_api_training_data(request: GetAPITrainingRequest, db: Session = Depends(get_db)):
    """
    Get all API training data for a specific agent
    Similar to get_all_training_data in training_center.py
    """
    try:
        def get_all_apis():
            return db.query(Training).filter(
                Training.agent_id == request.agent_id
            ).all()

        training_data_list = await run_in_threadpool(get_all_apis)
        
        result = []
        for item in training_data_list:
            result.append({
                "id": item.id,
                "name": item.name,
                "endpoint": item.endpoint,
                "base_params": item.base_params,
                "response_schema": item.response_schema,
                "description": item.description,
                "extra1": item.extra1,
                "created_at": item.created_at,
                "updated_at": item.updated_at
            })
        
        return {
            "success": True,
            "message": "API training data retrieved successfully",
            "details": result
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Error fetching API training data: {str(e)}",
            "details": {}
        }


@router.post("/api-training/update")
async def update_api_training_data(request: UpdateAPITrainingRequest, db: Session = Depends(get_db)):
    """
    Update API training data entry.
    Uses retry logic for database locks and ensures database connection is closed after response.
    Similar to update_training_data in training_center.py
    """
    max_retries = 5
    retry_delay = 0.1  # Start with 100ms delay
    
    # Use the provided session for the first attempt
    current_db = db
    manually_created_session = False
    
    try:
        for attempt in range(max_retries):
            try:
                # Create a new session for retries (after first attempt)
                if attempt > 0:
                    current_db = SessionLocal()
                    manually_created_session = True
                
                try:
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    
                    
                    # Find the training data entry
                    def get_training_data():
                        return current_db.query(Training).filter(Training.id == request.id).first()
                    
                    training_data = await run_in_threadpool(get_training_data)
                    if not training_data:
                        if manually_created_session:
                            await run_in_threadpool(current_db.close)
                        return {
                            "success": False,
                            "message": "API training data not found",
                            "details": {}
                        }
                    
                    # Check if name already exists for a different record (unique constraint)
                    def check_existing_update():
                        return current_db.query(Training).filter(
                            Training.name == request.name,
                            Training.id != request.id
                        ).first()

                    existing = await run_in_threadpool(check_existing_update)
                    if existing:
                        if manually_created_session:
                            await run_in_threadpool(current_db.close)
                        return {
                            "success": False,
                            "message": f"Training data with name '{request.name}' already exists",
                            "details": {}
                        }
                    
                    # Update fields
                    training_data.name = request.name
                    training_data.endpoint = request.endpoint
                    training_data.base_params = request.base_params
                    training_data.response_schema = request.response_schema
                    training_data.description = request.description
                    training_data.extra1 = "Y" if request.is_sap else "N"
                    training_data.updated_at = current_time
                    
                    await run_in_threadpool(current_db.commit)
                    await run_in_threadpool(current_db.refresh, training_data)
                    
                    result = {
                        "success": True,
                        "message": "API training data updated successfully",
                        "details": {
                            "id": training_data.id,
                            "name": training_data.name,
                            "endpoint": training_data.endpoint,
                            "base_params": training_data.base_params,
                            "response_schema": training_data.response_schema,
                            "description": training_data.description,
                            "extra1": training_data.extra1,
                            "extra2": training_data.extra2,
                            "updated_at": training_data.updated_at
                        }
                    }
                    
                    # Close manually created session before returning
                    if manually_created_session:
                        current_db.close()
                    
                    return result
                    
                except OperationalError as e:
                    current_db.rollback()
                    # Check if it's a database locked error
                    if "database is locked" in str(e).lower() and attempt < max_retries - 1:
                        # Close the current session before retry
                        if manually_created_session:
                            current_db.close()
                            manually_created_session = False
                        # Exponential backoff: wait longer for each retry
                        time.sleep(retry_delay * (2 ** attempt))
                        continue
                    else:
                        if manually_created_session:
                            current_db.close()
                        return {
                            "success": False,
                            "message": f"Error updating API training data: Database is locked. Please try again later. {str(e)}",
                            "details": {}
                        }
                except HTTPException:
                    if manually_created_session:
                        current_db.close()
                    raise
                except Exception as e:
                    current_db.rollback()
                    if manually_created_session:
                        current_db.close()
                    return {
                        "success": False,
                        "message": f"Error updating API training data: {str(e)}",
                        "details": {}
                    }
            except HTTPException:
                raise
            except Exception as e:
                if manually_created_session and current_db:
                    current_db.close()
                return {
                    "success": False,
                    "message": f"Error updating API training data: {str(e)}",
                    "details": {}
                }
        
        # If we've exhausted all retries
        return {
            "success": False,
            "message": "Error updating API training data: Database is locked after multiple retries. Please try again later.",
            "details": {}
        }
    finally:
        # Ensure any manually created session is closed
        if manually_created_session and current_db:
            try:
                current_db.close()
            except Exception:
                pass


@router.post("/api-training/delete")
async def delete_api_training_data(request: DeleteAPITrainingRequest, db: Session = Depends(get_db)):
    """
    Delete API training data entry.
    Uses retry logic for database locks and ensures database connection is closed after response.
    Similar to delete_training_data in training_center.py
    """
    max_retries = 5
    retry_delay = 0.1  # Start with 100ms delay
    
    # Use the provided session for the first attempt
    current_db = db
    manually_created_session = False
    
    try:
        for attempt in range(max_retries):
            try:
                # Create a new session for retries (after first attempt)
                if attempt > 0:
                    current_db = SessionLocal()
                    manually_created_session = True
                
                try:
                    training_data = current_db.query(Training).filter(Training.id == request.id).first()
                    if not training_data:
                        if manually_created_session:
                            current_db.close()
                        return {
                            "success": False,
                            "message": "API training data not found",
                            "details": {}
                        }
                    
                    current_db.delete(training_data)
                    current_db.commit()
                    
                    result = {
                        "success": True,
                        "message": "API training data deleted successfully",
                        "details": {"id": request.id}
                    }
                    
                    # Close manually created session before returning
                    if manually_created_session:
                        current_db.close()
                    
                    return result
                    
                except OperationalError as e:
                    current_db.rollback()
                    # Check if it's a database locked error
                    if "database is locked" in str(e).lower() and attempt < max_retries - 1:
                        # Close the current session before retry
                        if manually_created_session:
                            current_db.close()
                            manually_created_session = False
                        # Exponential backoff: wait longer for each retry
                        time.sleep(retry_delay * (2 ** attempt))
                        continue
                    else:
                        if manually_created_session:
                            current_db.close()
                        return {
                            "success": False,
                            "message": f"Error deleting API training data: Database is locked. Please try again later. {str(e)}",
                            "details": {}
                        }
                except HTTPException:
                    if manually_created_session:
                        current_db.close()
                    raise
                except Exception as e:
                    current_db.rollback()
                    if manually_created_session:
                        current_db.close()
                    return {
                        "success": False,
                        "message": f"Error deleting API training data: {str(e)}",
                        "details": {}
                    }
            except HTTPException:
                raise
            except Exception as e:
                if manually_created_session and current_db:
                    current_db.close()
                return {
                    "success": False,
                    "message": f"Error deleting API training data: {str(e)}",
                    "details": {}
                }
        
        # If we've exhausted all retries
        return {
            "success": False,
            "message": "Error deleting API training data: Database is locked after multiple retries. Please try again later.",
            "details": {}
        }
    finally:
        # Ensure any manually created session is closed
        if manually_created_session and current_db:
            try:
                current_db.close()
            except Exception:
                pass



@router.post("/upload-api-training_center", status_code = 201)
async def upload_api_training_center(
    file: UploadFile = FastAPIFile(...),
    agent_id: str = Form(...),
    db: Session = Depends(get_db)
):
    """
    Upload Excel file to bulk import API training data.
    Required columns: API Name, API Endpoint, Description, Base Params, Response Schema
    """
    try:
        service = get_askme_service(agent_id=agent_id)
        if not service:
            return {
                "success": False,
                "message": f"Agent '{agent_id}' not found",
                "details": {}
            }
            
        # Validate file type
        if not file.filename.endswith(('.xlsx', '.xls')):
            return {
                "success": False,
                "message": "Invalid file type. Please upload an Excel file (.xlsx or .xls)",
                "details": {}
            }
        
        # Read Excel file
        try:
            file_content = await file.read()
            df = pd.read_excel(BytesIO(file_content), engine='openpyxl')
        except Exception as e:
            return {
                "success": False,
                "message": f"Failed to read Excel file: {str(e)}",
                "details": {}
            }
            
        # Strip whitespace from column names
        df.columns = df.columns.astype(str).str.strip()
        
        # Map Excel columns to DB columns
        column_mapping = {
            "API Name": "name",
            "API Endpoint": "endpoint",
            "Description": "description",
            "Base Params": "base_params",
            "Response Schema": "response_schema",
            "Is_SAP": "extra1"
        }
        
        required_columns = list(column_mapping.keys())
        missing_columns = [col for col in required_columns if col not in df.columns]
        
        if missing_columns:
            return {
                "success": False,
                "message": f"Excel format error: Missing required columns: {missing_columns}",
                "details": {"columns_found": list(df.columns)}
            }

        # Check for extra columns
        extra_columns = [col for col in df.columns if col not in required_columns]
        if extra_columns:
            return {
                "success": False,
                "message": f"Excel format error: Invalid columns found: {extra_columns}. Only 'API Name', 'API Endpoint', 'Description', 'Base Params', 'Response Schema', and 'Is_SAP' are allowed.",
                "details": {"columns_found": list(df.columns)}
            }
            
        # DB Operation with retry logic
        max_retries = 5
        retry_delay = 0.1  # Start with 100ms delay
        
        # Use the provided session for the first attempt
        current_db = db
        manually_created_session = False
        
        try:
            for attempt in range(max_retries):
                try:
                    # Create a new session for retries (after first attempt)
                    if attempt > 0:
                        current_db = await run_in_threadpool(SessionLocal)
                        manually_created_session = True
                    
                    success_count = 0
                    failed_count = 0
                    errors = []
                    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    
                    for idx, row in df.iterrows():
                        try:
                            # row access is fast, no need to async
                            name_val = row["API Name"]
                            name = str(name_val).strip() if pd.notna(name_val) else ""
                            
                            if not name:
                               failed_count += 1
                               errors.append(f"Row {idx + 2}: 'API Name' is required.")
                               continue
            
                            # Check for duplicate name
                            def check_existing_name(n):
                                return current_db.query(Training).filter(Training.name == n).first()
                            
                            existing = await run_in_threadpool(check_existing_name, name)
                            if existing:
                                failed_count += 1
                                errors.append(f"Row {idx + 2}: API Name '{name}' already exists.")
                                continue
                            
                            is_sap_raw = str(row["Is_SAP"]).strip().upper() if pd.notna(row["Is_SAP"]) else "N"
                            # Map to Y/N, default to N if not Y
                            is_sap = "Y" if is_sap_raw == "Y" else "N"

                            training_data = Training(
                                agent_id=agent_id,
                                name=name,
                                endpoint=str(row["API Endpoint"]) if pd.notna(row["API Endpoint"]) else "",
                                base_params=str(row["Base Params"]) if pd.notna(row["Base Params"]) else "",
                                response_schema=str(row["Response Schema"]) if pd.notna(row["Response Schema"]) else "",
                                description=str(row["Description"]) if pd.notna(row["Description"]) else "",
                                extra1=is_sap, # From Is_SAP column
                                extra2="",
                                created_at=current_time,
                                updated_at=current_time
                            )
                            
                            await run_in_threadpool(current_db.add, training_data)
                            await run_in_threadpool(current_db.commit)
                            success_count += 1
                            
                        except Exception as e:
                            await run_in_threadpool(current_db.rollback)
                            # If it's a lock error, re-raise to trigger outer retry
                            if "database is locked" in str(e).lower():
                                raise OperationalError(str(e), params=None, orig=e)
                                
                            failed_count += 1
                            errors.append(f"Row {idx + 2}: {str(e)}")
                    
                    result = {
                        "success": True,
                        "message": f"Bulk upload completed. {success_count} records added, {failed_count} failed.",
                        "details": {
                            "success_count": success_count,
                            "failed_count": failed_count,
                            "total_rows": len(df),
                            "errors": errors if errors else None
                        }
                    }
                    
                    if manually_created_session:
                        await run_in_threadpool(current_db.close)
                        
                    return result
                    
                except OperationalError as e:
                    await run_in_threadpool(current_db.rollback)
                    # Check if it's a database locked error
                    if "database is locked" in str(e).lower() and attempt < max_retries - 1:
                        # Close the current session before retry
                        if manually_created_session:
                            await run_in_threadpool(current_db.close)
                            manually_created_session = False
                        # Exponential backoff: wait longer for each retry
                        time.sleep(retry_delay * (2 ** attempt))
                        continue
                    else:
                        if manually_created_session:
                            current_db.close()
                        return {
                            "success": False,
                            "message": f"Error uploading API training data: Database is locked. Please try again later. {str(e)}",
                            "details": {}
                        }
                except HTTPException:
                    if manually_created_session:
                        current_db.close()
                    raise
                except Exception as e:
                    current_db.rollback()
                    if manually_created_session:
                        current_db.close()
                    return {
                        "success": False,
                        "message": f"Error uploading API training data: {str(e)}",
                        "details": {}
                    }
                    
            # If we've exhausted all retries
            return {
                "success": False,
                "message": "Error uploading API training data: Database is locked after multiple retries. Please try again later.",
                "details": {}
            }
            
        finally:
             if manually_created_session and current_db:
                try:
                    await run_in_threadpool(current_db.close)
                except Exception:
                    pass
            
    except Exception as e:
        return {
            "success": False,
            "message": f"Error uploading file: {str(e)}",
            "details": {}
        }
    finally:
        # Close the injected dependency session if strictly needed, 
        # though the retry logic manages 'current_db'.
        # But 'db' from Depends is still open.
        try:
            db.close()
        except:
             pass


@router.post("/download-api-training-center")
async def download_api_training_center(request: GetAPITrainingRequest, db: Session = Depends(get_db)):
    """
    Download API training data as Excel file.
    Download API training data as Excel file.
    Columns: id, API Name, API Endpoint, Base Params, Response Schema, Description, Is_SAP
    """
    try:
        # Fetch all training data for this agent
        def get_training_data():
            return db.query(Training).filter(
                Training.agent_id == request.agent_id
            ).all()

        training_data_list = await run_in_threadpool(get_training_data)
        
        if not training_data_list:
             return {
                "success": False,
                "message": "No API training data found to download",
                "details": {}
            }
            
        data = []
        for item in training_data_list:
            data.append({
                "id": item.id,
                "API Name": item.name,
                "API Endpoint": item.endpoint,
                "Base Params": item.base_params,
                "Response Schema": item.response_schema,
                "Description": item.description,
                "Is_SAP": item.extra1
            })
            
        df = pd.DataFrame(data)
        
        # Ensure column order
        columns_order = ["id", "API Name", "API Endpoint", "Base Params", "Response Schema", "Description", "Is_SAP"]
        # Handle case where some columns might be missing if data list is empty (handled above)
        df = df[columns_order]
        
        output = BytesIO()
        with pd.ExcelWriter(output, engine='openpyxl') as writer:
            df.to_excel(writer, index=False, sheet_name='API Training Data')
        output.seek(0)
        
        return StreamingResponse(
            BytesIO(output.read()),
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            headers={"Content-Disposition": "attachment; filename=api_training_data.xlsx"}
        )
        
    except Exception as e:
        return {
            "success": False,
            "message": f"Error downloading API training data: {str(e)}",
            "details": {}
        }


# ===== Document Training Center Endpoints ===== #


def _open_scan_db() -> DocumentDbService:
    return DocumentDbService(db_file=FILE_SYS_METADATA_DB)


def _lookup_file_schedule_for_stop(
    *,
    agent_id: str,
    job_id: Optional[str] = None,
    file_path: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """Resolve a file_schedules row by job_id or file_path for schedule-stop."""
    jid = str(job_id or "").strip()
    if jid:
        row = _file_schedule_svc.get_file_schedule_by_job_id(jid)
        if row:
            return row
    fp = str(file_path or "").strip()
    if fp:
        return _file_schedule_svc.get_file_schedule_row(file_path=fp, agent_id=agent_id)
    return None


def _schedule_stop_success_payload(
    *,
    job_id: Optional[str],
    file_path: Optional[str],
    message: str = FOLDER_SCAN_JOB_STOPPED,
    cancelled_count: int = 1,
) -> Dict[str, Any]:
    """Immediate schedule-stop response (frontend can close the dialog)."""
    return {
        "success": True,
        "message": message,
        "job_id": job_id,
        "file_path": file_path,
        "status": "cancelled",
        "schedule_status": "cancelled",
        "stop_pending": False,
        "stop_complete": True,
        "scheduler_stopped": True,
        "scheduler_status": "cancelled",
        "cancelled": True,
        "cancelled_count": cancelled_count,
    }


def _deferred_schedule_stop_cleanup_sync(
    agent_id: str,
    job_id: str,
    file_path: str,
) -> None:
    """Persist cancel in DB (memory already cleared on the HTTP response path)."""
    jid = str(job_id or "").strip()
    fp = str(file_path or "").strip()
    if not jid:
        return
    try:
        _file_schedule_svc.cancel_file_schedule(jid)
        from app.utils.folder_scan_scheduler import _persist_folder_scan_job_status

        _persist_folder_scan_job_status(
            jid,
            status="cancelled",
            last_error="Upload schedule stopped by user",
        )
    except Exception as exc:
        import logging

        logging.getLogger(__name__).warning(
            "deferred schedule-stop cleanup failed job_id=%s: %s", jid, exc
        )


async def _deferred_schedule_stop_cleanup(
    agent_id: str,
    job_id: str,
    file_path: str,
) -> None:
    await run_in_threadpool(
        _deferred_schedule_stop_cleanup_sync,
        agent_id,
        job_id,
        file_path,
    )


def _schedule_stop_blocked_response(
    sched_row: Dict[str, Any],
    *,
    job_id: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """Return a response when file_schedules.status is already terminal (not cancelled)."""
    st = str(sched_row.get("status") or "").lower()
    if st == "completed":
        message = FOLDER_SCAN_JOB_STOP_ALREADY_COMPLETED
    elif st == "failed":
        message = FOLDER_SCAN_JOB_STOP_ALREADY_FAILED
    else:
        return None
    return {
        "success": False,
        "message": message,
        "status": st,
        "job_id": sched_row.get("job_id") or job_id,
        "file_path": sched_row.get("file_path"),
        "cancelled": False,
        "cancelled_count": 0,
    }


def _file_path_already_scheduled_response(
    file_path: str,
    existing: Dict[str, Any],
    *,
    reason: str = "active_schedule",
) -> Dict[str, Any]:
    """Reject duplicate folder upload/schedule for the same path."""
    details: Dict[str, Any] = {
        "file_path": file_path,
        "job_id": existing.get("job_id"),
        "schedule_status": existing.get("status"),
        "reason": reason,
    }
    if reason == "all_files_already_ingested":
        details["all_files_exist"] = True
    return {
        "success": False,
        "message": FOLDER_SCAN_FILE_PATH_ALREADY_EXISTS,
        "details": details,
    }


def _folder_upload_scheduler_rest_blocked() -> Optional[Dict[str, Any]]:
    """Reject new folder uploads while the global scheduler is in its rest window."""
    rest = global_scheduler_rest_phase()
    if not rest:
        return None
    remaining_label = str(rest.get("remaining_label") or "a few minutes")
    retry_after = int(rest.get("seconds_until_phase_end") or 0)
    return {
        "success": False,
        "message": FOLDER_SCAN_SCHEDULER_REST_UPLOAD_BLOCKED.format(
            remaining=remaining_label,
        ),
        "details": {
            "phase": "rest",
            "retry_after_seconds": retry_after,
            "seconds_until_phase_end": retry_after,
            "scheduler": schedule_details_payload(None),
        },
    }


def _folder_upload_path_blocked(
    scan_path: str,
    agent_id: str,
    doc_assistant: Any,
) -> Optional[Dict[str, Any]]:
    """
    Block upload/schedule when scheduler is resting or a job is already active for this path.

    Duplicate files are skipped per file during ingest (disk file_path + file name),
    so re-uploading the same folder with new files adds only the new files.
    """
    del doc_assistant  # reserved for future path-specific checks
    rest_blocked = _folder_upload_scheduler_rest_blocked()
    if rest_blocked:
        return rest_blocked

    aid = str(agent_id or "").strip()
    path_norm = os.path.normpath(os.path.abspath(str(scan_path or "").strip()))
    if not path_norm or not aid:
        return None

    existing = active_file_path_schedule(path_norm, aid)
    if existing:
        return _file_path_already_scheduled_response(
            path_norm, existing, reason="active_schedule"
        )
    return None


def _schedule_folder_response(
    *,
    success: bool,
    message: str,
    job_id: Optional[str] = None,
    file_path: Optional[str] = None,
    scheduled: bool = False,
    background_running: bool = False,
    scheduler_already_running: bool = False,
    scheduler_status: str = "",
    scheduler_message: Optional[str] = None,
    phase: Optional[str] = None,
) -> Dict[str, Any]:
    """Fixed keys for POST /document/schedule-folder responses."""
    return {
        "success": success,
        "message": message,
        "job_id": job_id,
        "file_path": file_path,
        "scheduled": scheduled,
        "background_running": background_running,
        "scheduler_already_running": scheduler_already_running,
        "scheduler_status": scheduler_status,
        "scheduler_message": scheduler_message,
        "phase": phase,
    }


def _flatten_schedule_folder_response(scheduled: Dict[str, Any]) -> Dict[str, Any]:
    """Map internal scan result to the public schedule-folder response shape."""
    if not scheduled.get("success"):
        return _schedule_folder_response(
            success=False,
            message=scheduled.get("message", "Could not start folder schedule"),
        )

    raw = scheduled.get("details") or {}
    jid = raw.get("job_id")
    file_path = raw.get("file_path") or raw.get("root_path")
    completed = bool(raw.get("completed", False))
    sched_block = dict(raw.get("schedule") or {})
    sched_enabled = bool(
        sched_block.get("schedule_enabled", raw.get("schedule_enabled", True))
    )
    phase = (sched_block.get("phase") or raw.get("phase") or "active").lower()
    sched_status = (raw.get("scheduler_status") or "").lower()

    bg = get_scheduled_job_status(str(jid)) if jid else None
    bg_status = (bg.get("status") or "").lower() if bg else ""
    active_statuses = frozenset(
        {"pending", "running", "starting", "resting", "sleeping", "cancelling"}
    )
    scheduler_already_running = bool(raw.get("scheduler_already_running")) or (
        bg_status in active_statuses
    )
    background_running = (
        not completed
        and sched_enabled
        and not raw.get("scheduler_stopped")
        and (
            scheduler_already_running
            or sched_status in active_statuses
            or (bool(raw.get("scheduled")) and not raw.get("scan_start_skipped"))
        )
    )
    scheduler_msg = (
        raw.get("scheduler_message")
        or (bg or {}).get("scheduler_message")
        or (scheduler_message_from_schedule_row(sched_block) if sched_block else None)
    )
    status = sched_status or bg_status or ("running" if background_running else "")
    display_message = scheduler_msg or scheduled.get("message") or ""

    return _schedule_folder_response(
        success=True,
        message=display_message,
        job_id=jid,
        file_path=file_path,
        scheduled=bool(raw.get("scheduled", True)) and sched_enabled and not completed,
        background_running=background_running,
        scheduler_already_running=scheduler_already_running,
        scheduler_status=status,
        scheduler_message=scheduler_msg,
        phase=phase,
    )


def _scheduled_scan_attach_response(
    *,
    agent_id: str,
    server_path_abs: str,
    job_id: str,
    message: str,
    scheduler_status: str,
    scheduler_already_running: bool = True,
    reuse_existing_job: bool = True,
    existing_bg: Optional[Dict[str, Any]] = None,
    completed: bool = False,
    uploaded: Optional[int] = None,
    skipped: Optional[int] = None,
) -> Dict[str, Any]:
    """Return when schedule-folder should not create a second scheduler for the same path."""
    cycle_s, active_s, rest_s = folder_scan_schedule_seconds()
    active_min = max(1, active_s // 60)
    rest_min = max(1, rest_s // 60)
    sched_row = _file_schedule_svc.get_file_schedule(
        file_path=server_path_abs, agent_id=str(agent_id)
    )
    sched_payload = schedule_details_payload(sched_row)
    phase = sched_payload.get("phase") or (existing_bg or {}).get("phase") or "active"
    sched_enabled = bool(sched_payload.get("schedule_enabled"))
    user_message = message
    if completed:
        user_message, _ = build_folder_scan_step_user_message(
            job_completed=True,
            uploaded=int(uploaded or 0),
            skipped=int(skipped or 0),
            schedule_row=sched_row,
            phase=phase,
        )
    if completed or not sched_enabled:
        scheduler_already_running = False
        scheduler_msg = scheduler_message_from_schedule_row(sched_row) if sched_row else message
    else:
        scheduler_msg = (existing_bg or {}).get("scheduler_message") or (
            scheduler_message_from_schedule_row(sched_row)
            if sched_row and sched_enabled
            else scheduler_status_message(
                phase=phase,
                rest_minutes=rest_min,
                active_minutes=active_min,
                seconds_remaining=sched_payload.get("seconds_until_phase_end"),
            )
        )
    return {
        "success": True,
        "message": user_message,
        "details": {
            "scheduled": sched_enabled and not completed,
            "scheduler_already_running": scheduler_already_running,
            "reuse_existing_job": reuse_existing_job,
            "scan_start_skipped": True,
            "job_id": job_id,
            "file_path": server_path_abs,
            "scheduler_status": scheduler_status,
            "scheduler_message": scheduler_msg,
            "schedule": sched_payload,
            "phase": phase,
            "resting": phase == "rest",
            "completed": completed,
            "uploaded": uploaded,
            "skipped": skipped,
            "hint": (
                "Background scheduler is already running for this folder path "
                "(20 min active / 10 min rest). Refresh the file list for progress, "
                "or call schedule-stop to cancel."
            ),
        },
    }


async def _start_scheduled_folder_scan(
    agent_id: str,
    server_path: str,
    *,
    batch_size: int = 100,
    schedule_enabled: bool = True,
    force_new: bool = False,
    free_time: Optional[int] = None,
) -> Dict[str, Any]:
    """Create a resumable scan job and run batches in a background thread (20 min on / 10 min off)."""
    from app.services.folder_scan_job_prompts import (
        FOLDER_SCAN_JOB_PATH_ALREADY_EXISTS,
        FOLDER_SCAN_JOB_REUSED,
        FOLDER_SCAN_JOB_SCHEDULE_ALREADY,
        FOLDER_SCAN_JOB_SCHEDULE_PATH_RUNNING,
        FOLDER_SCAN_JOB_SCHEDULED,
        FOLDER_SCAN_JOB_STARTED,
    )

    server_path_abs = os.path.abspath(str(server_path).strip())

    doc_assistant, error_msg = _get_document_assistant(agent_id)
    if not doc_assistant:
        return {"success": False, "message": error_msg, "details": {}}

    cycle_s, active_s, rest_s = folder_scan_schedule_seconds()
    active_min = max(1, active_s // 60)
    rest_min = max(1, rest_s // 60)

    existing_bg = find_scheduled_job_for_file_path(server_path_abs, str(agent_id))
    if not force_new and existing_bg and is_file_path_scheduler_running(
        server_path_abs, str(agent_id)
    ):
        return _scheduled_scan_attach_response(
            agent_id=agent_id,
            server_path_abs=server_path_abs,
            job_id=str(existing_bg.get("job_id") or ""),
            message=FOLDER_SCAN_JOB_SCHEDULE_PATH_RUNNING,
            scheduler_status=str(existing_bg.get("status") or "running"),
            existing_bg=existing_bg,
        )

    sched_row = _file_schedule_svc.get_file_schedule(
        file_path=server_path_abs, agent_id=str(agent_id)
    )
    if not force_new and sched_row and int(sched_row.get("schedule_enabled") or 0):
        if is_file_path_scheduler_running(server_path_abs, str(agent_id)):
            existing_bg = find_scheduled_job_for_file_path(server_path_abs, str(agent_id))
            return _scheduled_scan_attach_response(
                agent_id=agent_id,
                server_path_abs=server_path_abs,
                job_id=str((existing_bg or {}).get("job_id") or ""),
                message=FOLDER_SCAN_JOB_SCHEDULE_PATH_RUNNING,
                scheduler_status=str((existing_bg or {}).get("status") or "running"),
                scheduler_already_running=True,
                existing_bg=existing_bg,
            )
        scan_db_early = _open_scan_db()
        try:
            active_job = scan_db_early.get_active_folder_scan_job_for_root(server_path_abs)
        finally:
            scan_db_early.close()
        if active_job:
            jid = str(active_job.get("id") or "")
            try_resume_background_scheduler(
                str(agent_id),
                jid,
                file_path=server_path_abs,
                batch_size=batch_size,
                schedule_enabled=schedule_enabled,
            )
            bg = get_scheduled_job_status(jid) or {}
            return _scheduled_scan_attach_response(
                agent_id=agent_id,
                server_path_abs=server_path_abs,
                job_id=jid,
                message=FOLDER_SCAN_JOB_PATH_ALREADY_EXISTS,
                scheduler_status=str(active_job.get("status") or "running"),
                scheduler_already_running=bool(bg)
                or is_file_path_scheduler_running(server_path_abs, str(agent_id)),
                existing_bg=bg or None,
            )

    scan_db = _open_scan_db()
    try:
        if not force_new:
            active_job = scan_db.get_active_folder_scan_job_for_root(server_path_abs)
            if active_job:
                jid = str(active_job.get("id") or "")
                try_resume_background_scheduler(
                    str(agent_id),
                    jid,
                    file_path=server_path_abs,
                    batch_size=batch_size,
                    schedule_enabled=schedule_enabled,
                )
                bg = get_scheduled_job_status(jid) or {}
                print(
                    f"[SCAN] Reusing active job {jid} for {server_path_abs} (schedule-folder not duplicated)",
                    flush=True,
                )
                return _scheduled_scan_attach_response(
                    agent_id=agent_id,
                    server_path_abs=server_path_abs,
                    job_id=jid,
                    message=FOLDER_SCAN_JOB_REUSED,
                    scheduler_status=str(active_job.get("status") or bg.get("status") or "running"),
                    scheduler_already_running=bool(bg) or is_file_path_scheduler_running(
                        server_path_abs, str(agent_id)
                    ),
                    existing_bg=bg or None,
                )

            latest_job = scan_db.get_latest_folder_scan_job_for_root(server_path_abs)
            if latest_job:
                st = (latest_job.get("status") or "").lower()
                if st in ("completed", "failed", "cancelled"):
                    jid = str(latest_job.get("id") or "")
                    stats: Dict[str, Any] = {}
                    if latest_job.get("stats_json"):
                        try:
                            stats = json.loads(latest_job["stats_json"])
                        except Exception:
                            pass
                    print(
                        f"[SCAN] schedule-folder skipped new job — file_path already has {st} job "
                        f"{jid} for {server_path_abs}",
                        flush=True,
                    )
                    return _scheduled_scan_attach_response(
                        agent_id=agent_id,
                        server_path_abs=server_path_abs,
                        job_id=jid,
                        message=FOLDER_SCAN_JOB_PATH_ALREADY_EXISTS,
                        scheduler_status=st,
                        scheduler_already_running=is_file_path_scheduler_running(
                            server_path_abs, str(agent_id)
                        ),
                        reuse_existing_job=True,
                        existing_bg=get_scheduled_job_status(jid),
                        completed=st == "completed",
                        uploaded=stats.get("uploaded"),
                        skipped=stats.get("skipped"),
                    )
    finally:
        scan_db.close()

    start_result = await run_in_threadpool(
        doc_assistant.start_resumable_folder_scan,
        server_path_abs,
        batch_size,
    )
    if not start_result.get("success"):
        return {
            "success": False,
            "message": start_result.get("message", "Could not start folder scan job"),
            "details": start_result,
        }

    job_id = start_result.get("job_id")
    scheduled = schedule_folder_scan_job(
        agent_id,
        job_id,
        file_path=server_path_abs,
        batch_size=batch_size,
        schedule_enabled=schedule_enabled,
        free_time=free_time,
    )
    if not scheduled:
        if is_file_path_scheduler_running(server_path_abs, str(agent_id)):
            bg = find_scheduled_job_for_file_path(server_path_abs, str(agent_id)) or {}
            return {
                "success": True,
                "message": FOLDER_SCAN_JOB_SCHEDULE_PATH_RUNNING,
                "details": {
                    "scheduled": True,
                    "scheduler_already_running": True,
                    "job_id": bg.get("job_id") or job_id,
                    "file_path": server_path_abs,
                },
            }
        return {
            "success": False,
            "message": FOLDER_SCAN_JOB_SCHEDULE_ALREADY,
            "details": {"job_id": job_id},
        }

    sched_row = _file_schedule_svc.get_file_schedule(
        file_path=server_path_abs, agent_id=str(agent_id)
    )
    sched_payload = schedule_details_payload(sched_row)
    scheduler_msg = scheduler_message_from_schedule_row(sched_row)

    print(
        f"Scheduled folder upload started: batch_size={batch_size}, "
        f"file_path={server_path_abs}, cycle={cycle_s // 60} min "
        f"({active_min} min active / {rest_min} min rest), job_id={job_id}",
        flush=True,
    )
    print(scheduler_msg, flush=True)

    return {
        "success": True,
        "message": FOLDER_SCAN_JOB_SCHEDULED,
        "details": {
            "job_id": job_id,
            "file_path": server_path_abs,
            "root_path": server_path_abs,
            "root_folder_name": start_result.get("root_folder_name"),
            "dirs_total": start_result.get("dirs_total"),
            "batch_size": start_result.get("batch_size"),
            "max_file_size": start_result.get("max_file_size"),
            "scheduled": True,
            "schedule_enabled": schedule_enabled,
            "schedule": sched_payload,
            "phase": sched_payload.get("phase") or "active",
        },
    }


def _get_document_assistant(agent_id: str):
    """
    Helper function to get or create DocumentAssistant instance for an agent.
    Uses the shared document store (same SQLite + Chroma + disk roots for all FILE-capable agents).
    """
    if not DOCUMENT_ASSISTANT_AVAILABLE:
        return None, "Document assistant module is not available"

    if not agent_has_file_datasource(str(agent_id)):
        return None, "File datasource is not enabled for this agent"
    
    # Get agent config to get OpenAI API key
    agent_config = get_agent_config(agent_id=agent_id)
    if not agent_config:
        return None, f"Agent configuration not found for agent_id: {agent_id}"
    
    api_key = agent_config.get("api_key")
    if not api_key:
        return None, "OpenAI API key not configured for this agent (required for embeddings; Claude does not replace this)"

    try:
        doc_assistant = DocumentAssistant(
            upload_folder=FILE_UPLOAD_DISK_ROOT,
            image_folder=FILE_IMAGE_DISK_ROOT,
            extracted_folder=FILE_EXTRACT_DISK_ROOT,
            db_file=FILE_SYS_METADATA_DB,
            chroma_path=FILE_VECTORS_DB_DIR,
            openai_api_key=str(api_key) if api_key else None,
            claude_api_key=str(agent_config.get("claude_api_key") or ""),
            agent_id=agent_id,
        )
        return doc_assistant, None
    except Exception as e:
        return None, f"Error initializing DocumentAssistant: {str(e)}"


def _format_extracted_image_entries(
    image_rows: List[Tuple],
    agent_id: str,
) -> List[Dict[str, Any]]:
    """Format extracted_images rows as API image dicts (with served URLs)."""
    return [
        {
            "id": img[0],
            "path": convert_image_path_to_url(img[1], agent_id),
            "type": img[2],
            "page": img[3],
        }
        for img in image_rows
    ]


def _serve_document_by_id(doc_assistant: Any, file_id: str, *, inline: bool) -> FileResponse:
    """
    Resolve files.file_path from metadata DB and stream the document.
    Inline preview converts DOC/DOCX to PDF (application/pdf, Content-Disposition inline).
    Raises HTTPException(404) when the file is missing.
    """
    file_path = doc_assistant.get_file_path(file_id)
    if not file_path or not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail=f"File not found: {file_id}")

    file_info = doc_assistant.db.get_file(file_id)
    if not file_info:
        raise HTTPException(
            status_code=404, detail=f"File information not found: {file_id}"
        )

    filename = file_info[2] if len(file_info) > 2 else os.path.basename(file_path)
    try:
        serve_path, media_type, serve_name = resolve_preview_file_response(
            file_path,
            filename,
            file_id,
            inline=inline,
            cache_dir=DEFAULT_PREVIEW_PDF_CACHE_DIR,
        )
    except Exception as exc:
        raise HTTPException(
            status_code=503,
            detail=f"Could not prepare document preview: {exc}",
        ) from exc

    disposition = "inline" if inline else "attachment"
    return FileResponse(
        path=serve_path,
        media_type=media_type,
        filename=serve_name,
        headers={"Content-Disposition": f'{disposition}; filename="{serve_name}"'},
    )


def _format_browse_file_payload(
    file: Dict[str, Any],
    page_count: int,
    agent_id: str,
) -> Dict[str, Any]:
    """Build one browse file dict (thumbnail URL + page_count; image rows loaded via separate API)."""
    file_data: Dict[str, Any] = {
        "id": file["id"],
        "filename": file["filename"],
        "file_type": file["file_type"],
        "file_path": file["file_path"],
        "folder_path": file["folder_path"],
        "relative_path": file.get("relative_path", ""),
        "upload_date": file["upload_date"],
        "size": file["size"],
        "size_kb": round(file["size"] / 1024, 2) if file["size"] else 0,
        "preview": file.get("preview", ""),
        "page_count": page_count,
    }
    thumb = file.get("thumbnail_path") or None
    file_data["thumbnail_path"] = (
        convert_image_path_to_url(thumb, agent_id) if thumb else None
    )
    apply_preview_contract_to_mapping(
        file_data,
        file_id=str(file["id"]),
        agent_id=str(agent_id),
        filename=str(file.get("filename") or file.get("original_filename") or ""),
        file_type=str(file.get("file_type") or ""),
        relative_url=True,
    )
    return file_data


def _resolve_pages_count(file: Dict[str, Any]) -> int:
    """
    Resolve true page count for PDFs from the source file.
    """
    file_type = str(file.get("file_type") or "").lower()
    if file_type != "pdf":
        return 0

    file_path = str(file.get("file_path") or "").strip()
    if not file_path or not os.path.isfile(file_path):
        return 0

    try:
        import pymupdf  # lazy import to avoid hard dependency at module import time

        doc = pymupdf.open(file_path)
        try:
            return int(doc.page_count or 0)
        finally:
            doc.close()
    except Exception:
        return 0


def _enrich_browse_files_for_folder(doc_assistant: Any, folder_path: str, agent_id: str) -> List[Dict[str, Any]]:
    """Browse file rows plus resolved page_count values."""
    files = doc_assistant.get_files_in_folder_browse(folder_path)
    if not files:
        return []
    return [
        _format_browse_file_payload(
            f,
            _resolve_pages_count(f),
            agent_id,
        )
        for f in files
    ]


def _get_file_extracted_images_details(
    doc_assistant: Any,
    file_id: str,
    agent_id: str,
) -> Optional[Dict[str, Any]]:
    """Load extracted images for one file (agent-scoped). Runs on worker thread."""
    file_info = doc_assistant.db.get_file(file_id)
    if not file_info:
        return None
    rows = doc_assistant.db.get_extracted_images(file_id)
    return {
        "file_id": file_id,
        "images_count": len(rows),
        "images": _format_extracted_image_entries(rows, agent_id),
    }


def _sort_folder_children(node: Dict[str, Any]) -> None:
    node["children"].sort(key=lambda x: (x.get("folder_path") or "", x.get("folder_name") or ""))
    for ch in node["children"]:
        _sort_folder_children(ch)


def _count_tree_folders_and_files(node: Dict[str, Any]) -> tuple:
    """Return (folder_node_count, total_file_count) for subtree rooted at node."""
    files_here = len(node.get("files") or [])
    count = 1
    total_files = files_here
    for ch in node.get("children") or []:
        sub_c, sub_f = _count_tree_folders_and_files(ch)
        count += sub_c
        total_files += sub_f
    return count, total_files


def _norm_folder_path(p: Optional[str]) -> str:
    """Normalize logical folder_path / parent_path for stable parent-child linking."""
    if not p:
        return ""
    s = str(p).replace("\\", "/").strip()
    s = re.sub(r"/+", "/", s).strip("/")
    return s


def _build_all_files_folders_details(doc_assistant: Any, agent_id: str) -> Dict[str, Any]:
    """
    Build nested ``folder_tree`` and ``folders`` (single root). Uses light file listing and batched image counts.
    Per-file image rows are loaded via ``get-file-extracted-images``. Runs synchronously on the worker thread.
    """
    folders = doc_assistant.get_all_folders()
    by_path: Dict[str, Dict[str, Any]] = {}
    norm_to_node: Dict[str, Dict[str, Any]] = {}

    for folder in folders:
        if len(folder) == 5:
            folder_name, folder_path, file_count, parent_path, full_path = folder
        else:
            folder_name, folder_path, file_count, parent_path = folder
            full_path = ""

        files_data = _enrich_browse_files_for_folder(doc_assistant, folder_path, agent_id)
        node = {
            "folder_name": folder_name,
            "folder_path": folder_path,
            "file_count": file_count,
            "parent_path": parent_path or "",
            "full_path": full_path or "",
            "files": files_data,
            "children": [],
        }
        by_path[folder_path] = node
        nk = _norm_folder_path(folder_path)
        if nk:
            norm_to_node[nk] = node

    root_children: List[Dict[str, Any]] = []
    for _path, node in by_path.items():
        parent_raw = node.get("parent_path") or ""
        parent_key = _norm_folder_path(parent_raw)
        parent_node = norm_to_node.get(parent_key) if parent_key else None
        if parent_node is not None and parent_node is not node:
            parent_node["children"].append(node)
        else:
            root_children.append(node)

    root_children.sort(key=lambda x: (x.get("folder_path") or "", x.get("folder_name") or ""))
    for ch in root_children:
        _sort_folder_children(ch)

    root_files_data = _enrich_browse_files_for_folder(doc_assistant, "", agent_id)

    folder_tree: Dict[str, Any] = {
        "folder_name": "Root",
        "folder_path": "",
        "parent_path": "",
        "full_path": "",
        "file_count": len(root_files_data),
        "files": root_files_data,
        "children": root_children,
    }

    total_folders, total_files = _count_tree_folders_and_files(folder_tree)

    return {
        "total_folders": total_folders,
        "total_files": total_files,
        # Single nested tree; use folders[0] or folder_tree — same object graph.
        "folders": [folder_tree],
        "folder_tree": folder_tree,
    }


@router.post("/document/get-all-files-folders")
async def get_all_files_folders(request: GetFilesFoldersRequest):
    """
    Get all files and folders for document training center (Browse section).

    - ``details.folder_tree``: nested tree (Root with ``files`` and ``children``). Same graph as ``folders[0]``.
    - ``details.folders``: a one-element array ``[folder_tree]`` so legacy code that reads ``folders`` still
      receives a nested tree; recurse into each node's ``children`` and read ``files`` per folder.
    - ``details.total_folders`` / ``details.total_files``: counts over the full tree.
    - Each file includes ``page_count`` only; use ``POST /document/get-file-extracted-images`` for image rows.
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {},
            }

        details = await run_in_threadpool(
            _build_all_files_folders_details,
            doc_assistant,
            request.agent_id,
        )

        return {
            "success": True,
            "message": "Files and folders retrieved successfully",
            "details": details,
        }

    except Exception as e:
        return {
            "success": False,
            "message": f"Error getting files and folders: {str(e)}",
            "details": {},
        }


@router.post("/document/get-file-extracted-images")
async def get_file_extracted_images(request: GetFileExtractedImagesRequest):
    """
    Get extracted images for a single document file (lazy-load after browse).

    ``details.images``: list of ``{id, path, type, page}`` with ``path`` as a served URL.
    ``details.images_count``: same as ``len(details.images)``.
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {},
            }

        file_id = (request.file_id or "").strip()
        if not file_id:
            return {
                "success": False,
                "message": "Please provide a file_id",
                "details": {},
            }

        details = await run_in_threadpool(
            _get_file_extracted_images_details,
            doc_assistant,
            file_id,
            request.agent_id,
        )
        if details is None:
            return {
                "success": False,
                "message": f"File not found: {file_id}",
                "details": {},
            }

        return {
            "success": True,
            "message": "Extracted images retrieved successfully",
            "details": details,
        }

    except Exception as e:
        return {
            "success": False,
            "message": f"Error getting extracted images: {str(e)}",
            "details": {},
        }


_UPLOAD_FOLDER_SCHEDULE_HINT = (
    "Files saved. Call POST /training/document/schedule-folder with the same file_path "
    "to start background ingestion (20 min active / 10 min rest)."
)


@router.post(
    "/document/upload-folder",
    openapi_extra={
        "requestBody": {
            "content": {
                "multipart/form-data": {
                    "schema": {
                        "type": "object",
                        "required": ["agent_id"],
                        "properties": {
                            "agent_id": {"type": "string", "description": "Agent ID (required)"},
                            "server_folder_path": {"type": "string", "nullable": True, "description": "Server folder path (optional). When set, leave files empty."},
                            "files": {"type": "array", "items": {"type": "string", "format": "binary"}, "nullable": True, "description": "Files (optional). When uploading files, omit server_folder_path."},
                        },
                    }
                }
            }
        }
    },
)
@folder_upload_request_timeout
async def upload_folder(request: Request):
    """
    Save folder files to disk (multipart/form-data) and synchronously process them.
    Blocks and returns a single response only after ingestion is fully complete.
    """
    form = await request.form()
    agent_id = form.get("agent_id")
    if not agent_id:
        raise HTTPException(status_code=422, detail=[{"type": "missing", "loc": ["body", "agent_id"], "msg": "Field required", "input": None}])
    server_folder_path = form.get("server_folder_path")
    server_folder_path = server_folder_path.strip() or None if isinstance(server_folder_path, str) else None
    raw_files = form.getlist("files") if "files" in form else []
    files = [f for f in raw_files if hasattr(f, "read") and not isinstance(f, str)]

    try:
        doc_assistant, error_msg = _get_document_assistant(agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }

        rest_blocked = _folder_upload_scheduler_rest_blocked()
        if rest_blocked:
            return rest_blocked

        if server_folder_path:
            server_path = os.path.abspath(server_folder_path)
            if not os.path.exists(server_path) or not os.path.isdir(server_path):
                return {
                    "success": False,
                    "message": f"Server folder path does not exist or is not a directory: {server_folder_path}",
                    "details": {}
                }

            aid = str(agent_id).strip()
            blocked = await run_in_threadpool(
                _folder_upload_path_blocked, server_path, aid, doc_assistant
            )
            if blocked:
                return blocked

            import uuid
            job_id = str(uuid.uuid4())
            _file_schedule_svc.get_scheduler_status()

            print(f"📋 Scheduling background folder scan: {server_path} for agent {agent_id}")
            scheduled_job_id = await run_in_threadpool(
                schedule_folder_scan_job,
                aid,
                job_id,
                file_path=server_path,
            )
            if scheduled_job_id:
                return {
                    "success": True,
                    "message": "Folder scan job scheduled successfully.",
                    "job_id": scheduled_job_id,
                    "file_path": server_path
                }
            rest_blocked = _folder_upload_scheduler_rest_blocked()
            if rest_blocked:
                return rest_blocked
            return {
                "success": False,
                "message": "Could not start folder schedule.",
                "details": {}
            }

        if not files and not server_folder_path:
            return {
                "success": False,
                "message": "No files provided and no server folder path specified.",
                "details": {}
            }

        from pathlib import Path

        base_upload_dir = Path(doc_assistant.upload_folder)
        base_upload_dir.mkdir(parents=True, exist_ok=True)
        temp_base_dir = base_upload_dir
        aid = str(agent_id).strip()

        uploaded_count = 0
        skipped_count = 0
        total_chunks = 0
        saved_files = []

        try:
            for idx, file in enumerate(files or []):
                raw_name = (file.filename or "").strip() or f"file_{idx}"
                rel_path = Path(raw_name)
                if rel_path.is_absolute() or (rel_path.parts and rel_path.parts[0].endswith(":")):
                    rel_path = Path(*rel_path.parts[1:]) if len(rel_path.parts) > 1 else Path(rel_path.name or f"file_{idx}")
                if not rel_path.name:
                    rel_path = Path(f"file_{idx}")
                rel_path = rel_path.as_posix()

                file_path = temp_base_dir / rel_path
                file_path.parent.mkdir(parents=True, exist_ok=True)

                try:
                    content = await file.read()
                    with open(file_path, "wb") as buffer:
                        buffer.write(content)
                    saved_files.append(str(file_path))
                except Exception as e:
                    print(f"Error saving file {rel_path}: {str(e)}")
                    skipped_count += 1
                    continue

            if not saved_files:
                return {
                    "success": False,
                    "message": "Failed to save any uploaded files.",
                    "details": {}
                }

            direct_children = [p for p in temp_base_dir.iterdir()]
            single_dir = None
            if len(direct_children) == 1 and direct_children[0].is_dir():
                single_dir = direct_children[0]
            scan_path = single_dir if single_dir is not None else temp_base_dir
            scan_path_str = str(scan_path.resolve())

            blocked = await run_in_threadpool(
                _folder_upload_path_blocked, scan_path_str, aid, doc_assistant
            )
            if blocked:
                return blocked

            import uuid
            job_id = str(uuid.uuid4())
            _file_schedule_svc.get_scheduler_status()

            print(f"📋 Scheduling background folder scan for uploaded files: {scan_path_str}")
            scheduled_job_id = await run_in_threadpool(
                schedule_folder_scan_job,
                aid,
                job_id,
                file_path=scan_path_str,
            )
            if scheduled_job_id:
                return {
                    "success": True,
                    "message": "Folder scan job scheduled successfully.",
                    "job_id": scheduled_job_id,
                    "file_path": scan_path_str
                }
            return {
                "success": False,
                "message": "Could not start folder schedule.",
                "details": {}
            }

        except Exception as e:
            print(f"Error during folder upload processing: {str(e)}")
            raise e
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error uploading folder: {str(e)}",
            "details": {}
        }


@router.get("/document/bm25-status")
async def document_bm25_status(
    current_user=Depends(get_current_user),
):
    """Global BM25 sparse index rebuild / corruption status."""
    from app.utils.bm25_checkpoint import get_bm25_status, pkl_path_for_chroma

    if not DOCUMENT_ASSISTANT_AVAILABLE:
        return {
            "success": False,
            "message": "Document service unavailable",
            "details": {},
        }
    details = get_bm25_status(pkl_path_for_chroma(FILE_VECTORS_DB_DIR))
    return {"success": True, "message": details.get("message", ""), "details": details}


@router.post("/document/rebuild-bm25")
async def document_rebuild_bm25(
    agent_id: str = Form(..., description="Agent ID (file datasource; API key used for Chroma)"),
    current_user=Depends(get_current_user),
):
    """Start background rebuild of BM25 pickle from Chroma (idempotent if already rebuilding)."""
    from app.utils.bm25_checkpoint import get_bm25_status, pkl_path_for_chroma, schedule_background_rebuild

    if not DOCUMENT_ASSISTANT_AVAILABLE:
        return {
            "success": False,
            "message": "Document service unavailable",
            "details": {},
        }

    doc_assistant, error_msg = _get_document_assistant(agent_id)
    if not doc_assistant:
        return {"success": False, "message": error_msg, "details": {}}

    api_key = None
    agent_config = get_agent_config(agent_id=agent_id)
    if agent_config:
        api_key = agent_config.get("api_key")

    if not api_key:
        return {
            "success": False,
            "message": "OpenAI API key required for Chroma access during rebuild",
            "details": {},
        }

    started = schedule_background_rebuild(
        FILE_VECTORS_DB_DIR,
        str(api_key),
        pkl_path=pkl_path_for_chroma(FILE_VECTORS_DB_DIR),
    )
    details = get_bm25_status(pkl_path_for_chroma(FILE_VECTORS_DB_DIR))
    msg = (
        "BM25 rebuild started in background"
        if started
        else "BM25 rebuild already in progress or could not start"
    )
    return {"success": True, "message": msg, "details": details}


def _build_stop_schedule_response(
    stop_result: Dict[str, Any],
    *,
    agent_id: str,
) -> Dict[str, Any]:
    """API payload after stop_folder_upload_scheduler."""
    jid = str(stop_result.get("job_id") or "")
    file_path = str(stop_result.get("file_path") or "")
    sched = (
        _file_schedule_svc.get_file_schedule(file_path=file_path, agent_id=agent_id)
        if file_path
        else None
    )
    sched_payload = schedule_details_payload(sched)
    cancelled = bool(stop_result.get("cancelled"))
    stop_pending = bool(stop_result.get("stop_pending"))
    scheduler_cancel = bool(stop_result.get("scheduler_cancel_requested"))
    if stop_pending:
        message = str(stop_result.get("message") or FOLDER_SCAN_JOB_STOP_PENDING)
    elif scheduler_cancel or stop_result.get("schedule_disabled"):
        message = FOLDER_SCAN_JOB_STOPPED
    else:
        message = FOLDER_SCAN_JOB_STOP_NOT_RUNNING
    stats = {}
    if jid:
        svc = _open_scan_db()
        try:
            row = svc.get_folder_scan_job(jid)
            if row and row.get("stats_json"):
                try:
                    stats = json.loads(row["stats_json"])
                except Exception:
                    pass
        finally:
            svc.close()

    sched_row = (
        _file_schedule_svc.get_file_schedule_by_job_id(jid)
        if jid
        else None
    )
    if not sched_row and file_path:
        sched_row = _file_schedule_svc.get_file_schedule_row(
            file_path=file_path,
            agent_id=agent_id,
        )
    queue_status = str((sched_row or {}).get("status") or "").lower()
    job_db_status = get_folder_scan_job_status(jid) if jid else ""
    if stop_pending:
        schedule_status = "cancelling"
    elif job_db_status in ("cancelled", "completed", "failed"):
        schedule_status = job_db_status
    elif queue_status in ("cancelled", "completed", "failed"):
        schedule_status = queue_status
    else:
        schedule_status = "cancelled" if cancelled else (queue_status or "cancelled")
    stop_complete = not stop_pending

    return {
        "success": True,
        "message": message,
        "job_id": jid or None,
        "file_path": file_path or None,
        "status": schedule_status,
        "schedule_status": schedule_status,
        "completed": False,
        "cancelled": cancelled or schedule_status == "cancelled",
        "stop_pending": stop_pending,
        "stop_complete": stop_complete,
        "stop_schedule": True,
        "background": not stop_pending,
        "scheduler_status": stop_result.get("scheduler_status") or schedule_status,
        "scheduler_stopped": stop_complete,
        "schedule_disabled": bool(stop_result.get("schedule_disabled")),
        "schedule": sched_payload,
        "phase": sched_payload.get("phase"),
        "uploaded": stats.get("uploaded"),
        "skipped": stats.get("skipped"),
        "cancelled_count": 1 if cancelled or schedule_status == "cancelled" else 0,
    }

def _get_scheduler_info() -> Dict[str, Any]:
    """Return current scheduler phase info for API responses."""
    status = _file_schedule_svc.get_scheduler_status()
    if not status:
        return {
            "phase": "active",
            "active_seconds": 1200,
            "rest_seconds": 600,
            "phase_started_at": None,
            "next_phase_at": None,
        }
    return {
        "phase": (status.get("phase") or "active").lower(),
        "active_seconds": status.get("active_seconds", 1200),
        "rest_seconds": status.get("rest_seconds", 600),
        "phase_started_at": status.get("phase_started_at"),
        "next_phase_at": status.get("next_phase_at"),
    }


@router.post("/document/schedule-folder/list")
async def list_schedule_folders(request: ListSchedulesRequest):
    """
    List folder upload schedules for an agent.

    Optional ``file_path`` limits results to that path (and related subpaths).
    Returns ``file_path``, ``job_id``, and ``schedule_status`` for each row
    (including completed, failed, and cancelled). Also includes top-level
    ``scheduler`` info with current phase details.
    """
    if not DOCUMENT_ASSISTANT_AVAILABLE:
        return {
            "success": False,
            "message": "Document service unavailable",
            "schedules": [],
            "count": 0,
        }
    aid = request.agent_id.strip()
    path_filter = (request.file_path or "").strip() or None
    if not agent_has_file_datasource(aid):
        return {
            "success": False,
            "message": "File datasource is not enabled for this agent",
            "schedules": [],
            "count": 0,
        }
    try:
        rows = await run_in_threadpool(
            lambda: _file_schedule_svc.list_file_schedules_for_agent(
                agent_id=aid,
                file_path=path_filter,
            )
        )
        schedules = []
        for row in rows:
            schedules.append({
                "id": row.get("id"),
                "file_path": row.get("file_path"),
                "job_id": row.get("job_id"),
                "schedule_status": effective_schedule_status(row, agent_id=aid),
                "created_at": row.get("created_at"),
            })
        schedules = dedupe_schedule_list_by_file_path(schedules)

        return {
            "success": True,
            "message": f"Found {len(schedules)} schedule(s) for agent {aid}.",
            "schedules": schedules,
            "count": len(schedules),
            "scheduler": _get_scheduler_info(),
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Error listing schedules: {str(e)}",
            "schedules": [],
            "count": 0,
        }


@router.post("/document/schedule-folder")
@folder_upload_request_timeout
async def schedule_folder(request: ScheduleFolderRequest):
    """
    Start background folder ingestion for a path (20 min active / 10 min rest per cycle).
    Generates a new job_id, updates scheduler_status heartbeat, registers task in file_schedules,
    and returns immediately.
    """
    if not DOCUMENT_ASSISTANT_AVAILABLE:
        return _schedule_folder_response(
            success=False,
            message="Document service unavailable",
        )
    path = os.path.abspath(request.file_path.strip())
    if not os.path.exists(path) or not os.path.isdir(path):
        return _schedule_folder_response(
            success=False,
            message=f"Folder path does not exist or is not a directory: {request.file_path}",
        )
    try:
        aid = request.agent_id.strip()
        doc_assistant, error_msg = _get_document_assistant(aid)
        if not doc_assistant:
            return _schedule_folder_response(success=False, message=error_msg or "Agent not found")
        blocked = await run_in_threadpool(_folder_upload_path_blocked, path, aid, doc_assistant)
        if blocked:
            if blocked.get("details", {}).get("phase") == "rest":
                return _schedule_folder_response(
                    success=False,
                    message=blocked.get("message") or "Scheduler is on rest.",
                    phase="rest",
                    scheduler_status="rest",
                    scheduler_message=blocked.get("message"),
                )
            return blocked

        import uuid
        job_id = str(uuid.uuid4())
        
        # Ensure status initialized
        _file_schedule_svc.get_scheduler_status()

        # Enqueue task and alert background loop
        scheduled_job_id = await run_in_threadpool(
            schedule_folder_scan_job,
            aid,
            job_id,
            file_path=path,
        )
        
        if scheduled_job_id:
            return {
                "success": True,
                "message": "Folder scan job scheduled successfully.",
                "job_id": scheduled_job_id,
                "file_path": path
            }
        rest_blocked = _folder_upload_scheduler_rest_blocked()
        if rest_blocked:
            return _schedule_folder_response(
                success=False,
                message=rest_blocked.get("message") or "Scheduler is on rest.",
                phase="rest",
                scheduler_status="rest",
                scheduler_message=rest_blocked.get("message"),
            )
        return _schedule_folder_response(
            success=False,
            message="Could not start folder schedule",
        )
    except Exception as e:
        return _schedule_folder_response(
            success=False,
            message=f"Error starting folder schedule: {str(e)}",
        )


@router.post("/document/schedule-stop")
async def schedule_stop(request: StopScheduleFolderRequest):
    """
    Stop the background folder ingest scheduler.

    **Single-job mode** (``job_id`` provided): cancels that specific schedule when
    ``file_schedules.status`` is ``pending`` or ``processing``. Returns ``status`` from
    the queue row and a message when the job is already ``completed`` or ``failed``.

    **Stop-all mode** (``job_id`` omitted / null): cancels every active (pending/processing)
    schedule for the agent — use this when the frontend shows a list of active schedules
    and the user clicks "Stop All".

    Returns immediately after the job is marked cancelled (``stop_complete: true``).
    Any file still open in the worker stops at the next cancel check in the background.
    """
    aid = request.agent_id.strip()
    jid_raw = (request.job_id or "").strip()
    _INVALID = frozenset({"string", "null", "none", "undefined", "job_id", "uuid", "id"})

    # ---------- STOP-ALL MODE ----------
    if not jid_raw or jid_raw.lower() in _INVALID:
        try:
            active_rows = await run_in_threadpool(
                lambda: [
                    r
                    for r in _file_schedule_svc.list_file_schedules_for_agent(agent_id=aid)
                    if str(r.get("status") or "").lower() in ("pending", "processing")
                ]
            )

            for row in active_rows:
                row_jid = str(row.get("job_id") or "").strip()
                row_fp = str(row.get("file_path") or "")
                if row_jid:
                    await run_in_threadpool(
                        signal_immediate_folder_scan_stop,
                        aid,
                        row_jid,
                        row_fp,
                    )

            cancelled_job_ids = {
                str(r.get("job_id")).strip()
                for r in active_rows
                if r.get("job_id")
            }
            status = _file_schedule_svc.get_scheduler_status()
            if status and status.get("current_job_id") in cancelled_job_ids:
                _file_schedule_svc.update_scheduler_status(
                    phase=status.get("phase", "active"),
                    phase_started_at=status.get("phase_started_at") or datetime.now().isoformat(),
                    next_phase_at=status.get("next_phase_at") or datetime.now().isoformat(),
                    clear_job_id=True,
                )

            async def _stop_all_background() -> None:
                await run_in_threadpool(
                    lambda: _file_schedule_svc.cancel_all_active_schedules(agent_id=aid)
                )
                for row in active_rows:
                    row_jid = str(row.get("job_id") or "").strip()
                    row_fp = str(row.get("file_path") or "")
                    if row_jid:
                        await _deferred_schedule_stop_cleanup(aid, row_jid, row_fp)

            asyncio.create_task(_stop_all_background())

            return {
                "success": True,
                "message": f"{len(active_rows)} active schedule(s) stopped successfully.",
                "job_id": None,
                "cancelled_count": len(active_rows),
                "cancelled": active_rows,
                "scheduler_status": "cancelled",
                "status": "cancelled",
                "schedule_status": "cancelled",
                "stop_pending": False,
                "stop_complete": True,
                "scheduler_stopped": True,
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"Error stopping all schedules: {str(e)}",
                "cancelled_count": 0,
                "cancelled": [],
            }

    # ---------- SINGLE-JOB MODE ----------
    jid = jid_raw
    fp_req = (request.file_path or "").strip() or None
    try:
        sched_row = await run_in_threadpool(
            _lookup_file_schedule_for_stop,
            agent_id=aid,
            job_id=jid,
            file_path=fp_req,
        )
        if sched_row:
            prior_status = str(sched_row.get("status") or "").lower()
            if prior_status == "cancelled":
                return _schedule_stop_success_payload(
                    job_id=sched_row.get("job_id") or jid,
                    file_path=sched_row.get("file_path") or fp_req,
                )
            blocked = _schedule_stop_blocked_response(sched_row, job_id=jid)
            if blocked:
                return blocked
        else:
            prior_status = ""

        if prior_status not in ("pending", "processing", ""):
            # Row exists but status is unexpected; do not cancel silently.
            return {
                "success": False,
                "message": FOLDER_SCAN_JOB_STOP_NOT_RUNNING,
                "status": prior_status or None,
                "job_id": jid,
                "file_path": (sched_row or {}).get("file_path") or fp_req,
                "cancelled": False,
                "cancelled_count": 0,
            }

        stop_fp = (sched_row or {}).get("file_path") or fp_req or ""

        # 1. In-memory cancel first (instant; does not wait on scheduler DB lock)
        await run_in_threadpool(
            signal_immediate_folder_scan_stop,
            aid,
            jid,
            stop_fp,
        )

        # 2. Clear heartbeat so the UI does not show this job as still active
        status = _file_schedule_svc.get_scheduler_status()
        if status and status.get("current_job_id") == jid:
            _file_schedule_svc.update_scheduler_status(
                phase=status.get("phase", "active"),
                phase_started_at=status.get("phase_started_at") or datetime.now().isoformat(),
                next_phase_at=status.get("next_phase_at") or datetime.now().isoformat(),
                clear_job_id=True,
            )

        # 3. Persist cancel + full teardown in background (HTTP returns immediately)
        asyncio.create_task(_deferred_schedule_stop_cleanup(aid, jid, stop_fp))

        return _schedule_stop_success_payload(
            job_id=jid,
            file_path=stop_fp or None,
        )
    except Exception as e:
        return {
            "success": False,
            "message": f"Error stopping folder schedule: {str(e)}",
            "details": {},
        }


@router.post("/document/schedule-stop/status")
async def schedule_stop_status(request: ScheduleStopStatusRequest):
    """
    Optional status check for ``schedule-stop`` (not required for the stop dialog).

    Use after ``schedule-stop`` only if you need to know when background ingest
    has fully finished (optional; the stop dialog can close on ``schedule-stop``).
    """
    aid = request.agent_id.strip()
    jid = request.job_id.strip()
    try:
        pending = await run_in_threadpool(is_folder_scan_stop_pending, jid)
        sched_row = await run_in_threadpool(
            _file_schedule_svc.get_file_schedule_by_job_id,
            jid,
        )
        queue_status = str((sched_row or {}).get("status") or "").lower()
        job_status = await run_in_threadpool(get_folder_scan_job_status, jid) or ""
        if pending:
            schedule_status = "cancelling"
        elif job_status in ("cancelled", "completed", "failed"):
            schedule_status = job_status
        elif queue_status in ("cancelled", "completed", "failed"):
            schedule_status = queue_status
        else:
            schedule_status = "cancelled"
        return {
            "success": True,
            "job_id": jid,
            "stop_pending": pending,
            "stop_complete": not pending,
            "status": schedule_status,
            "schedule_status": schedule_status,
            "scheduler_stopped": not pending,
            "cancelled": schedule_status == "cancelled",
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Error reading stop status: {str(e)}",
            "job_id": jid,
            "stop_pending": True,
            "stop_complete": False,
        }


@router.post("/document/upload-files")
async def upload_files(
    agent_id: str = Form(...),
    folder_path: str = Form(""),
    files: List[UploadFile] = FastAPIFile(...)
):
    """
    Upload one or more files to the specified folder.
    Accepts both single file and multiple files in a single request.
    
    Args:
        agent_id: Agent ID (required)
        folder_path: Destination folder path (e.g., "MyFolder" or "" for Root)
        files: One or more files to upload
    
    Returns:
        Success response with upload details for each file
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }
        
        if not files or len(files) == 0:
            return {
                "success": False,
                "message": "No files provided",
                "details": {}
            }
        
        folder_path = folder_path.strip() if folder_path else ""
        supported_extensions = ['pdf', 'docx', 'doc', 'xlsx', 'xls', 'csv', 'txt', 'png', 'jpg', 'jpeg']
        
        success_count = 0
        failed_count = 0
        total_chunks = 0
        failed_files = []
        successful_files = []
        
        import io
        
        for file in files:
            try:
                # Skip macOS AppleDouble/resource-fork files (._*); they are not real documents
                if file.filename and file.filename.startswith('._'):
                    failed_count += 1
                    failed_files.append({
                        "filename": file.filename,
                        "reason": "Resource-fork metadata file (._*), not a document"
                    })
                    continue
                file_ext = file.filename.split('.')[-1].lower() if file.filename else ""
                
                if file_ext not in supported_extensions:
                    failed_count += 1
                    failed_files.append({
                        "filename": file.filename,
                        "reason": f"Unsupported file type: {file_ext}"
                    })
                    continue
                
                # Read file content asynchronously and convert to BytesIO
                file_content = await file.read()
                file_bytes = io.BytesIO(file_content)
                file_bytes.filename = file.filename
                
                # Upload document
                result = await run_in_threadpool(doc_assistant.upload_document, file_bytes, folder_path)
                
                if result.get('success'):
                    success_count += 1
                    total_chunks += result.get('chunks_created', 0)
                    successful_files.append({
                        "filename": file.filename,
                        "file_id": result.get('file_id', ''),
                        "chunks_created": result.get('chunks_created', 0),
                        "images_extracted": result.get('images_extracted', 0)
                    })
                else:
                    failed_count += 1
                    failed_files.append({
                        "filename": file.filename,
                        "reason": result.get('message', 'Upload failed')
                    })
            
            except Exception as e:
                failed_count += 1
                failed_files.append({
                    "filename": file.filename if file.filename else "Unknown",
                    "reason": str(e)
                })
        
        total_files = len(files)
        if success_count == 0 and total_files > 0:
            return {
                "success": False,
                "message": "All files were skipped because they are already present.",
                "details": {
                    "success_count": success_count,
                    "failed_count": failed_count,
                    "total_files": total_files,
                    "total_chunks": total_chunks,
                    "folder_path": folder_path or "Root",
                    "successful_files": successful_files,
                    "failed_files": failed_files if failed_files else []
                }
            }
        return {
            "success": success_count > 0,
            "message": f"Upload completed. {success_count}/{total_files} file(s) uploaded successfully.",
            "details": {
                "success_count": success_count,
                "failed_count": failed_count,
                "total_files": total_files,
                "total_chunks": total_chunks,
                "folder_path": folder_path or "Root",
                "successful_files": successful_files,
                "failed_files": failed_files if failed_files else []
            }
        }
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error uploading files: {str(e)}",
            "details": {}
        }


@router.post("/document/create-folder")
async def create_folder(request: CreateFolderRequest):
    """
    Create a new folder in the document system.
    This is used when users want to create a new folder before uploading files.
    Similar to the "Create New Folder" functionality in file_traning_centre.py
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }
        
        if not request.folder_name or not request.folder_name.strip():
            return {
                "success": False,
                "message": "Please provide a folder name",
                "details": {}
            }
        
        folder_name = request.folder_name.strip()
        folder_path = folder_name  # Use folder_name as folder_path for root-level folders
        
        # Check if folder already exists
        folders = await run_in_threadpool(doc_assistant.get_all_folders)
        for folder in folders:
            if len(folder) >= 2 and folder[1] == folder_path:
                return {
                    "success": False,
                    "message": f"Folder '{folder_name}' already exists",
                    "details": {}
                }
        
        # Create folder in database
        await run_in_threadpool(doc_assistant.db.save_folder, folder_name, folder_path, "", "")
        
        return {
            "success": True,
            "message": f"Folder '{folder_name}' created successfully",
            "details": {
                "folder_name": folder_name,
                "folder_path": folder_path
            }
        }
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error creating folder: {str(e)}",
            "details": {}
        }


@router.get("/document/preview-file/{file_id}")
async def preview_document_file(file_id: str, agent_id: str):
    """
    Preview the original document by file_id (inline in browser).
    Resolves path from file_sys_metadata.db via DocumentService.get_file_path.
    """
    doc_assistant, error_msg = _get_document_assistant(agent_id)
    if not doc_assistant:
        raise HTTPException(status_code=400, detail=error_msg or "Document assistant unavailable")

    return await run_in_threadpool(
        lambda: _serve_document_by_id(doc_assistant, file_id, inline=True)
    )


def _serve_document_page(
    doc_assistant: Any,
    file_id: str,
    page_number: int,
    *,
    inline: bool = False,
) -> FileResponse:
    """Download or inline-view a single page as PDF."""
    file_path = doc_assistant.get_file_path(file_id)
    if not file_path or not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail=f"File not found: {file_id}")

    file_info = doc_assistant.db.get_file(file_id)
    if not file_info:
        raise HTTPException(
            status_code=404, detail=f"File information not found: {file_id}"
        )

    filename = file_info[2] if len(file_info) > 2 else os.path.basename(file_path)
    try:
        full_pdf_path, media_type, serve_name = resolve_preview_file_response(
            file_path,
            filename,
            file_id,
            inline=True,
            cache_dir=DEFAULT_PREVIEW_PDF_CACHE_DIR,
        )
        if media_type != "application/pdf":
            raise HTTPException(
                status_code=400,
                detail="Page download is only available for PDF-based documents",
            )
        page_path = extract_and_cache_page_pdf(
            full_pdf_path, file_id, page_number
        )
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(
            status_code=503,
            detail=f"Could not prepare page PDF: {exc}",
        ) from exc

    base = os.path.splitext(serve_name)[0]
    page_name = f"{base}_page_{int(page_number)}.pdf"
    disposition = "inline" if inline else "attachment"
    return FileResponse(
        path=page_path,
        media_type="application/pdf",
        filename=page_name,
        headers={"Content-Disposition": f'{disposition}; filename="{page_name}"'},
    )


def _serve_document_page_image(
    doc_assistant: Any,
    file_id: str,
    page_number: int,
) -> Response:
    """Render a single PDF page to PNG in memory (no disk cache)."""
    file_path = doc_assistant.get_file_path(file_id)
    if not file_path or not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail=f"File not found: {file_id}")

    file_info = doc_assistant.db.get_file(file_id)
    if not file_info:
        raise HTTPException(
            status_code=404, detail=f"File information not found: {file_id}"
        )

    filename = file_info[2] if len(file_info) > 2 else os.path.basename(file_path)
    try:
        full_pdf_path, media_type, serve_name = resolve_preview_file_response(
            file_path,
            filename,
            file_id,
            inline=True,
            cache_dir=DEFAULT_PREVIEW_PDF_CACHE_DIR,
        )
        if media_type != "application/pdf":
            raise HTTPException(
                status_code=400,
                detail="Page image is only available for PDF-based documents",
            )
        png_bytes = render_page_png_bytes(full_pdf_path, page_number)
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(
            status_code=503,
            detail=f"Could not render page image: {exc}",
        ) from exc

    base = os.path.splitext(serve_name)[0]
    page_name = f"{base}_page_{int(page_number)}.png"
    return Response(
        content=png_bytes,
        media_type="image/png",
        headers={"Content-Disposition": f'inline; filename="{page_name}"'},
    )


@router.get("/document/download-page/{file_id}")
async def download_document_page(file_id: str, agent_id: str, page_number: int):
    """Download one PDF page by file_id and 1-based page_number."""
    doc_assistant, error_msg = _get_document_assistant(agent_id)
    if not doc_assistant:
        raise HTTPException(status_code=400, detail=error_msg or "Document assistant unavailable")
    if page_number < 1:
        raise HTTPException(status_code=400, detail="page_number must be >= 1")

    return await run_in_threadpool(
        lambda: _serve_document_page(
            doc_assistant, file_id, page_number, inline=False
        )
    )


@router.get("/document/page-image/{file_id}")
async def render_document_page_image(file_id: str, agent_id: str, page_number: int):
    """Render one PDF page as PNG at request time (no persistent image cache)."""
    doc_assistant, error_msg = _get_document_assistant(agent_id)
    if not doc_assistant:
        raise HTTPException(status_code=400, detail=error_msg or "Document assistant unavailable")
    if page_number < 1:
        raise HTTPException(status_code=400, detail="page_number must be >= 1")

    return await run_in_threadpool(
        lambda: _serve_document_page_image(doc_assistant, file_id, page_number)
    )


@router.post("/document/download-file")
async def download_document_file(request: RemoveRequest):
    """
    Download a document file by file_id.
    Returns the file for download.
    Similar to download functionality in file_traning_centre.py
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }

        return await run_in_threadpool(
            lambda: _serve_document_by_id(doc_assistant, request.id, inline=False)
        )

    except HTTPException as e:
        return {
            "success": False,
            "message": e.detail,
            "details": {}
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Error downloading file: {str(e)}",
            "details": {}
        }


@router.post("/document/delete-file")
async def delete_document_file(request: RemoveRequest):
    """
    Delete a document file by file_id.
    Similar to delete_document in file_traning_centre.py
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }
        
        # Delete document
        await run_in_threadpool(doc_assistant.delete_document, request.id)
        
        return {
            "success": True,
            "message": "File deleted successfully",
            "details": {
                "file_id": request.id
            }
        }
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error deleting file: {str(e)}",
            "details": {}
        }


@router.post("/document/delete-folder")
async def delete_document_folder(request: DeleteFolderRequest):
    """
    Delete a folder and all its files (including subfolders).
    Similar to delete_document but for folders.
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }
        
        if not request.folder_path or not request.folder_path.strip():
            return {
                "success": False,
                "message": "Please provide a folder path",
                "details": {}
            }
        
        folder_path = request.folder_path.strip()
        
        # Delete folder and all contents
        # Uses the unified delete_folder method which handles disk, vector DB, and SQL cleanup
        result = await run_in_threadpool(doc_assistant.delete_folder, folder_path)
        
        if result:
            return {
                "success": True,
                "message": "Folder deleted successfully",
                "details": {
                    "folder_path": folder_path
                }
            }
        else:
            return {
                "success": False,
                "message": "Failed to delete folder",
                "details": {
                    "folder_path": folder_path
                }
            }
    
    except Exception as e:
        return {
            "success": False,
            "message": f"Error deleting folder: {str(e)}",
            "details": {}
        }


@router.post("/document/search-files")
async def search_uploaded_files(request: SearchFilesRequest):
    """
    Search across all uploaded files by keyword.
    Matches against the original filename, folder path, and content preview.
    Optionally restrict the search to a specific folder.
    """
    try:
        doc_assistant, error_msg = _get_document_assistant(request.agent_id)
        if not doc_assistant:
            return {
                "success": False,
                "message": error_msg,
                "details": {}
            }

        keyword = (request.keyword or "").strip().lower()
        if not keyword:
            return {
                "success": False,
                "message": "Please provide a keyword to search for",
                "details": {}
            }

        folder_filter = request.folder_path.strip() if request.folder_path else None

        # SQL-level filter: uses agent_id + folder_path indexes, omits full_content, bounded by LIMIT.
        # Tuple indices from search_files:
        # 0=id, 1=filename, 2=original_filename, 3=file_type, 4=file_path,
        # 5=folder_path, 6=relative_path, 7=upload_date, 8=file_size,
        # 9=content_preview, 10=image_path, 11=thumbnail_path
        rows = await run_in_threadpool(
            doc_assistant.db.search_files,
            keyword,
            folder_filter,
        )

        matched = [
            {
                "id": f[0],
                "filename": f[2] or f[1],
                "file_type": f[3],
                "folder_path": f[5] or "Root",
                "relative_path": f[6],
                "upload_date": f[7],
                "size": f[8],
                "size_kb": round((f[8] or 0) / 1024, 2),
                "preview": (f[9] or "")[:300],
            }
            for f in rows
        ]

        return {
            "success": True,
            "message": f"{len(matched)} file(s) found for keyword '{request.keyword}'",
            "details": {
                "keyword": request.keyword,
                "folder_filter": request.folder_path or None,
                "total_matches": len(matched),
                "files": matched,
            }
        }

    except Exception as e:
        return {
            "success": False,
            "message": f"Error searching files: {str(e)}",
            "details": {}
        }
