import gradio as gr
from Bio import Entrez
import os  # For environment variables and file paths
from components import federated_learning

# ---------------------------- Configuration ----------------------------
ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "ENTREZ_EMAIL") # Use environment variable, default fallback
HUGGINGFACE_API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN", "HUGGINGFACE_API_TOKEN") # Use environment variable, default fallback

# ---------------------------- Helper Functions ----------------------------

def log_error(message: str):
    """Logs an error message to the console and a file (if possible)."""
    print(f"ERROR: {message}")
    try:
        with open("error_log.txt", "a") as f:
            f.write(f"{message}\n")
    except:
        print("Couldn't write to error log file.")  #If logging fails, still print to console

# ---------------------------- Tool Functions ----------------------------

def search_pubmed(query: str) -> list:
    """Searches PubMed and returns a list of article IDs."""
    try:
        Entrez.email = ENTREZ_EMAIL
        handle = Entrez.esearch(db="pubmed", term=query, retmax="5")
        record = Entrez.read(handle)
        handle.close()
        return record["IdList"]
    except Exception as e:
        log_error(f"PubMed search error: {e}")
        return [f"Error during PubMed search: {e}"]

def fetch_abstract(article_id: str) -> str:
    """Fetches the abstract for a given PubMed article ID."""
    try:
        Entrez.email = ENTREZ_EMAIL
        handle = Entrez.efetch(db="pubmed", id=article_id, rettype="abstract", retmode="text")
        abstract = handle.read()
        handle.close()
        return abstract
    except Exception as e:
        log_error(f"Error fetching abstract for {article_id}: {e}")
        return f"Error fetching abstract for {article_id}: {e}"

# ---------------------------- Agent Function ----------------------------

def medai_agent(query: str) -> str:
    """Orchestrates the medical literature review and presents abstract."""
    article_ids = search_pubmed(query)

    if isinstance(article_ids, list) and article_ids:
        results = []
        for article_id in article_ids:
            abstract = fetch_abstract(article_id)
            if "Error" not in abstract:
                results.append(f"<div class='article'>\n"
                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
                               f"  <p class='abstract'><strong>Abstract:</strong> {abstract}</p>\n"
                               f"</div>\n")
            else:
                results.append(f"<div class='article error'>\n"
                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
                               f"  <p class='error-message'>Error processing article: {abstract}</p>\n"
                               f"</div>\n")
        return "\n".join(results)
    else:
        return f"No articles found or error occurred: {article_ids}"

# ---------------------------- Gradio Interface ----------------------------

def launch_gradio():
    """Launches the Gradio interface."""

    # CSS to style the article output
    css = """
    .article {
      border: 1px solid #ddd;
      margin-bottom: 10px;
      padding: 10px;
      border-radius: 5px;
    }
    .article.error {
      border-color: #f00;
    }
    .article-id {
      font-size: 1.2em;
      margin-bottom: 5px;
    }
    .abstract {
      font-style: italic;
    }
    .error-message {
      color: #f00;
    }
    """

    with gr.Blocks(css=css) as iface:
        gr.Markdown("# MedAI: Medical Literature Review")
        gr.Markdown("Enter a medical query to retrieve abstracts from PubMed.")

        query_input = gr.Textbox(lines=3, placeholder="Enter your medical query to get abstract from PubMed.")
        submit_button = gr.Button("Submit")
        output_results = gr.HTML()  # Use HTML for formatted output
        federated_learning_output = gr.HTML()

        # Get data
        submit_button.click(medai_agent, inputs=query_input, outputs=output_results)
        run_fl_button = gr.Button("Run Federated Learning (Conceptual)")
        run_fl_button.click(federated_learning.run_federated_learning, outputs = federated_learning_output)

    iface.launch()

# ---------------------------- Main Execution ----------------------------

if __name__ == "__main__":
    launch_gradio()