import os
import json
from pathlib import Path
from typing import List, Dict
from tqdm import tqdm
from ragformance.generators.generators.data_generator_interface import CorpusGeneratorInterface
from ragformance.generators.generators.llm_prompt.file_processors.md_processor import MarkdownProcessor
from ragformance.generators.generators.llm_prompt.file_processors.pdf_processor import PdfProcessor
from ragformance.models.corpus import DocModel
from ragformance.generators.generators.llm_prompt.file_processors.file_processor_interface import FileProcessor

class CorpusBuilder(CorpusGeneratorInterface):
    def __init__(self, output_file: str, processors: List[FileProcessor]):
        self.output_file = output_file
        self.processors = processors
        self.documents: List[DocModel] = []

    def generate_corpus(self, docs_folder: Path, config: Dict = {}) -> List[DocModel]:
        if config is None:
            config = {}
        all_documents = []
        print(docs_folder)
        for file_path in tqdm(docs_folder.glob("*"), desc="Processing files"):
            for processor in self.processors:
                if processor.can_process(file_path):
                    documents = processor.process(file_path)
                    all_documents.extend(documents)
                    break
        self.documents = all_documents
        self._save_to_jsonl()

    def _save_to_jsonl(self):
        if not self.documents:
            print("No documents to save.")
            return

        output_dir = os.path.dirname(self.output_file)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        with open(self.output_file, "w", encoding="utf-8") as fout:
            for doc in self.documents:
                fout.write(json.dumps(doc.to_dict(), ensure_ascii=False) + "\n")
        print(f"Saved {len(self.documents)} documents to {self.output_file}")

if __name__ == "__main__":
    doc_folder = Path("/Users/to125419/CRT/projects/RAGFORmance/data/test/")
    output_file = "/Users/to125419/CRT/projects/RAGFORmance/data/DatabaseToloxa/Bosch.jsonl"
    processors = [PdfProcessor(), MarkdownProcessor()]
    builder = CorpusBuilder(output_file=output_file, processors=processors)
    builder.generate_corpus(doc_folder)
