main.py

# pyright: strict
import argparse
import asyncio
import logging
import sys
from pathlib import Path

from src.config.logging import setup_logging, update_logging_from_config

setup_logging()
logger = logging.getLogger(__name__)

from pydantic import ValidationError  # noqa: E402

from src.adapters.docling_adapter import DoclingAdapter  # noqa: E402
from src.adapters.embeddings import embedding_mode  # noqa: E402
from src.adapters.http_client import HTTPClientPool  # noqa: E402
from src.adapters.minio_adapter import MinIOAdapter  # noqa: E402
from src.api.factory import ServerFactory  # noqa: E402
from src.api.rest.auth.jwt_handler import JWTHandler  # noqa: E402
from src.config.logging.interceptor import FrameworkInterceptor  # noqa: E402
from src.core.accessor import ConfigAccessor  # noqa: E402
from src.core.admin import WorkspaceManager  # noqa: E402
from src.core.health import run_preflight_checks  # noqa: E402
from src.core.rag_builder import RagFactory  # noqa: E402

# Восстанавливаем propagate=True для всех логгеров сразу после их физического импорта
FrameworkInterceptor.enforce()


async def index_command(path: str) -> None:
    """Index files from the given path into the RAG pipeline."""
    embedding_mode.set("document")

    secrets = ConfigAccessor.get_secrets()
    if not secrets.llm_api_key:
        logger.error("LLM_API_KEY is not set in secrets")
        return
    if not secrets.embed_api_key:
        logger.error("EMBED_API_KEY is not set in secrets")
        return
    if not secrets.qdrant_url:
        logger.error("QDRANT_URL is not set in secrets")
        return

    path_obj = Path(path)
    if not path_obj.exists():
        logger.error(f"Path does not exist: {path}")
        return

    files_to_process: list[Path] = []
    if path_obj.is_file():
        files_to_process.append(path_obj)
    else:
        files_to_process = (
            list(path_obj.rglob("*.md"))
            + list(path_obj.rglob("*.txt"))
            + list(path_obj.rglob("*.pdf"))
            + list(path_obj.rglob("*.docx"))
        )

    if not files_to_process:
        logger.warning(f"No supported files found in {path}")
        return

    logger.info(f"Found {len(files_to_process)} files to process")

    rag_pipeline = None
    try:
        await HTTPClientPool.init_client()
        try:
            await run_preflight_checks()
        except RuntimeError as e:
            logger.critical(str(e))
            sys.exit(1)

        rag_pipeline = await RagFactory.create_rag()
        FrameworkInterceptor.enforce()
        docling_adapter = DoclingAdapter()
        FrameworkInterceptor.enforce()

        minio_adapter = None
        if secrets.minio_endpoint:
            minio_adapter = MinIOAdapter()

        core_config = ConfigAccessor.get_core_config()
        workspace = core_config.workspace_name or "default"

        for file_path in files_to_process:
            logger.info(f"Processing: {file_path}")
            try:
                ext = file_path.suffix.lower()
                doc_id = None
                if ext in (".md", ".txt"):
                    content_list = docling_adapter.process_file(str(file_path))
                else:
                    content_list, doc_id = await rag_pipeline.parse_with_fallback(
                        str(file_path), workspace
                    )

                for content in content_list:
                    if minio_adapter:
                        for key in ["img_path", "table_img_path", "equation_img_path"]:
                            if content.get("type") in [
                                "image",
                                "table",
                                "equation",
                            ] and content.get(key):
                                image_path = content[key]
                                if image_path:
                                    uri = await minio_adapter.upload_file(
                                        workspace, image_path
                                    )
                                    content[key] = uri

                await rag_pipeline.insert_content_list(
                    content_list=content_list,
                    file_path=str(file_path),
                    doc_id=doc_id,
                )
                logger.info(f"Successfully indexed: {file_path}")
            except Exception as e:
                logger.error(f"Error processing {file_path}: {e}")

    finally:
        if rag_pipeline:
            await asyncio.shield(rag_pipeline.shutdown_workers())
            await asyncio.shield(rag_pipeline.finalize_storages())
        await asyncio.shield(HTTPClientPool.close_client())

        import shutil

        core_config = ConfigAccessor.get_core_config()
        workspace = core_config.workspace_name or "default"
        shutil.rmtree(Path("output") / workspace, ignore_errors=True)


async def query_command(question: str) -> None:
    """Query the RAG system with the given question."""
    embedding_mode.set("query")

    secrets = ConfigAccessor.get_secrets()
    if not secrets.llm_api_key:
        logger.error("LLM_API_KEY is not set in secrets")
        return
    if not secrets.embed_api_key:
        logger.error("EMBED_API_KEY is not set in secrets")
        return
    if not secrets.qdrant_url:
        logger.error("QDRANT_URL is not set in secrets")
        return

    rag_pipeline = None
    try:
        await HTTPClientPool.init_client()
        try:
            await run_preflight_checks()
        except RuntimeError as e:
            logger.critical(str(e))
            sys.exit(1)

        rag_pipeline = await RagFactory.create_rag()
        FrameworkInterceptor.enforce()
        result = await rag_pipeline.aquery(question, mode="hybrid", stream=False)
        print(f"\nAnswer: {result}\n")
    finally:
        if rag_pipeline:
            await asyncio.shield(rag_pipeline.shutdown_workers())
            await asyncio.shield(rag_pipeline.finalize_storages())
        await asyncio.shield(HTTPClientPool.close_client())


