from django.core.management.base import BaseCommand
from ai_agent.rag_service import RAGService
import os

class Command(BaseCommand):
    help = 'Ingest data into the AI Knowledge Base'

    def add_arguments(self, parser):
        parser.add_argument('--text', type=str, help='Text content to ingest')
        parser.add_argument('--file', type=str, help='Path to text file to ingest')
        parser.add_argument('--source', type=str, default='manual', help='Source metadata')
        parser.add_argument('--type', type=str, default='knowledge', help='Type of content (knowledge/answer_bank)')

    def handle(self, *args, **options):
        text = options.get('text')
        file_path = options.get('file')
        source = options.get('source')
        content_type = options.get('type')

        if not text and not file_path:
            self.stdout.write(self.style.ERROR('Please provide either --text or --file'))
            return

        rag = RAGService()

        if file_path:
            try:
                if not os.path.exists(file_path):
                     self.stdout.write(self.style.ERROR(f'File not found: {file_path}'))
                     return
                
                with open(file_path, 'r', encoding='utf-8') as f:
                    text_content = f.read()
                    chunks = self._chunk_text(text_content)
                    self.stdout.write(f"Ingesting {len(chunks)} chunks from file...")
                    for i, chunk in enumerate(chunks):
                        rag.add_document(chunk, {"source": source, "type": content_type, "filename": file_path, "chunk_id": i})
                    self.stdout.write(self.style.SUCCESS(f'Successfully ingested file: {file_path}'))
            except Exception as e:
                self.stdout.write(self.style.ERROR(f'Error reading file: {e}'))
                return

        if text:
            rag.add_document(text, {"source": source, "type": content_type})
            self.stdout.write(self.style.SUCCESS('Successfully ingested text snippet'))

    def _chunk_text(self, text, chunk_size=1000):
        """Simple chunking by characters for now"""
        return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]