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
    
    # Convert float timestamps to datetime objects
    try:
        if isinstance(start_date, float):
            start_date = datetime.fromtimestamp(start_date)
        if isinstance(end_date, float):
            end_date = datetime.fromtimestamp(end_date)
        
        start_date_str = start_date.strftime("%Y-%m-%d %H:%M:%S")
        end_date_str = end_date.strftime("%Y-%m-%d %H:%M:%S")
    except ValueError as e:
        return f"Error processing dates: {str(e)}", None, None

    # Join MeSH terms for query
    query = " AND ".join([f'"{term}"[MeSH Terms]' for term in mesh_terms.split(",")])
    
    # Fetch PubMed IDs
    pubmed_ids = fetch_pubmed(query, email, start_date_str, end_date_str)
    if not pubmed_ids:
        return "No articles found for the given search terms and date range.", None, None

    # Fetch article details
    try:
        records = fetch_details(pubmed_ids)
    except Exception as e:
        return f"Error fetching article details: {str(e)}", None, None

    # Compile markdown content
    markdown_content = compile_summaries(records)

    # Create temporary files
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    markdown_filename = f"pubmed_results_{timestamp}.md"
    text_filename = f"pubmed_results_{timestamp}.txt"

    # Save files
    try:
        with open(markdown_filename, "w", encoding="utf-8") as f:
            f.write(markdown_content)
        with open(text_filename, "w", encoding="utf-8") as f:
            f.write(markdown_content)
    except Exception as e:
        return f"Error saving files: {str(e)}", None, None

    return "Search completed successfully!", markdown_filename, text_filename

# 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

# 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
        )
    
    with gr.Row():
        start_date = gr.DateTime(label="Start Date", value="2025-01-01 00:00:00")
        end_date = gr.DateTime(label="End Date", value="2025-01-07 00:00:00")
        
    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)

    # 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]
    )

app.launch()