async def check_command() -> None:
    """Run preflight health checks."""
    secrets = ConfigAccessor.get_secrets()
    if not secrets.llm_api_key:
        logger.error("LLM_API_KEY is not set in secrets")
        return
    if not secrets.embed_api_key:
        logger.error("EMBED_API_KEY is not set in secrets")
        return
    if not secrets.qdrant_url:
        logger.error("QDRANT_URL is not set in secrets")
        return

    try:
        await HTTPClientPool.init_client()
        try:
            await run_preflight_checks()
        except RuntimeError as e:
            logger.critical(str(e))
            sys.exit(1)
        FrameworkInterceptor.enforce()
        logger.info("All checks passed.")
    finally:
        await asyncio.shield(HTTPClientPool.close_client())


async def serve_command(server_type: str, port: int, webui: bool = False) -> None:
    """Start the API server (REST or MCP)."""

    secrets = ConfigAccessor.get_secrets()
    if not secrets.llm_api_key:
        logger.error("LLM_API_KEY is not set in secrets")
        return
    if not secrets.embed_api_key:
        logger.error("EMBED_API_KEY is not set in secrets")
        return
    if not secrets.qdrant_url:
        logger.error("QDRANT_URL is not set in secrets")
        return

    server_config = ConfigAccessor.get_server_config()

    token_secret = secrets.token_secret

    if not token_secret or not token_secret.get_secret_value():
        logger.error(
            "TOKEN_SECRET is required but not set in .env. "
            "Server will not start without a secret. "
            "Please configure it and use 'python main.py trivialtoken' to generate an access token."
        )
        sys.exit(1)

    jwt_handler = JWTHandler(
        token_secret=token_secret.get_secret_value(),
        algorithm=server_config.jwt_algorithm,
        expire_hours=server_config.token_expire_hours,
    )
    logger.info("JWT authentication configured")

    # Validate server_type
    if server_type != "rest":
        raise ValueError(f"Invalid server type: {server_type}")

    # REST server uses lifespan for graceful shutdown
    server = ServerFactory.create(
        server_type=server_type,  # type: ignore[arg-type]
        server_config=server_config,
        port=port,
        jwt_handler=jwt_handler,
        webui=webui,
    )

    logger.info(f"Starting {server_type.upper()} server on {server_config.host}:{port}")
    FrameworkInterceptor.enforce()
    await server.start()


def trivialtoken_command() -> None:
    """Generate a long-lived trivial JWT token for CLI usage."""
    secrets = ConfigAccessor.get_secrets()
    server_config = ConfigAccessor.get_server_config()

    if not secrets.token_secret or not secrets.token_secret.get_secret_value():
        logger.error(
            "TOKEN_SECRET is not set or empty in secrets. Cannot generate trivial token."
        )
        sys.exit(1)
        return

    try:
        jwt_handler = JWTHandler(
            token_secret=secrets.token_secret.get_secret_value(),
            algorithm=server_config.jwt_algorithm,
            expire_hours=server_config.token_expire_hours,
        )
        token = jwt_handler.create_token(username="admin", custom_expire_hours=8760)
        print(token)
    except Exception as e:
        logger.error(f"Failed to generate trivial token: {e}")
        sys.exit(1)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="My RAG Pipeline Orchestrator",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Required arguments for all commands:
  --workspace NAME       Logical name for the collection/namespace (prevents data corruption)
  --max-async N          Maximum concurrent API requests (I/O-bound limit)
  --max-workers N        Batch processing threads (CPU-bound limit)

