import sys, os
import json
import string
import traceback
from optparse import OptionParser
import requests
import glob
import random
import time
import gzip
import shutil
import subprocess
from email.utils import formatdate
from collections import Counter
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, as_completed
import urllib3
urllib3.disable_warnings()
#import util

# ---------------------------------------------------------------------------
# This step used to call the RCSB "sequence-coordinates" GraphQL API per
# chain to get PDB <-> UniProt residue-level mappings. It now downloads the
# PDBe/SIFTS per-entry XML files instead and parses the residue-level
# crossRefDb mappings directly, which additionally gives us the *author*
# (auth_seq_id) residue number for free -- something the RCSB API call
# never provided.
#
# SIFTS XML numbering cheat-sheet:
#   <residue dbSource="PDBe" dbResNum="X">                     -> "pdbe" numbering
#                                                                  (equivalent to mmCIF label_seq_id;
#                                                                  same numbering used for
#                                                                  glycosylation_site_pdb in step_2,
#                                                                  since that came from label_seq_id)
#     <crossRefDb dbSource="PDB" dbCoordSys="PDBresnum"
#                 dbChainId="A" dbResNum="Y"/>                  -> "auth" numbering (mmCIF auth_seq_id,
#                                                                  what's in the PDB file / PyMOL)
#     <crossRefDb dbSource="UniProt" dbAccessionId="P.." dbResNum="Z"/>  -> UniProt numbering
# ---------------------------------------------------------------------------

SIFTS_XML_URL = "https://ftp.ebi.ac.uk/pub/databases/msd/sifts/xml/%s.xml.gz"
BASE_DIR = "/data/projects/glygen/downloads/rcsb_pdb/temp"
SIFTS_XML_DIR = "/data/projects/glygen/downloads/rcsb_pdb/sifts_xml"
STEP_2_DIR = BASE_DIR + "/step_2"
STEP_3_OUT_DIR = BASE_DIR + "/step_3"

# The SIFTS/EBI servers can be slow or overloaded, so be patient: long timeouts
# and a fixed pause between attempts.
MAX_RETRIES = 6        # attempts per file before it is skipped
RETRY_DELAY = 10       # seconds to wait between attempts (fixed)
REQUEST_TIMEOUT = 300  # seconds to wait for a response
MAX_WORKERS = 8

# A cached SIFTS file younger than this is trusted as-is: the server isn't even
# asked about it. UniProt (and so SIFTS) is refreshed about every 8 weeks, so
# files older than this are re-checked with the server. 0 = always check.
# Override with:  --max-cache-age-days N
MAX_CACHE_AGE_DAYS = 60

# Optional folder holding OLDER copies of SIFTS files (e.g. a previous download).
# Only used for an entry whose file on the SIFTS server has no residue records
# and which has no usable copy in the cache. Set with:  --backup-dir /path/to/folder
BACKUP_SIFTS_DIR = None

# Files that could not be downloaded; read by --retry-failed. Lives in BASE_DIR
# (not in step_3/) so it is never mistaken for a sifts_map file.
FAILED_FILE = BASE_DIR + "/step_3_failed_sifts.tsv"


def strip_ns(tag):
    """Strip the XML namespace prefix SIFTS puts on every tag."""
    return tag.split("}")[-1] if "}" in tag else tag


