import time
import requests

BASE = "https://rest.uniprot.org"

def search_protein(query: str, fmt: str = "json", size: int = 5):
    """
    Search UniProtKB using Lucene-style query syntax.
    """
    url = f"{BASE}/uniprotkb/search"
    params = {"query": query, "format": fmt, "size": size}
    r = requests.get(url, params=params, timeout=30)
    r.raise_for_status()
    return r.json() if fmt == "json" else r.text

def retrieve_entry(accession: str, fmt: str = "json"):
    """
    Retrieve a UniProtKB entry by accession.
    """
    url = f"{BASE}/uniprotkb/{accession}"
    params = {"format": fmt}
    r = requests.get(url, params=params, timeout=30)
    r.raise_for_status()
    return r.json() if fmt == "json" else r.text

def id_mapping(from_db: str, to_db: str, ids, poll_interval_s: float = 1.0):
    """
    Map identifiers using UniProt ID Mapping.
    ids can be a list of strings or a comma-separated string.
    """
    if isinstance(ids, (list, tuple)):
        ids = ",".join(ids)

    # 1) Submit mapping job
    submit_url = f"{BASE}/idmapping/run"
    r = requests.post(
        submit_url,
        data={"from": from_db, "to": to_db, "ids": ids},
        timeout=30,
    )
    r.raise_for_status()
    job_id = r.json()["jobId"]

    # 2) Poll job status
    status_url = f"{BASE}/idmapping/status/{job_id}"
    while True:
        s = requests.get(status_url, timeout=30)
        s.raise_for_status()
        payload = s.json()
        if payload.get("jobStatus") in (None, "FINISHED"):
            break
        if payload.get("jobStatus") == "FAILED":
            raise RuntimeError(f"ID mapping failed: {payload}")
        time.sleep(poll_interval_s)

    # 3) Fetch results (JSON)
    results_url = f"{BASE}/idmapping/results/{job_id}"
    res = requests.get(results_url, params={"format": "json"}, timeout=30)
    res.raise_for_status()
    return res.json()

if __name__ == "__main__":
    # Search example: human BRCA1
    search = search_protein("gene:BRCA1 AND organism_id:9606", size=3)
    print("Search results (first accessions):",
          [item["primaryAccession"] for item in search.get("results", [])])

    # Retrieve entry example
    entry = retrieve_entry("P38398")  # UniProt accession for human BRCA1 (example)
    print("Entry primaryAccession:", entry.get("primaryAccession"))
    print("Protein name:", entry.get("proteinDescription", {}).get("recommendedName", {}).get("fullName", {}).get("value"))

    # ID mapping example: gene name -> UniProtKB
    mapping = id_mapping(from_db="Gene_Name", to_db="UniProtKB", ids=["BRCA1"])
    print("Mapping results keys:", mapping.keys())