import gradio as gr
import json
import PyPDF2
import markdown
import tempfile

def merge_files(files):
    merged_content = []
    for file in files:
        content = ""
        if file.name.endswith('.txt'):
            content = file.read().decode('utf-8')
        elif file.name.endswith('.pdf'):
            pdf_reader = PyPDF2.PdfReader(file)
            for page in pdf_reader.pages:
                content += page.extract_text()
        elif file.name.endswith('.md'):
            content = file.read().decode('utf-8')
            content = markdown.markdown(content)
        
        lines = content.split('\n')
        for line in lines:
            if line.strip():
                merged_content.append(line.strip())
    
    jsonl_output = ""
    for i, text in enumerate(merged_content):
        json_obj = {"id": i, "text": text}
        jsonl_output += json.dumps(json_obj) + "\n"
    
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jsonl", prefix="merged_")
    temp_file.write(jsonl_output.encode('utf-8'))
    temp_file.close()
    
    return temp_file.name

iface = gr.Interface(
    fn=merge_files,
    inputs=gr.File(label="Upload .txt, .pdf, or .md files", file_count="multiple"),
    outputs=gr.File(label="Download merged JSONL file"),
    title="Multiple Files to Merged Indexed JSONL Converter",
    description="Upload multiple .txt, .pdf, or .md files to convert and merge them into a single indexed JSONL file for LLM training."
)

if __name__ == "__main__":
    iface.launch()