def get_pdb_ids_by_taxname(step2_dir=STEP_2_DIR):
    """
    Scan every step_2/<tax_name>.csv file and return:
      tax_pdb_ids : dict  tax_name -> set(pdb_id)
      all_pdb_ids : set(pdb_id) across all taxa (for a single dedup'd download pass)
      tax_needed  : dict  tax_name -> {"chain_keys": set, "pos_keys": set}
          the only SIFTS lookups step_4 will make for this species, so the
          map files can be cut down to just those residues:
            chain_keys "<pdb_id>.<chain>|<label_pos>"  (auth chain and label chain)
            pos_keys   "<pdb_id>|<label_pos>"
    """
    tax_pdb_ids = {}
    tax_needed = {}
    all_pdb_ids = set()
    file_list = glob.glob(os.path.join(step2_dir, "*.csv"))
    for in_file in file_list:
        tax_name = os.path.basename(in_file)[:-4]  # strip ".csv"
        pdb_ids = set()
        chain_keys = set()
        pos_keys = set()
        with open(in_file, "r") as FR:
            header_line = FR.readline().strip()
            step2_header = header_line[1:-1].split("\",\"")
            idx = {name: step2_header.index(name) for name in
                   ("src_xref_id", "amino_acid_chain", "glycosylation_site_pdb", "amino_acid_chain_auth")}
            for line in FR:
                line = line.strip()
                if not line:
                    continue
                row = line[1:-1].split("\",\"")
                pdb_id = row[idx["src_xref_id"]].strip().lower()
                if not pdb_id:
                    continue
                pdb_ids.add(pdb_id)
                pos = row[idx["glycosylation_site_pdb"]]
                for chain in (row[idx["amino_acid_chain_auth"]], row[idx["amino_acid_chain"]]):
                    if chain:
                        chain_keys.add("%s.%s|%s" % (pdb_id, chain, pos))
                pos_keys.add("%s|%s" % (pdb_id, pos))
        tax_pdb_ids[tax_name] = pdb_ids
        tax_needed[tax_name] = {"chain_keys": chain_keys, "pos_keys": pos_keys}
        all_pdb_ids |= pdb_ids
    return tax_pdb_ids, all_pdb_ids, tax_needed


def cached_xml_looks_empty(xml_path, nbytes=65536):
    """
    True if an existing cached SIFTS file has no <residue> records. Such a file
    parses fine but is useless (no mappings), so it must not be trusted -- and
    must not be protected by the "not modified" check below.
    (Residue records start within the first few KB of a real SIFTS file.)
    """
    try:
        with open(xml_path, "rb") as f:
            return b"<residue" not in f.read(nbytes)
    except OSError:
        return True


def restore_from_backup(pdb_id, xml_path):
    """Copy a good (non-empty) older copy from BACKUP_SIFTS_DIR into the cache."""
    if not BACKUP_SIFTS_DIR:
        return False
    src = os.path.join(BACKUP_SIFTS_DIR, "%s.xml" % pdb_id)
    if not os.path.exists(src) or cached_xml_looks_empty(src):
        return False
    tmp = xml_path + ".tmp"
    shutil.copy2(src, tmp)
    # Date the restored copy as older than the cache-age limit (never newer than
    # the backup file itself), so every later run asks the server again and
    # picks up the real file as soon as SIFTS fixes it.
    stale = time.time() - (MAX_CACHE_AGE_DAYS + 1) * 86400
    t = min(os.path.getmtime(src), stale)
    os.utime(tmp, (t, t))
    os.replace(tmp, xml_path)
    return True


