main.py

"""
ForgeRAG one-shot launcher.

Running `python main.py` in the project root does the right thing
in three cases:

    1. ./forgerag.yaml exists  -> use it, start uvicorn.
    2. $FORGERAG_CONFIG set    -> use it, start uvicorn.
    3. Nothing configured        -> write a minimal default config
                                    to ./forgerag.yaml, then start.

Overrides:
    python main.py --host 127.0.0.1 --port 9000 --reload
    python main.py --config path/to/my.yaml
    python main.py --init-only       # write default config and exit

The generated default uses zero external services:
    - SQLite at ./storage/forgerag.db
    - ChromaDB at ./storage/chroma
    - Local blob store at ./storage/blobs
    - Embedder / LLM = litellm with OPENAI_API_KEY (override via setup wizard)

For a richer wizard, run:  python scripts/setup.py
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path
from typing import Any

# Ensure repo root is on sys.path so `import config` etc. work when
# the user invokes `python main.py` directly.
_ROOT = Path(__file__).resolve().parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))


log = logging.getLogger("forgerag.main")


# ---------------------------------------------------------------------------
# Default minimal config
# ---------------------------------------------------------------------------


DEFAULT_CONFIG_PATH = _ROOT / "forgerag.yaml"


def _minimal_default_config() -> dict[str, Any]:
    """
    Smallest config that runs out of the box.

    Infrastructure only. LLM models, API keys, retrieval strategy,
    and parsing options are managed via /settings (DB-backed,
    frontend-editable). Defaults are seeded on first startup.
    """
    return {
        "parser": {
            "backends": {"pymupdf": {"enabled": True}},
        },
        "storage": {
            "mode": "local",
            "local": {"root": "./storage/blobs"},
        },
        "files": {
            "hash_algorithm": "sha256",
            "max_bytes": 524288000,
        },
        "persistence": {
            "relational": {
                "backend": "sqlite",
                "sqlite": {
                    "path": "./storage/forgerag.db",
                    "journal_mode": "wal",
                },
            },
            "vector": {
                "backend": "chromadb",
                "chromadb": {
                    "mode": "persistent",
                    "persist_directory": "./storage/chroma",
                    "collection_name": "forgerag",
                    "dimension": 1536,
                    "distance": "cosine",
                },
            },
        },
        "embedder": {
            "dimension": 1536,
        },
    }


def _write_default_config(path: Path) -> None:
    try:
        import yaml
    except ImportError as e:
        raise RuntimeError("pyyaml not installed. pip install pyyaml") from e
    path.parent.mkdir(parents=True, exist_ok=True)
    header = (
        "# ForgeRAG config (auto-generated by main.py).\n"
        "# Infrastructure only — LLM models, API keys, retrieval strategy\n"
        "# are managed via /settings API after startup.\n"
        "#\n"
        "# For an interactive wizard: python scripts/setup.py\n"
        "# To configure models/keys: visit http://localhost:8000/settings\n"
        "\n"
    )
    with open(path, "w", encoding="utf-8") as f:
        f.write(header)
        yaml.safe_dump(
            _minimal_default_config(),
            f,
            sort_keys=False,
            default_flow_style=False,
            allow_unicode=True,
        )


# ---------------------------------------------------------------------------
# Config resolution
# ---------------------------------------------------------------------------


def resolve_config_path(cli_path: Path | None) -> Path:
    """
    Precedence (highest first):
        1. --config argument
        2. FORGERAG_CONFIG env var
        3. ./forgerag.yaml (current directory)
        4. None -> caller writes the default to ./forgerag.yaml
    """
    if cli_path is not None:
        return cli_path.resolve()

    env_path = os.environ.get("FORGERAG_CONFIG")
    if env_path:
        p = Path(env_path)
        return p.resolve()

    cwd_candidate = Path.cwd() / "forgerag.yaml"
    if cwd_candidate.exists():
        return cwd_candidate.resolve()

    return DEFAULT_CONFIG_PATH.resolve()


# ---------------------------------------------------------------------------
# Pre-flight checks on the resolved config
# ---------------------------------------------------------------------------


def _preflight(cfg_path: Path) -> None:
    """Validate the config and print a concise startup banner."""
    try:
        from config import load_config
    except ImportError as e:
        print(f"error: cannot import config package: {e}", file=sys.stderr)
        print(
            "hint: did you install dependencies? pip install sqlalchemy pydantic pyyaml",
            file=sys.stderr,
        )
        raise SystemExit(2)

    try:
        cfg = load_config(cfg_path)
    except Exception as e:
        print(f"error: failed to load {cfg_path}: {e}", file=sys.stderr)
        raise SystemExit(1)

    print()
    print(f"  config        {cfg_path}")
    print(f"  relational    {cfg.persistence.relational.backend}")
    print(f"  vector        {cfg.persistence.vector.backend} (dim={cfg.embedder.dimension})")
    print(f"  blob          {cfg.storage.mode}")
    print(f"  embedder      {cfg.embedder.backend}")
    print(f"  generator     {cfg.answering.generator.backend} / {cfg.answering.generator.model}")
    print()

    print(
        "  note: LLM models and API keys are configured via /settings\n"
        "        Visit http://localhost:8000/settings after startup.\n"
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="main.py",
        description=(
            "Launch ForgeRAG. If forgerag.yaml exists in the current "
            "directory (or FORGERAG_CONFIG points at one), it is used. "
            "Otherwise a minimal default config is written to ./forgerag.yaml "
            "and the server starts."
        ),
    )
    p.add_argument(
        "--config",
        type=Path,
        default=None,
        help="Path to forgerag.yaml. Overrides env var and cwd discovery.",
    )
    p.add_argument(
        "--host",
        type=str,
        default=os.environ.get("FORGERAG_HOST", "0.0.0.0"),
        help="Bind host (default: 0.0.0.0, or $FORGERAG_HOST)",
    )
    p.add_argument(
        "--port",
        type=int,
        default=int(os.environ.get("FORGERAG_PORT", "8000")),
        help="Bind port (default: 8000, or $FORGERAG_PORT)",
    )
    p.add_argument(
        "--reload",
        action="store_true",
        help="Enable uvicorn --reload for development.",
    )
    p.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Uvicorn worker count (default: 4). Ignored when --reload is set.",
    )
    p.add_argument(
        "--log-level",
        type=str,
        default="info",
        help="uvicorn log level (critical / error / warning / info / debug / trace).",
    )
    p.add_argument(
        "--init-only",
        action="store_true",
        help="Write the default config (if missing) and exit without starting the server.",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()

    cfg_path = resolve_config_path(args.config)

    if not cfg_path.exists():
        print(f"no config found, writing default to {cfg_path}")
        _write_default_config(cfg_path)
        print("default config ready. edit it to customize.")
        print("for an interactive wizard: python scripts/setup.py")

    if args.init_only:
        return 0

    # Point everything downstream at the resolved config
    os.environ["FORGERAG_CONFIG"] = str(cfg_path)

    # Initialise logging early so preflight and all downstream modules
    # write to the daily-rotated log file under logs/.
    try:
        from config import load_config as _load_cfg
        from config.logging import setup_logging

        _early_cfg = _load_cfg(cfg_path)
        setup_logging(_early_cfg.logging)
    except Exception as _log_err:
        # Fallback: basic stderr logging so the server still starts
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s %(levelname)-8s %(name)s  %(message)s",
        )
        logging.getLogger("forgerag.main").warning(
            "Failed to initialise structured logging, using basicConfig: %s",
            _log_err,
        )

    _preflight(cfg_path)

    try:
        import uvicorn
    except ImportError:
        print(
            "error: uvicorn not installed. pip install 'uvicorn[standard]' fastapi python-multipart",
            file=sys.stderr,
        )
        return 2

    uvicorn_kwargs: dict[str, Any] = dict(
        host=args.host,
        port=args.port,
        log_level=args.log_level,
    )
    if args.reload:
        uvicorn_kwargs["reload"] = True
        # --reload is incompatible with --workers
    elif args.workers > 1:
        uvicorn_kwargs["workers"] = args.workers

    # Pass the factory string so uvicorn's --reload watcher can
    # re-import after file changes.
    uvicorn.run("api.app:app", **uvicorn_kwargs)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())