from __future__ import annotations

from pathlib import Path

from django.conf import settings
from django.core.management.base import BaseCommand, CommandError

from review_agent.regulatory_review.services.rag_embedding import get_embedding_provider
from review_agent.regulatory_review.services.rag_index import build_chroma_index
from review_agent.regulatory_review.services.rule_loader import load_rule_file


class Command(BaseCommand):
    help = "构建 NMPA 法规材料本地 ChromaDB RAG 索引。"

    def add_arguments(self, parser):
        parser.add_argument("--provider", default=None, help="覆盖 REGULATORY_RAG_PROVIDER。")

    def handle(self, *args, **options):
        rule_set = load_rule_file()
        source_dir = Path(settings.BASE_DIR) / rule_set["source_material_dir"]
        if not source_dir.exists():
            raise CommandError(f"法规材料目录不存在：{source_dir}")
        try:
            provider = get_embedding_provider(options["provider"])
            count = build_chroma_index(source_dir=source_dir, embedding_provider=provider, reset=True)
        except Exception as exc:
            raise CommandError(str(exc)) from exc
        self.stdout.write(
            self.style.SUCCESS(
                f"已构建法规 RAG 索引：collection={settings.REGULATORY_RAG_COLLECTION}, chunks={count}"
            )
        )