def download_and_unzip_sifts(pdb_id, outdir=SIFTS_XML_DIR):
    """
    Download <pdb_id>.xml.gz from the SIFTS mirror and unzip it.

    If a cached .xml already exists we still ask the server whether the file
    has changed since we saved it (HTTP If-Modified-Since, using the cached
    file's modification time):
      - 304 Not Modified -> keep the cached copy, nothing is downloaded
      - 200              -> SIFTS has updated this entry, so replace the copy
    That way each release only re-downloads entries SIFTS actually changed,
    rather than everything (or nothing). If the server ignored the header it
    would answer 200 every time, which just means a full re-download -- safe,
    only slower.

    Returns (pdb_id, status) with status one of:
      "ok" (new file), "updated" (cached file replaced), "unchanged",
      "not_found", or "failed: <reason>".
    """
    os.makedirs(outdir, exist_ok=True)
    xml_path = os.path.join(outdir, "%s.xml" % pdb_id)
    gz_path = os.path.join(outdir, "%s.xml.gz" % pdb_id)
    tmp_path = xml_path + ".tmp"

    headers = {}
    had_cache = os.path.exists(xml_path)
    bad_cache = had_cache and cached_xml_looks_empty(xml_path)
    if bad_cache:
        # cached copy has no residue records: ignore it and download unconditionally
        had_cache = False
    if had_cache and MAX_CACHE_AGE_DAYS > 0:
        age_days = (time.time() - os.path.getmtime(xml_path)) / 86400.0
        if age_days < MAX_CACHE_AGE_DAYS:
            return pdb_id, "fresh_cache"     # recent and has residues: no request needed
    if had_cache:
        headers["If-Modified-Since"] = formatdate(os.path.getmtime(xml_path), usegmt=True)

    url = SIFTS_XML_URL % pdb_id
    last_err = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            res = requests.get(url, headers=headers, timeout=REQUEST_TIMEOUT, verify=False)
            if res.status_code == 304:
                try:
                    os.utime(xml_path, None)   # checked today: start its age again
                except OSError:
                    pass
                return pdb_id, "unchanged"
            if res.status_code == 404:
                return pdb_id, "not_found"
            res.raise_for_status()
            with open(gz_path, "wb") as f:
                f.write(res.content)
            # unzip to a temp file first so a failed/partial unzip can never
            # replace a good cached copy
            with gzip.open(gz_path, "rb") as gz_in, open(tmp_path, "wb") as xml_out:
                shutil.copyfileobj(gz_in, xml_out)
            tree = ET.parse(tmp_path)   # raises if the file is truncated/corrupt -> retried
            if not any(strip_ns(e.tag) == "residue" for e in tree.getroot().iter()):
                # the server's own file has no residue mappings: not a download
                # problem, so don't retry or overwrite anything -- just report it
                os.remove(tmp_path)
                os.remove(gz_path)
                if had_cache:
                    # a good cached copy exists: never replace it with an empty file
                    return pdb_id, "kept_old_copy"
                if restore_from_backup(pdb_id, xml_path):
                    return pdb_id, "restored_from_backup"
                return pdb_id, "empty_in_sifts"
            os.replace(tmp_path, xml_path)
            os.remove(gz_path)
            if bad_cache:
                return pdb_id, "repaired"
            return pdb_id, ("updated" if had_cache else "ok")
        except Exception as e:
            last_err = e
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
    for leftover in (gz_path, tmp_path):   # never leave partial files behind
        if os.path.exists(leftover):
            os.remove(leftover)
    return pdb_id, "failed after %d attempts: %s" % (MAX_RETRIES, str(last_err)[:150])


def parse_sifts_xml(xml_path):
    """
    Parse one SIFTS per-entry XML file.

    Returns:
      dict keyed by (chain_id, pdbe_resnum) -> list of
        {"auth_resnum": ..., "uniprot_ac": ..., "uniprot_resnum": ...}

    A residue can map to more than one chain (homo-oligomers sharing one
    entity) and, rarely, to more than one UniProt accession, so both are
    captured as lists rather than assumed to be 1:1.
    """
    mapping = {}
    try:
        tree = ET.parse(xml_path)
    except ET.ParseError:
        return mapping
    root = tree.getroot()

    for entity in root:
        if strip_ns(entity.tag) != "entity":
            continue
        ent_res = {}       # pdbe_resnum (int) -> (pdb_aa, {uniprot_ac: uniprot_aa}); ALL residues,
                           # observed or not, so neighbours can be judged correctly
        ent_entries = []   # (pdbe_resnum int, entry dict) created for this entity
        for segment in entity:
            if strip_ns(segment.tag) != "segment":
                continue
            for list_residue in segment:
                if strip_ns(list_residue.tag) != "listResidue":
                    continue
                for residue in list_residue:
                    if strip_ns(residue.tag) != "residue":
                        continue
                    pdbe_resnum = residue.get("dbResNum")
                    pdb_aa = (residue.get("dbResName") or "").upper()   # 3-letter, e.g. ASN

                    auth_pairs = []      # (chain_id, auth_resnum)
                    uniprot_pairs = []   # (uniprot_ac, uniprot_resnum)

                    for cross in residue:
                        if strip_ns(cross.tag) != "crossRefDb":
                            continue
                        src = cross.get("dbSource")
                        if src == "PDB" and cross.get("dbCoordSys") == "PDBresnum":
                            auth_pairs.append((cross.get("dbChainId"), cross.get("dbResNum")))
                        elif src == "UniProt":
                            uniprot_pairs.append((cross.get("dbAccessionId"), cross.get("dbResNum"),
                                                  (cross.get("dbResName") or "").upper()))   # 1-letter, e.g. N

                    try:
                        ent_res[int(pdbe_resnum)] = (pdb_aa, {u[0]: u[2] for u in uniprot_pairs})
                    except (TypeError, ValueError):
                        pass

                    for chain_id, auth_resnum in auth_pairs:
                        if chain_id is None or pdbe_resnum is None:
                            continue
                        key = (chain_id, pdbe_resnum)
                        if uniprot_pairs:
                            for uniprot_ac, uniprot_resnum, uniprot_aa in uniprot_pairs:
                                ent = {
                                    "auth_resnum": auth_resnum,
                                    "auth_chain_id": chain_id,
                                    "uniprot_ac": uniprot_ac,
                                    "uniprot_resnum": uniprot_resnum,
                                    "pdb_aa": pdb_aa,
                                    "uniprot_aa": uniprot_aa,
                                }
                                mapping.setdefault(key, []).append(ent)
                                ent_entries.append((pdbe_resnum, ent))
                        else:
                            # residue exists in the structure but isn't covered
                            # by a UniProt segment (e.g. tags, unmapped termini)
                            mapping.setdefault(key, []).append({
                                "auth_resnum": auth_resnum,
                                "auth_chain_id": chain_id,
                                "uniprot_ac": None,
                                "uniprot_resnum": None,
                                "pdb_aa": pdb_aa,
                                "uniprot_aa": "",
                            })
        add_neighbour_info(ent_res, ent_entries)
    return mapping


