from django.core.management.base import BaseCommand
from shaoApp.rag_service import RAGChatbotService


class Command(BaseCommand):
    help = 'Build FAISS vector store from database data'

    def add_arguments(self, parser):
        parser.add_argument(
            '--force',
            action='store_true',
            help='Force rebuild even if vector store exists',
        )

    def handle(self, *args, **options):
        self.stdout.write(
            self.style.SUCCESS('Starting vector store build...')
        )
        
        try:
            chatbot = RAGChatbotService()
            
            if options['force']:
                self.stdout.write('Force rebuild requested...')
                chatbot.build_vector_store()
            else:
                # Try to load existing vector store first
                if chatbot.load_vector_store():
                    self.stdout.write(
                        self.style.SUCCESS('Vector store already exists and loaded successfully!')
                    )
                else:
                    self.stdout.write('No existing vector store found. Building new one...')
                    chatbot.build_vector_store()
            
            # Get statistics
            stats = chatbot.get_statistics()
            self.stdout.write(
                self.style.SUCCESS(
                    f'Vector store build completed!\n'
                    f'Applications: {stats.get("total_applications", 0)}\n'
                    f'PDF Documents: {stats.get("total_pdf_documents", 0)}\n'
                    f'Extracted Data: {stats.get("total_extracted_data", 0)}\n'
                    f'Councils: {len(stats.get("councils", []))}'
                )
            )
            
        except Exception as e:
            self.stdout.write(
                self.style.ERROR(f'Error building vector store: {str(e)}')
            )