import sys, os
import json
import string
import traceback
from optparse import OptionParser
import requests
import glob
import random
import urllib3
urllib3.disable_warnings()
#import util
import subprocess
import csv
import time


# The RCSB APIs are often overloaded, so be patient: long timeouts and
# a fixed pause between attempts.
MAX_RETRIES = 6        # attempts per search request
GQL_RETRIES = 8        # attempts per GraphQL request
RETRY_DELAY = 10       # seconds to wait between attempts (fixed)
REQUEST_TIMEOUT = 300  # seconds to wait for a response (search + GraphQL)
PAUSE_BETWEEN_REQUESTS = 2  # seconds, to avoid adding to the load

# Root of the step_1..step_4 working folders (used to clear the previous release)
BASE_DIR = "/data/projects/glygen/downloads/rcsb_pdb/temp"
# List of batches that could not be fetched. Kept OUTSIDE step_1/ so step 2
# never mistakes it for a batch file. Used by:  python3 pdb-download-step-1.py --retry-failed
MANIFEST = BASE_DIR + "/step_1_failed_batches.manifest"
BAD_IDS_FILE = BASE_DIR + "/step_1_bad_ids.tsv"   # written by --diagnose-failed
DIAG_RETRIES = 3   # attempts per single ID in --diagnose-failed mode
LAST_ERROR = [""]  # reason for the most recent failed GraphQL attempt
GRAPHQL_TEMPLATE = "/data/projects/glygen/downloads/rcsb_pdb/queries/step_1_graphql.txt"


def backoff(attempt):
    return RETRY_DELAY


def clear_previous_release():
    """
    Remove the previous release's output files (steps 1-4) before a new run,
    so later steps never pick up stale leftovers (e.g. extra batch files from
    a species that now has fewer entries, or species that failed this time).
    Only the generated .json/.csv files inside the step folders are removed --
    the folders themselves and the SIFTS XML cache are left alone.
    """
    targets = [
        (BASE_DIR + "/step_1", "*.json"),
        (BASE_DIR + "/step_2", "*.csv"),
        (BASE_DIR + "/step_3", "*.json"),
        (BASE_DIR + "/step_4", "*.csv"),
    ]
    for folder, pattern in targets:
        os.makedirs(folder, exist_ok=True)
        removed = 0
        for f in glob.glob(os.path.join(folder, pattern)):
            os.remove(f)
            removed += 1
        print("Cleared %d previous file(s) from %s" % (removed, folder))
    for stale in (MANIFEST, BAD_IDS_FILE):
        if os.path.exists(stale):
            os.remove(stale)


def fetch_with_retry(method, url, **kwargs):
    """Fetch a URL with retries on failure. Returns (status_code, content) or (None, None)."""
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            if method == "post":
                res = requests.post(url, verify=False, timeout=REQUEST_TIMEOUT, **kwargs)
            else:
                res = requests.get(url, verify=False, timeout=REQUEST_TIMEOUT, **kwargs)

            if res.status_code == 200:
                return res.status_code, res.content
            if res.status_code == 204:
                # the RCSB search API answers 204 (empty body) when a query
                # matches nothing -- a valid "no results", not a failure
                return res.status_code, b""
            print(f"    [Attempt {attempt}] HTTP {res.status_code} from {url}")
        except Exception as e:
            print(f"    [Attempt {attempt}] Exception: {e}")

        if attempt < MAX_RETRIES:
            print(f"    Retrying in {backoff(attempt)}s...")
            time.sleep(backoff(attempt))

    return None, None