def add_neighbour_info(ent_res, ent_entries):
    """
    For every residue mapped to UniProt, record the next two residues along the
    entity (n1_*, n2_*): the PDB residue name, the UniProt residue, and a state:
      mapped    the neighbour is mapped to the SAME UniProt accession
      unmapped  the neighbour exists in the entity but is not mapped to that
                accession (tag, linker, insertion, other protein)
      end       there is no residue there (end of the chain)
    Unobserved residues count, because the table is built from every residue
    SIFTS lists for the entity, not only those with coordinates.
    For an N-linked site n1 is the X of N-X-S/T and n2 is the S/T slot.
    """
    for pdbe_resnum, e in ent_entries:
        try:
            base = int(pdbe_resnum)
        except (TypeError, ValueError):
            continue
        for k in (1, 2):
            nb = ent_res.get(base + k)
            if nb is None:
                e["n%d_state" % k], e["n%d_pdb_aa" % k], e["n%d_uniprot_aa" % k] = "end", "", ""
            elif e["uniprot_ac"] in nb[1]:
                e["n%d_state" % k], e["n%d_pdb_aa" % k], e["n%d_uniprot_aa" % k] = "mapped", nb[0], nb[1][e["uniprot_ac"]]
            else:
                e["n%d_state" % k], e["n%d_pdb_aa" % k], e["n%d_uniprot_aa" % k] = "unmapped", nb[0], ""


