import gradio as gr
from Bio import Entrez
from datetime import datetime
import os

# Function to fetch PubMed articles
def fetch_pubmed(query, email, start_date, end_date, max_results=100):
    Entrez.email = email
    handle = Entrez.esearch(
        db="pubmed",
        term=query,
        retmax=max_results,
        sort="pub_date",
        mindate=start_date,
        maxdate=end_date,
        datetype="pdat"
    )
    results = Entrez.read(handle)
    handle.close()
    ids = results["IdList"]
    return ids

# Function to fetch article details
def fetch_details(pubmed_ids):
    handle = Entrez.efetch(db="pubmed", id=",".join(pubmed_ids), retmode="xml")
    records = Entrez.read(handle)
    handle.close()
    return records

# Function to compile Markdown content
def compile_summaries(records):
    content = "# PubMed Search Results\n\n"
    for record in records["PubmedArticle"]:
        title = record["MedlineCitation"]["Article"]["ArticleTitle"]
        abstract = record["MedlineCitation"]["Article"].get("Abstract", {}).get("AbstractText", ["No abstract available."])[0]
        url = f'https://pubmed.ncbi.nlm.nih.gov/{record["MedlineCitation"]["PMID"]}/'
        content += f"## [{title}]({url})\n\n{abstract}\n\n"
    return content

# Main function for the Gradio interface
def pubmed_search(mesh_terms, email, start_date, end_date):
    if not mesh_terms or not email:
        return "Please provide MeSH terms and email.", None, None
    
    # Join MeSH terms for query
    query = " AND ".join([f'"{term}"[MeSH Terms]' for term in mesh_terms])
    
    # Fetch PubMed IDs
    pubmed_ids = fetch_pubmed(query, email, start_date, end_date)
    if not pubmed_ids:
        return "No articles found for the given search terms and date range.", None, None
    
    # Fetch article details
    records = fetch_details(pubmed_ids)
    if not records["PubmedArticle"]:
        return "No detailed articles found.", None, None
    
    # Compile summaries
    markdown_content = compile_summaries(records)
    text_content = markdown_content.replace("\n", "\r\n")
    
    # Save files temporarily for download
    with open("results.md", "w") as md_file:
        md_file.write(markdown_content)
    with open("results.txt", "w") as txt_file:
        txt_file.write(text_content)
    
    return "Search completed successfully. Download your files below.", "results.md", "results.txt"

# Gradio interface components
with gr.Blocks() as app:
    gr.Markdown("# PubMed Search Tool for Thyroid AI Research")
    
    with gr.Row():
        mesh_term_input = gr.Textbox(label="Enter a MeSH Term", placeholder="e.g., thyroid", interactive=True)
        add_button = gr.Button("Add MeSH Term")
        mesh_terms_box = gr.Textbox(label="Added MeSH Terms", interactive=False, lines=2)
    
    start_date = gr.Date(label="Start Date")
    end_date = gr.Date(label="End Date")
    email_input = gr.Textbox(label="Email", placeholder="Your email (required by PubMed API)", interactive=True)
    
    search_button = gr.Button("Search PubMed")
    status_output = gr.Textbox(label="Status")
    
    markdown_file = gr.File(label="Markdown File", interactive=False)
    text_file = gr.File(label="Text File", interactive=False)

    # Logic for adding MeSH terms
    def add_mesh_term(term, terms):
        if term.strip():
            terms = terms.split(",") if terms else []
            terms.append(term.strip())
            return ", ".join(terms)
        return terms

    # Bind functions to interface
    add_button.click(fn=add_mesh_term, inputs=[mesh_term_input, mesh_terms_box], outputs=mesh_terms_box)
    search_button.click(fn=pubmed_search, inputs=[mesh_terms_box, email_input, start_date, end_date], outputs=[status_output, markdown_file, text_file])

# Launch app
app.launch()