def fetch_graphql_entries(id_list, template, retries=None):
    """
    One GraphQL request for id_list, retried up to GQL_RETRIES times with
    backoff. Returns the list of entries, or None if every attempt failed.

    A response only counts as good if it is HTTP 200, has no "errors" block
    and returns one non-null entry per requested id.
    """
    q = template % (json.dumps(id_list))
    url_two = "https://data.rcsb.org/graphql?query=" + q

    retries = retries or GQL_RETRIES
    for attempt in range(1, retries + 1):
        try:
            res = requests.get(url_two, verify=False, timeout=REQUEST_TIMEOUT)
            if res.status_code == 200:
                obj = json.loads(res.content)
                entries = (obj.get("data") or {}).get("entries")
                if not obj.get("errors") and entries is not None \
                        and len(entries) == len(id_list) and all(e is not None for e in entries):
                    return entries
                if obj.get("errors"):
                    LAST_ERROR[0] = "GraphQL error: %s" % str(obj["errors"][0].get("message", obj["errors"][0]))[:200]
                elif entries is None or any(e is None for e in entries):
                    LAST_ERROR[0] = "entry returned null"
                else:
                    LAST_ERROR[0] = "wrong number of entries returned"
                print("    [GraphQL %d ids, attempt %d] incomplete/errored response" % (len(id_list), attempt))
            else:
                LAST_ERROR[0] = "HTTP %s" % res.status_code
                print("    [GraphQL %d ids, attempt %d] HTTP %s" % (len(id_list), attempt, res.status_code))
        except Exception as e:
            LAST_ERROR[0] = "Exception: %s" % str(e)[:200]
            print("    [GraphQL %d ids, attempt %d] Exception: %s" % (len(id_list), attempt, e))
        if attempt < retries:
            print("    Waiting %ds before retrying..." % backoff(attempt))
            time.sleep(backoff(attempt))
    return None


def batch_file(tax_name, batch_no, part):
    suffix = ("." + part) if part else ""
    return BASE_DIR + "/step_1/%s.batch.%s%s.json" % (tax_name, batch_no, suffix)


def save_batch(tax_name, batch_no, part, entries):
    """
    A whole batch is saved as <species>.batch.<n>.json (as before).

    Anything recovered after a failure (a half-batch from a split, or IDs
    re-fetched by --retry-failed / --diagnose-failed) is MERGED into a single
    <species>.batch.recovered.json instead of getting a file of its own, so
    failures never multiply the number of files: at most one extra file per
    species. Entries are de-duplicated by PDB id when merged.
    """
    if not part:
        out_file = batch_file(tax_name, batch_no, "")
        with open(out_file, "w") as FW:
            FW.write("%s\n" % (json.dumps({"data": {"entries": entries}}, indent=4)))
        return out_file

    out_file = BASE_DIR + "/step_1/%s.batch.recovered.json" % tax_name
    merged = {}
    if os.path.exists(out_file):
        for e in json.load(open(out_file))["data"]["entries"]:
            merged[e["entry"]["id"]] = e
    for e in entries:
        merged[e["entry"]["id"]] = e
    with open(out_file, "w") as FW:
        FW.write("%s\n" % (json.dumps({"data": {"entries": list(merged.values())}}, indent=4)))
    return out_file


def process_batch(tax_name, batch_no, part, id_list, template, failed):
    """
    Fetch one batch and save it. If the whole batch fails (after GQL_RETRIES
    attempts), it is split in half ONCE and each half is tried on its own
    (again GQL_RETRIES attempts). A half that fails is skipped -- it is NOT
    split again -- and appended to `failed` so it is listed at the end of the
    run and can be re-fetched later with --retry-failed.

    `part` is "" for a whole batch and "a"/"b" for a half; recovered halves are
    merged into one <species>.batch.recovered.json (see save_batch).
    A part that already came from a split is never split again.
    """
    def label(prt):
        return "%s batch %s%s" % (tax_name, batch_no, ("." + prt) if prt else "")

    entries = fetch_graphql_entries(id_list, template)
    if entries is not None:
        out_file = save_batch(tax_name, batch_no, part, entries)
        print("  [%s] Saved %d entries -> %s" % (label(part), len(id_list), out_file))
        return

    if part or len(id_list) < 2:
        # already a half (or too small to split): skip it, do not split again
        print("  [%s] SKIPPED after %d attempts (%d ids)" % (label(part), GQL_RETRIES, len(id_list)))
        failed.append({"species": tax_name, "batch_no": batch_no, "part": part, "ids": id_list})
        return

    mid = len(id_list) // 2
    print("  [%s] failed -- splitting once into %d + %d ids" % (label(part), mid, len(id_list) - mid))
    for sub, ids in (("a", id_list[:mid]), ("b", id_list[mid:])):
        time.sleep(PAUSE_BETWEEN_REQUESTS)
        sub_entries = fetch_graphql_entries(ids, template)
        if sub_entries is not None:
            out_file = save_batch(tax_name, batch_no, sub, sub_entries)
            print("  [%s] Saved %d entries -> %s" % (label(sub), len(ids), out_file))
        else:
            print("  [%s] SKIPPED after %d attempts (%d ids)" % (label(sub), GQL_RETRIES, len(ids)))
            failed.append({"species": tax_name, "batch_no": batch_no, "part": sub, "ids": ids})