def build_sifts_map_for_taxon(pdb_ids, parsed_cache, needed=None):
    """
    Combine per-entry mappings into two indices for a tax_name's PDB IDs:

      by_chain_pos : "<pdb_id>.<auth_chain>|<pdbe_resnum>" -> list of entries
          Exact match. Works whenever step_2's label-based chain code
          happens to equal SIFTS' auth chain code (the common case).

      by_pos_only  : "<pdb_id>|<pdbe_resnum>" -> list of entries (all chain
          copies at that position, across every auth chain of the entity)
          Fallback for when label_asym_id (step_2's "amino_acid_chain")
          differs from auth_asym_id (SIFTS' dbChainId) -- SIFTS' classic
          XML has no field for label_asym_id, so this is as close as we can
          get. step_4 only uses this fallback when every candidate at that
          position agrees on the UniProt residue, to avoid guessing between
          genuinely different chain copies of a homo-oligomer.
    """
    # If `needed` is given, only residues step_4 will actually look up are kept
    # (every glycosylated residue from step_2), which keeps the map files small.
    # Lookups return exactly the same results as with the full map: by_pos_only
    # still keeps ALL chain copies at a needed position, so step_4's
    # "do all candidates agree?" check is unchanged.
    by_chain_pos = {}
    by_pos_only = {}
    for pdb_id in pdb_ids:
        if pdb_id not in parsed_cache:
            xml_path = os.path.join(SIFTS_XML_DIR, "%s.xml" % pdb_id)
            if not os.path.exists(xml_path):
                parsed_cache[pdb_id] = {}
                continue
            parsed_cache[pdb_id] = parse_sifts_xml(xml_path)

        entry_map = parsed_cache[pdb_id]
        for (chain_id, pdbe_resnum), entries in entry_map.items():
            chain_key = "%s.%s|%s" % (pdb_id, chain_id, pdbe_resnum)
            pos_key = "%s|%s" % (pdb_id, pdbe_resnum)
            if needed is not None:
                if chain_key in needed["chain_keys"]:
                    by_chain_pos.setdefault(chain_key, []).extend(entries)
                if pos_key in needed["pos_keys"]:
                    by_pos_only.setdefault(pos_key, []).extend(entries)
                continue
            by_chain_pos.setdefault(chain_key, []).extend(entries)
            by_pos_only.setdefault(pos_key, []).extend(entries)

    return {"by_chain_pos": by_chain_pos, "by_pos_only": by_pos_only}


def download_all(pdb_ids):
    """Download/refresh every id in pdb_ids. Returns (status_counts, problems)."""
    problems = []   # (pdb_id, status)
    status_counts = Counter()
    pdb_ids = list(pdb_ids)
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(download_and_unzip_sifts, pdb_id): pdb_id for pdb_id in pdb_ids}
        for i, future in enumerate(as_completed(futures), 1):
            pdb_id, status = future.result()
            status_counts[status.split(":")[0].split(" after")[0]] += 1
            if status.startswith("failed") or status in ("not_found", "empty_in_sifts"):
                problems.append((pdb_id, status))
            if i % 200 == 0 or i == len(pdb_ids):
                print("  ... %d/%d processed" % (i, len(pdb_ids)))
    return status_counts, problems


def report_problems(problems, tax_pdb_ids):
    """Print and save the files that could not be downloaded."""
    if os.path.exists(FAILED_FILE):
        os.remove(FAILED_FILE)
    if not problems:
        print("\nAll SIFTS XML files were downloaded/checked successfully.")
        return
    problems = sorted(problems)
    with open(FAILED_FILE, "w") as FW:
        FW.write("pdb_id\tstatus\told_cached_copy_used\tspecies\n")
        for pdb_id, status in problems:
            cached = os.path.exists(os.path.join(SIFTS_XML_DIR, "%s.xml" % pdb_id))
            species = ",".join(sorted(t for t, ids in tax_pdb_ids.items() if pdb_id in ids))
            FW.write("%s\t%s\t%s\t%s\n" % (pdb_id, status, "yes" if cached else "no", species))

    failed = [x for x in problems if x[1].startswith("failed")]
    missing = [x for x in problems if x[1] == "not_found"]
    print("\n!!! WARNING: %d SIFTS XML file(s) could not be downloaded (full list: %s)" % (len(problems), FAILED_FILE))
    if failed:
        print("  %d failed after %d attempts (server/network problem):" % (len(failed), MAX_RETRIES))
        for pdb_id, status in failed[:30]:
            cached = os.path.exists(os.path.join(SIFTS_XML_DIR, "%s.xml" % pdb_id))
            print("    %s: %s%s" % (pdb_id, status, "  [older cached copy used]" if cached else "  [NO copy -- sites from this entry will be dropped]"))
        if len(failed) > 30:
            print("    ... and %d more (see %s)" % (len(failed) - 30, FAILED_FILE))
    empty = [x for x in problems if x[1] == "empty_in_sifts"]
    if empty:
        print("  %d have NO residue mappings in the SIFTS file on the server itself (nothing to map): %s%s" % (
            len(empty), ", ".join(x[0] for x in empty[:30]), " ..." if len(empty) > 30 else ""))
    if missing:
        print("  %d not found on the SIFTS server (404 -- often new/obsolete entries): %s%s" % (
            len(missing), ", ".join(x[0] for x in missing[:30]), " ..." if len(missing) > 30 else ""))
    if failed:
        print("  Re-fetch just the failed ones (rebuilds the maps afterwards) with:  python3 pdb-download-step-3.py --retry-failed")


