from __future__ import annotations

import argparse
import asyncio
from typing import Iterable

from app.services.lightrag_service import lightrag_service
from app.services.supabase_rest import supabase_rest


def fetch_documents(
    workspace_id: str | None,
    limit: int | None,
) -> Iterable[dict]:
    filters = {}
    if workspace_id:
        filters["workspace_id"] = workspace_id
    docs = supabase_rest.select(
        "documents",
        filters,
        columns="id,raw_text,user_id,workspace_id,title",
        order="updated_at.desc",
        limit=limit,
    )
    return docs


async def rebuild_indexes(workspace_id: str | None, limit: int | None) -> None:
    documents = fetch_documents(workspace_id, limit)
    for doc in documents:
        text = str(doc.get("raw_text") or "").strip()
        if not text:
            continue
        document_id = str(doc.get("id"))
        user_id = str(doc.get("user_id"))
        workspace = str(doc.get("workspace_id")) if doc.get("workspace_id") else None
        title = doc.get("title") if isinstance(doc.get("title"), str) else None
        print(f"[LightRAG:init] indexing {document_id} ({workspace or 'user'})")  # noqa: T201
        supabase_rest.update(
            "documents", {"id": document_id}, {"index_status": "pending"}
        )
        await lightrag_service.index_document_async(
            document_id=document_id,
            user_id=user_id,
            workspace_id=workspace,
            text=text,
            title=title,
        )
        supabase_rest.update(
            "documents", {"id": document_id}, {"index_status": "completed"}
        )


def main() -> None:
    parser = argparse.ArgumentParser(description="初始化 LightRAG 索引")
    parser.add_argument("--workspace", help="仅处理指定 workspace", default=None)
    parser.add_argument("--limit", type=int, help="限制处理文档数量", default=None)
    args = parser.parse_args()
    asyncio.run(rebuild_indexes(args.workspace, args.limit))


if __name__ == "__main__":
    main()