def report_failures(failed, failed_species, no_structures=()):
    if no_structures:
        print("\nNote: no structures in RCSB for: %s (not an error; no files written for them)" % ", ".join(no_structures))
    if failed_species:
        print("\n!!! WARNING: the RCSB search itself failed for: %s" % ", ".join(failed_species))
        print("    Their data is incomplete -- re-run the whole of step 1 for a full refresh.")
    if failed:
        with open(MANIFEST, "w") as FW:
            json.dump(failed, FW, indent=2)
        n = sum(len(f["ids"]) for f in failed)
        print("\n!!! WARNING: %d batch(es) (%d PDB IDs) were SKIPPED after retries:" % (len(failed), n))
        for f in failed:
            print("    %s batch %s%s: %d ids (e.g. %s)" % (
                f["species"], f["batch_no"], ("." + f["part"]) if f["part"] else "", len(f["ids"]), ", ".join(f["ids"][:3])))
        print("    Saved to %s" % MANIFEST)
        print("    Re-fetch just these (nothing else is touched) with:  python3 pdb-download-step-1.py --retry-failed")
        print("    To find WHICH PDB IDs fail (fetches them one at a time):  python3 pdb-download-step-1.py --diagnose-failed")
    elif os.path.exists(MANIFEST):
        os.remove(MANIFEST)
    if not failed and not failed_species:
        print("\nAll PDB IDs from the search were fetched from the data API.")


def retry_failed():
    """
    Re-fetch only the batches listed in the manifest. Does NOT clear anything:
    all other step_1 files (and steps 2-4 output) are left exactly as they are.
    A recovered batch is saved under its own file name; batches that still
    fail stay in the manifest for the next attempt. Afterwards re-run steps
    2, 3 and 4 so the recovered entries flow into the final datasets.
    """
    if not os.path.exists(MANIFEST):
        print("No failed batches recorded (%s not found). Nothing to retry." % MANIFEST)
        return
    previous = json.load(open(MANIFEST))
    template = open(GRAPHQL_TEMPLATE, "r").read()
    print("Retrying %d previously failed batch(es)..." % len(previous))

    still_failed = []
    for f in previous:
        process_batch(f["species"], f["batch_no"], f["part"], f["ids"], template, still_failed)
        time.sleep(PAUSE_BETWEEN_REQUESTS)

    subprocess.getoutput("chmod -R 775 %s/step_1/" % BASE_DIR)
    print("\nRecovered %d of %d batch(es)." % (len(previous) - len(still_failed), len(previous)))
    report_failures(still_failed, [])
    if len(still_failed) < len(previous):
        print("Now re-run steps 2, 3 and 4 to include the recovered data.")


def diagnose_failed():
    """
    Take the skipped groups from the manifest and fetch every PDB ID on its
    own, so the IDs that really fail can be named. IDs that fetch fine are
    saved (so they are recovered); IDs that fail are written to
    step_1_bad_ids.tsv with the reason, and stay in the manifest. Nothing
    else is touched. Re-run steps 2, 3 and 4 afterwards.
    """
    if not os.path.exists(MANIFEST):
        print("No failed batches recorded (%s not found). Nothing to diagnose." % MANIFEST)
        return
    previous = json.load(open(MANIFEST))
    template = open(GRAPHQL_TEMPLATE, "r").read()
    total = sum(len(f["ids"]) for f in previous)
    print("Fetching %d PDB IDs one at a time from %d skipped group(s)..." % (total, len(previous)))

    remaining, bad_rows = [], []
    for f in previous:
        good_entries, bad_ids = [], []
        for pid in f["ids"]:
            entries = fetch_graphql_entries([pid], template, retries=DIAG_RETRIES)
            if entries is not None:
                good_entries.extend(entries)
            else:
                bad_ids.append(pid)
                bad_rows.append((f["species"], f["batch_no"], pid, LAST_ERROR[0]))
                print("  BAD ID: %s (%s)" % (pid, LAST_ERROR[0]))
            time.sleep(PAUSE_BETWEEN_REQUESTS)
        if good_entries:
            out_file = save_batch(f["species"], f["batch_no"], f["part"], good_entries)
            print("  %s batch %s: recovered %d of %d IDs -> %s" % (
                f["species"], f["batch_no"], len(good_entries), len(f["ids"]), out_file))
        if bad_ids:
            remaining.append({"species": f["species"], "batch_no": f["batch_no"], "part": f["part"], "ids": bad_ids})

    subprocess.getoutput("chmod -R 775 %s/step_1/" % BASE_DIR)
    with open(BAD_IDS_FILE, "w") as FW:
        FW.write("species\tbatch_no\tpdb_id\treason\n")
        for row in bad_rows:
            FW.write("%s\t%s\t%s\t%s\n" % row)
    if remaining:
        with open(MANIFEST, "w") as FW:
            json.dump(remaining, FW, indent=2)
        print("\n!!! %d PDB ID(s) still fail on their own -- listed with reasons in %s:" % (len(bad_rows), BAD_IDS_FILE))
        for sp, b, pid, why in bad_rows:
            print("    %s (%s batch %s): %s" % (pid, sp, b, why))
    else:
        os.remove(MANIFEST)
        print("\nEvery ID was fetched successfully on its own -- nothing left to retry.")
    if len(bad_rows) < total:
        print("Now re-run steps 2, 3 and 4 to include the recovered data.")


