# ==========================================================
# PATHOGENAGENT - MODULE 2: TOOL EXECUTOR
# ==========================================================

import requests
import time
from typing import List, Dict

class ToolExecutor:
    """اجرای ابزارها و بازیابی شواهد از NCBI"""
    
    def __init__(self):
        self.headers = {"User-Agent": "PathogenAgent/1.0"}
        self.last_request_time = 0
    
    def _wait_for_rate_limit(self):
        """Rate Limiting (حداکثر ۳ درخواست در ثانیه)"""
        current_time = time.time()
        if current_time - self.last_request_time < 0.333:
            time.sleep(0.333 - (current_time - self.last_request_time))
        self.last_request_time = time.time()
    
    def execute(self, query: str, tools: List[str]) -> Dict:
        """اجرای ابزارهای درخواستی"""
        results = {}
        
        if "PubMed" in tools:
            results["PubMed"] = self._search_pubmed(query)
        
        if "ClinVar" in tools:
            results["ClinVar"] = self._search_clinvar(query)
        
        if "GenBank" in tools:
            results["GenBank"] = self._search_genbank(query)
        
        return results
    
    def _search_pubmed(self, query: str) -> List[Dict]:
        """جستجوی PubMed"""
        try:
            self._wait_for_rate_limit()
            
            url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
            params = {"db": "pubmed", "term": query, "retmax": 3, "retmode": "json"}
            resp = requests.get(url, params=params, headers=self.headers, timeout=10)
            ids = resp.json().get("esearchresult", {}).get("idlist", [])
            
            if not ids:
                return []
            
            self._wait_for_rate_limit()
            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=self.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({
                        "title": title,
                        "source": "PubMed",
                        "identifier": f"PMID:{uid}"
                    })
            
            return results
            
        except Exception as e:
            print(f"⚠️ PubMed Error: {e}")
            return []
    
    def _search_clinvar(self, query: str) -> List[Dict]:
        """جستجوی ClinVar"""
        try:
            self._wait_for_rate_limit()
            
            url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
            params = {"db": "clinvar", "term": query, "retmax": 3, "retmode": "json"}
            resp = requests.get(url, params=params, headers=self.headers, timeout=10)
            ids = resp.json().get("esearchresult", {}).get("idlist", [])
            
            if not ids:
                return []
            
            self._wait_for_rate_limit()
            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=self.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({
                        "title": title,
                        "source": "ClinVar",
                        "identifier": f"RCV:{uid}"
                    })
            
            return results
            
        except Exception as e:
            print(f"⚠️ ClinVar Error: {e}")
            return []
    
    def _search_genbank(self, query: str) -> List[Dict]:
        """جستجوی GenBank"""
        try:
            self._wait_for_rate_limit()
            
            url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
            params = {"db": "nucleotide", "term": f"{query}[Organism]", "retmax": 1, "retmode": "json"}
            resp = requests.get(url, params=params, headers=self.headers, timeout=10)
            ids = resp.json().get("esearchresult", {}).get("idlist", [])
            
            if not ids:
                return []
            
            return [{
                "title": f"{query} genome",
                "source": "GenBank",
                "identifier": f"GenBank:{ids[0]}"
            }]
            
        except Exception as e:
            print(f"⚠️ GenBank Error: {e}")
            return []