from pathlib import Path
import tempfile
import zipfile
import gradio as gr

from splitter import split_fasta_per_sequence


def process_fasta(file_obj):
    if file_obj is None:
        return "Please upload a FASTA or TXT file.", None

    input_path = Path(file_obj.name)
    text = input_path.read_text(encoding="utf-8", errors="replace")

    if ">" not in text:
        return "No FASTA headers found. Your file should contain lines starting with '>'.", None

    work_dir = Path(tempfile.mkdtemp())
    out_dir = work_dir / "split_sequences"
    created_files, seq_count = split_fasta_per_sequence(text, out_dir)

    if seq_count == 0:
        return "No sequences were found.", None

    zip_path = work_dir / "split_sequences.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for file_path in created_files:
            zipf.write(file_path, arcname=file_path.name)

    message = f"Done. Created {seq_count} FASTA files. Download the ZIP below."
    return message, str(zip_path)


with gr.Blocks() as demo:
    gr.Markdown("# FASTA Splitter\nUpload one FASTA/TXT file and get one FASTA file per sequence.")

    fasta_input = gr.File(label="Upload FASTA or TXT file", file_types=[".fasta", ".fas", ".fa", ".txt", ".fna"])
    run_btn = gr.Button("Split sequences")
    status = gr.Textbox(label="Status", interactive=False)
    output_zip = gr.File(label="Download split FASTA files")

    run_btn.click(fn=process_fasta, inputs=fasta_input, outputs=[status, output_zip])


demo.launch()