def main():

    if "--retry-failed" in sys.argv:
        retry_failed()
        return
    if "--diagnose-failed" in sys.argv:
        diagnose_failed()
        return

    clear_previous_release()

    url = "https://search.rcsb.org/rcsbsearch/v2/query"

    in_file = "/data/projects/glygen/generated/misc/species_info.csv"
    sp_dict = {}
    with open(in_file, "r") as FR:
        reader = csv.reader(FR, delimiter=",")
        for row in reader:
            if "tax_id" in row:
                continue
            tax_id, short_name, ref_flag = row[0], row[1], row[6]
            if ref_flag.lower() != "yes":
                continue
            sp_dict[tax_id] = short_name

    in_file_one = "/data/projects/glygen/downloads/rcsb_pdb/queries/step_1_rcsbsearch.json"
    req_obj_one = json.load(open(in_file_one))
    template = open(GRAPHQL_TEMPLATE, "r").read()

    tax_id_list = sorted(list(sp_dict.keys()))
    total_species = len(tax_id_list)
    failed_species = []   # species whose search itself failed
    no_structures = []    # species with genuinely no matching structures (not an error)
    failed = []           # skipped batches/halves: {species, batch_no, part, ids}

    for sp_idx, tax_id in enumerate(tax_id_list, 1):
        batch_size = 500
        total_count = 0
        tax_name = sp_dict[tax_id]

        print(f"\n[{sp_idx}/{total_species}] Processing {tax_name} (tax_id={tax_id})")

        if "query" not in req_obj_one or "nodes" not in req_obj_one["query"]:
            print(f"  Skipping: malformed query object")
            continue

        req_obj_one["query"]["nodes"][0]["parameters"]["value"] = tax_id

        species_failed = False
        for i in range(0, 100000):
            start = i * batch_size
            end = (i + 1) * batch_size

            req_obj_one["request_options"]["paginate"]["start"] = start
            req_obj_one["request_options"]["paginate"]["rows"] = batch_size

            # --- Step 1: RCSB search ---
            status, content = fetch_with_retry("post", url, json=req_obj_one)
            if status is None:
                print(f"  [Batch {i+1}] Step 1 failed after {MAX_RETRIES} retries. Skipping species.")
                species_failed = True
                break

            time.sleep(PAUSE_BETWEEN_REQUESTS)
            if status == 204 or not content.strip():
                print(f"  No structures found in RCSB for {tax_name} (tax_id={tax_id}) -- nothing to fetch.")
                no_structures.append(tax_name)
                break
            res_obj = json.loads(content)

            if "total_count" in res_obj and total_count == 0:
                total_count = res_obj["total_count"]
                print(f"  Total entries: {total_count}")

            if "result_set" not in res_obj or not res_obj["result_set"]:
                print(f"  [Batch {i+1}] Empty result_set. Stopping.")
                break

            val_list = [obj["identifier"] for obj in res_obj["result_set"]]

            # --- Step 2: GraphQL fetch (retried; split once on failure) ---
            process_batch(tax_name, i + 1, "", val_list, template, failed)

            if total_count > 0 and end >= total_count:
                print(f"  Completed all batches for {tax_name}.")
                break

        if species_failed:
            print(f"  WARNING: {tax_name} was skipped due to repeated API failures.")
            failed_species.append(tax_name)

    subprocess.getoutput("chmod -R 775 %s/step_1/" % BASE_DIR)
    report_failures(failed, failed_species, no_structures)
    print("\nDone.")


if __name__ == "__main__":
    main()
    print("Script finished successfully.", flush=True)