from django.core.management.base import BaseCommand

from apps.knowledge.models import KnowledgeDocument, KnowledgeItem
from apps.knowledge.services import (
    index_published_solution,
    index_resolved_ticket,
    rebuild_document_chunks,
    rebuild_item_chunks,
)
from apps.solutions.models import Solution
from apps.tickets.models import Ticket


class Command(BaseCommand):
    help = "使用当前 Embedding 部署重建知识、文档和已发布解决方案向量"

    def handle(self, *args, **options):
        item_count = 0
        document_count = 0
        solution_count = 0
        ticket_count = 0
        for item in KnowledgeItem.objects.select_related("knowledge_base").iterator():
            rebuild_item_chunks(item)
            item_count += 1
        for document in KnowledgeDocument.objects.select_related("knowledge_base").iterator():
            rebuild_document_chunks(document)
            document_count += 1
        for solution in Solution.objects.filter(status="published").iterator():
            index_published_solution(solution)
            solution_count += 1
        for ticket in (
            Ticket.objects.filter(
                status__in=["resolved", "closed"],
                archived_at__isnull=True,
            )
            .select_related("category", "resolution")
            .iterator()
        ):
            index_resolved_ticket(ticket)
            ticket_count += 1
        self.stdout.write(
            self.style.SUCCESS(
                "知识向量重建完成："
                f"条目 {item_count}，文档 {document_count}，已发布方案 {solution_count}，"
                f"历史工单 {ticket_count}"
            )
        )