from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel
from Bio import Entrez
from fpdf import FPDF
import tempfile

# Initialize Clinical AI Agent
clinical_agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=HfApiModel())

# Your PubMed Email (IMPORTANT: Use your registered PubMed email)
Entrez.email = "Pub_Email"  # <-- Your registered PubMed email

def generate_clinical_content(topic, audience):
    """Generates medical content based on topic and audience."""
    prompt = f"Write a detailed medical article on: {topic}.\nTarget Audience: {audience}.\nInclude latest medical research insights."
    return clinical_agent.run(prompt)

def generate_summary(content):
    """Summarizes a given clinical document or research paper."""
    summary_prompt = f"Summarize the following medical research:\n{content}"
    return clinical_agent.run(summary_prompt)

def generate_soap_note(symptoms, history):
    """Generates a SOAP Note for clinicians based on symptoms and patient history."""
    soap_prompt = (
        f"Create a structured SOAP Note for a patient with:\n"
        f"Symptoms: {symptoms}\n"
        f"History: {history}\n"
        f"Include Assessment & Plan."
    )
    return clinical_agent.run(soap_prompt)

def fetch_pubmed_articles(query):
    """Fetches latest medical research from PubMed based on a query."""
    try:
        handle = Entrez.esearch(db="pubmed", term=query, retmax=5)
        record = Entrez.read(handle)
        handle.close()
        article_ids = record["IdList"]

        articles = []
        for article_id in article_ids:
            handle = Entrez.efetch(db="pubmed", id=article_id, retmode="xml")
            article_record = Entrez.read(handle)
            handle.close()

            article_data = article_record["PubmedArticle"][0]["MedlineCitation"]["Article"]
            title = article_data.get("ArticleTitle", "No Title Available")
            abstract = article_data.get("Abstract", {}).get("AbstractText", ["No Abstract Available"])[0]
            authors = ", ".join([author["LastName"] for author in article_data.get("AuthorList", []) if "LastName" in author])
            url = f"https://pubmed.ncbi.nlm.nih.gov/{article_id}/"

            articles.append({"title": title, "abstract": abstract, "authors": authors, "url": url})

        return articles
    except Exception as e:
        return [{"title": "Error Fetching Articles", "abstract": str(e), "authors": "N/A", "url": "#"}]

def save_as_text(content, filename):
    """Saves AI-generated content as a TXT file."""
    with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as tmp_file:
        tmp_file.write(content.encode())
        return tmp_file.name, filename

def save_as_pdf(content, filename):
    """Saves AI-generated content as a PDF file."""
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", size=12)
        pdf.multi_cell(0, 10, content)
        pdf.output(tmp_file.name)
        return tmp_file.name, filename