import gradio as gr
import requests
import time
import re
from Bio import Entrez, SeqIO
from io import StringIO

Entrez.email = "anonymous@pathogenagent.ai"
Entrez.tool = "PathogenAgent-Lite"

HEADERS = {"User-Agent": "PathogenAgent-Lite/1.0"}

def search_pubmed(query, retmax=3):
    try:
        time.sleep(0.3)
        url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
        params = {"db": "pubmed", "term": query, "retmax": retmax, "retmode": "json"}
        resp = requests.get(url, params=params, headers=HEADERS, timeout=10)
        ids = resp.json().get("esearchresult", {}).get("idlist", [])
        if not ids:
            return []
        url2 = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
        params2 = {"db": "pubmed", "id": ",".join(ids), "retmode": "json"}
        resp2 = requests.get(url2, params=params2, headers=HEADERS, timeout=10)
        data = resp2.json()
        results = []
        for uid, rec in data.get("result", {}).items():
            if uid == "uids":
                continue
            title = rec.get("title", "")
            if title:
                results.append(f"- {title}")
        return results
    except Exception:
        return []

def search_genbank(query):
    try:
        time.sleep(0.5)
        # بهبود جستجو: حذف کلمه "genome" و جستجوی نام موجودیت
        search_term = query.replace("genome", "").strip()
        if not search_term:
            search_term = query
        
        # جستجو با نام موجودیت
        handle = Entrez.esearch(db="nucleotide", term=f"{search_term}[Organism]", retmax=3)
        record = Entrez.read(handle)
        handle.close()
        
        if not record.get("IdList"):
            return None
        
        # گرفتن اولین نتیجه
        seq_id = record["IdList"][0]
        handle = Entrez.efetch(db="nucleotide", id=seq_id, rettype="fasta", retmode="text")
        seq = SeqIO.read(StringIO(handle.read()), "fasta")
        return {
            "id": seq.id,
            "length": len(seq.seq),
            "snippet": str(seq.seq[:200])
        }
    except Exception:
        return None

def search_clinvar(query):
    try:
        time.sleep(0.3)
        url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
        # برای ClinVar، فقط در صورتی که query شبیه واریانت باشد جستجو کن
        if not re.search(r"c\.[0-9]+|rs[0-9]+|del|dup|ins|[A-Z][0-9]+[A-Z]", query):
            return []
        params = {"db": "clinvar", "term": query, "retmax": 3, "retmode": "json"}
        resp = requests.get(url, params=params, headers=HEADERS, timeout=10)
        ids = resp.json().get("esearchresult", {}).get("idlist", [])
        if not ids:
            return []
        url2 = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
        params2 = {"db": "clinvar", "id": ",".join(ids), "retmode": "json"}
        resp2 = requests.get(url2, params=params2, headers=HEADERS, timeout=10)
        data = resp2.json()
        results = []
        for uid, rec in data.get("result", {}).items():
            if uid == "uids":
                continue
            title = rec.get("title", "")
            if title:
                results.append(f"- {title}")
        return results
    except Exception:
        return []

def route_query(query):
    q = query.lower()
    if re.search(r"c\.[0-9]+|rs[0-9]+|del|dup|ins|[A-Z][0-9]+[A-Z]", q):
        return "variant"
    if re.search(r"genome|sequence|virus|organism|strain|cov|sars|ebola|influenza|mers|hiv|dengue|zika", q):
        return "pathogen"
    return "general"

def pathogen_agent(query):
    route = route_query(query)
    pubmed = search_pubmed(query)
    clinvar = search_clinvar(query) if route == "variant" else []
    genbank = search_genbank(query) if route == "pathogen" else None
    
    pubmed_text = "\n".join(pubmed) if pubmed else "No PubMed results."
    clinvar_text = "\n".join(clinvar) if clinvar else "No ClinVar results."
    
    genbank_text = "Not applicable."
    if genbank:
        genbank_text = f"ID: {genbank['id']}\nLength: {genbank['length']:,} bp\n{genbank['snippet']}..."

    return f"""
# 🧬 PathogenAgent (Lite - Hugging Face)

**Query:** {query}
**Type:** {route}

---

## 📄 PubMed
{pubmed_text}

---

## 🧪 ClinVar
{clinvar_text}

---

## 🧬 GenBank
{genbank_text}

---

Generated by PathogenAgent
"""

def run(query):
    if not query.strip():
        return "Please enter a query."
    return pathogen_agent(query)

demo = gr.Interface(
    fn=run,
    inputs=gr.Textbox(label="Query", placeholder="e.g., SARS-CoV-2 genome or CFTR F508del"),
    outputs=gr.Markdown(),
    title="PathogenAgent (Hugging Face)",
    description="Agentic AI for Pathogen Genomics"
)

demo.launch()