Example usage:
  python main.py index --path ./docs --workspace openwrt --max-async 8 --max-workers 4
  python main.py query --question "How to configure VLAN?" --workspace openwrt --max-async 8 --max-workers 4
  python main.py serve --type rest --port 8000 --workspace openwrt --max-async 8 --max-workers 4
  python main.py check --workspace openwrt --max-async 8 --max-workers 4
        """,
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Common arguments for all commands
    def add_common_args(subparser: argparse.ArgumentParser) -> None:
        subparser.add_argument(
            "--workspace",
            type=str,
            required=True,
            help="Workspace name (logical namespace for data isolation in Qdrant/Neo4j)",
        )
        subparser.add_argument(
            "--max-async",
            type=int,
            required=True,
            help="Maximum concurrent API requests (prevents rate limiting)",
        )
        subparser.add_argument(
            "--max-workers",
            type=int,
            required=True,
            help="Batch processing threads (controls CPU-bound parallelism)",
        )

    index_parser = subparsers.add_parser("index", help="Index files or directories")
    index_parser.add_argument(
        "--path", required=True, help="Path to file or directory to index"
    )
    add_common_args(index_parser)

    query_parser = subparsers.add_parser("query", help="Query the RAG system")
    query_parser.add_argument("--question", required=True, help="Question to ask")
    add_common_args(query_parser)

    check_parser = subparsers.add_parser("check", help="Run preflight health checks")
    check_parser.add_argument(
        "--workspace", type=str, required=False, help="Workspace name"
    )
    check_parser.add_argument(
        "--max-async", type=int, required=False, help="Maximum concurrent API requests"
    )
    check_parser.add_argument(
        "--max-workers", type=int, required=False, help="Batch processing threads"
    )

    serve_parser = subparsers.add_parser("serve", help="Start the API server")
    serve_parser.add_argument(
        "--type",
        choices=["rest"],
        default="rest",
        help="Server type: rest (FastAPI)",
    )
    serve_parser.add_argument(
        "--port",
        type=int,
        required=True,
        help="Server port (required)",
    )
    serve_parser.add_argument(
        "--webui", action="store_true", help="Mount WebUI on the same port"
    )
    add_common_args(serve_parser)

    # Trivial token subcommand
    subparsers.add_parser(
        "trivialtoken",
        help="Generate a long-lived trivial JWT token",
        description="Generate a long-lived trivial JWT token",
    )

    # Workspaces subcommand
    workspaces_parser = subparsers.add_parser("workspaces", help="Manage workspaces")
    workspaces_subparsers = workspaces_parser.add_subparsers(
        dest="workspace_command", required=True
    )

    workspaces_subparsers.add_parser("list", help="List all workspaces")

    delete_parser = workspaces_subparsers.add_parser(
        "delete", help="Delete a workspace"
    )
    delete_parser.add_argument(
        "--name", required=True, help="Name of the workspace to delete"
    )
    delete_parser.add_argument("--force", action="store_true", help="Confirm deletion")

    args = parser.parse_args()

    # Build overrides dictionary from CLI arguments
    overrides: dict[str, str | int] = {}
    if hasattr(args, "workspace") and args.workspace is not None:
        overrides["workspace_name"] = args.workspace
    if hasattr(args, "max_async") and args.max_async is not None:
        overrides["max_async"] = args.max_async
    if hasattr(args, "max_workers") and args.max_workers is not None:
        overrides["max_workers"] = args.max_workers

    if args.command == "check" or args.command == "workspaces":
        overrides.setdefault("workspace_name", "default")
        overrides.setdefault("max_async", 1)
        overrides.setdefault("max_workers", 1)

    # Load configuration with CLI overrides before any command execution
    try:
        ConfigAccessor.load(overrides=overrides)
        update_logging_from_config(ConfigAccessor.get_logging_config())
    except ValidationError as e:
        logger.error(f"Configuration validation failed:\n{e}")
        sys.exit(1)

    try:
        if args.command == "index":
            asyncio.run(index_command(args.path))
        elif args.command == "query":
            asyncio.run(query_command(args.question))
        elif args.command == "check":
            asyncio.run(check_command())
        elif args.command == "serve":
            asyncio.run(
                serve_command(args.type, args.port, getattr(args, "webui", False))
            )
        elif args.command == "trivialtoken":
            trivialtoken_command()
        elif args.command == "workspaces":
            if args.workspace_command == "list":
                try:
                    workspaces = asyncio.run(WorkspaceManager.list_workspaces())
                    if not workspaces:
                        print("No workspaces found.")
                    else:
                        print("Workspaces:")
                        for w in workspaces:
                            print(f"- {w}")
                except Exception as e:
                    logger.error(f"Failed to list workspaces: {e}")
            elif args.workspace_command == "delete":
                try:
                    workspaces = asyncio.run(WorkspaceManager.list_workspaces())
                    if args.name not in workspaces:
                        logger.warning(
                            f"Workspace '{args.name}' not found in active workspaces."
                        )
                        sys.exit(0)

                    if not args.force:
                        user_input = input(
                            f"Are you sure you want to delete workspace '{args.name}'? (y/n): "
                        )
                        if user_input.strip().lower() != "y":
                            print("Deletion cancelled.")
                            sys.exit(0)

                    asyncio.run(WorkspaceManager.delete_workspace(args.name))
                    print(f"Workspace '{args.name}' deleted successfully.")
                except Exception as e:
                    logger.error(f"Failed to delete workspace '{args.name}': {e}")
    except KeyboardInterrupt:
        logger.info("Application shutdown requested via KeyboardInterrupt")
        try:
            core_config = ConfigAccessor.get_core_config()
            workspace = core_config.workspace_name
            if workspace:
                import shutil

                shutil.rmtree(Path("output") / workspace, ignore_errors=True)
        except Exception:
            pass
        sys.exit(0)


if __name__ == "__main__":
    main()