def build_maps(tax_pdb_ids, tax_needed):
    os.makedirs(STEP_3_OUT_DIR, exist_ok=True)
    parsed_cache = {}
    for tax_name, pdb_ids in tax_pdb_ids.items():
        print("Building SIFTS residue map for %s (%d PDB entries)..." % (tax_name, len(pdb_ids)))
        sifts_map = build_sifts_map_for_taxon(pdb_ids, parsed_cache, tax_needed.get(tax_name))
        out_file = os.path.join(STEP_3_OUT_DIR, "%s.sifts_map.json" % tax_name)
        with open(out_file, "w") as FW:
            json.dump(sifts_map, FW)
        print("  -> %s (%d chain-position entries, %d position entries; %.1f MB)" % (
            out_file, len(sifts_map["by_chain_pos"]), len(sifts_map["by_pos_only"]),
            os.path.getsize(out_file) / 1e6))
    subprocess.getoutput("chmod -R 775 %s/" % STEP_3_OUT_DIR)


def main():

    global BACKUP_SIFTS_DIR, MAX_CACHE_AGE_DAYS
    if "--max-cache-age-days" in sys.argv:
        MAX_CACHE_AGE_DAYS = float(sys.argv[sys.argv.index("--max-cache-age-days") + 1])
    print("Cached SIFTS files younger than %s days are used without asking the server (0 = always check)." % MAX_CACHE_AGE_DAYS)
    if "--backup-dir" in sys.argv:
        BACKUP_SIFTS_DIR = sys.argv[sys.argv.index("--backup-dir") + 1]
        print("Backup SIFTS folder for entries the server has no mappings for: %s" % BACKUP_SIFTS_DIR)

    tax_pdb_ids, all_pdb_ids, tax_needed = get_pdb_ids_by_taxname()
    print("Found %d unique PDB IDs across %d taxa." % (len(all_pdb_ids), len(tax_pdb_ids)))

    if "--retry-failed" in sys.argv:
        if not os.path.exists(FAILED_FILE):
            print("No failed downloads recorded (%s not found). Nothing to retry." % FAILED_FILE)
            return
        with open(FAILED_FILE) as FR:
            next(FR)
            retry_ids = [line.split("\t")[0] for line in FR if line.strip() and "\tfailed" in line]
        print("Retrying %d previously failed SIFTS download(s)..." % len(retry_ids))
        status_counts, problems = download_all(retry_ids)
        print("SIFTS XML retry summary: %s" % dict(status_counts))
        # keep earlier 404s on the list; only the retried ids are re-evaluated
        old_missing = []
        with open(FAILED_FILE) as FR:
            next(FR)
            for line in FR:
                parts = line.rstrip("\n").split("\t")
                if len(parts) > 1 and parts[1] == "not_found":
                    old_missing.append((parts[0], "not_found"))
        report_problems(problems + old_missing, tax_pdb_ids)
        build_maps(tax_pdb_ids, tax_needed)
        print("Now re-run step 4 to include the recovered data.")
        return

    print("Checking SIFTS XML files (cached entries are only re-downloaded if SIFTS has changed them)...")
    status_counts, problems = download_all(all_pdb_ids)
    print("SIFTS XML summary: %s" % dict(status_counts))
    if status_counts.get("kept_old_copy") or status_counts.get("restored_from_backup"):
        print("NOTE: %d entries use an OLDER copy of the SIFTS file (the server's file had no residue mappings): "
              "%d kept from the cache, %d restored from the backup folder." % (
                  status_counts.get("kept_old_copy", 0) + status_counts.get("restored_from_backup", 0),
                  status_counts.get("kept_old_copy", 0), status_counts.get("restored_from_backup", 0)))

    build_maps(tax_pdb_ids, tax_needed)

    # reported last so it is the final thing in the log
    report_problems(problems, tax_pdb_ids)

    return


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