"""
Read-path .bz2-fallback tests for common/pipeline.download_atom + the
savvy-provisioner completeness gate (.npy <-> .npy.bz2 / .pt <-> .pt.bz2).

Runs inside the chunker/extractor image (boto3 + numpy + redis + bz2 + pbzip2
present). Drives the FOLDER backend (SAVVY_LOCAL_DIR) so no live Garage is
needed; the GLOBAL side is stubbed to "absent" so a folder miss is a genuine
miss. Exercises:
  (i)   a raw .npy present                 -> download_atom returns its bytes
  (ii)  only .npy.bz2 present (pbzip2)      -> decompressed transparently
  (ii') only .npy.bz2 present (python bz2)  -> decompressed transparently
  (iii) neither present                     -> raises NOT-FOUND, as before
  (iv)  np.load round-trips through the fallback (chunker/extractor inherit it)
  (v)   provisioner gate_complete COMPLETE when the .npy exists as .npy.bz2
        + _present_in_global accepts .npy.bz2 / .pt.bz2

Exit 0 on all-pass, non-zero (with a traceback) on any failure.
"""
import bz2
import io
import os
import subprocess
import sys
import tempfile

# ---------------------------------------------------------------------------
# Force the FOLDER backend BEFORE importing pipeline (module-level constants).
# ---------------------------------------------------------------------------
SAVVY_ROOT = tempfile.mkdtemp(prefix="savvy_root_")
os.environ["SAVVY_LOCAL_DIR"] = SAVVY_ROOT
# Distinct global so SAVVY_ENABLED is True even without it (folder mode flips it).
os.environ.setdefault("GLOBAL_S3_URL", "http://global.invalid:3900")

sys.path.insert(0, "/app")            # chunker/extractor image layout
import numpy as np                    # noqa: E402
from common import pipeline           # noqa: E402

assert pipeline.SAVVY_ENABLED, "folder backend should set SAVVY_ENABLED"
assert isinstance(pipeline.get_primary_store(), pipeline._FolderStore)

# ---------------------------------------------------------------------------
# Stub the GLOBAL side to ALWAYS miss: a folder-store miss must surface as the
# genuine NOT-FOUND (no live Garage in the test). _copy_in_object returns False
# (absent in global) and the global stream raises a not-found error.
# ---------------------------------------------------------------------------
class _StubNotFound(Exception):
    response = {"Error": {"Code": "NoSuchKey"}}

class _GlobalAbsent:
    def get_object(self, *a, **k):
        raise _StubNotFound()
    def download_fileobj(self, *a, **k):
        raise _StubNotFound()
    def head_object(self, *a, **k):
        raise _StubNotFound()

pipeline.get_global_s3 = lambda: _GlobalAbsent()
pipeline._copy_in_object = lambda key: False          # nothing copies in from global

failures = []

def check(name, cond):
    print(("PASS" if cond else "FAIL"), name)
    if not cond:
        failures.append(name)

# ---------------------------------------------------------------------------
# Build a real .npy and its .bz2 sibling on the folder store.
# ---------------------------------------------------------------------------
KEY = "pipeline/ctx0/doc0/stage.hidden.json.npy"
arr = np.arange(4096, dtype=np.float32).reshape(64, 64) * 1.5
buf = io.BytesIO()
np.save(buf, arr, allow_pickle=False)
RAW_BYTES = buf.getvalue()
store = pipeline.get_primary_store()

def put_raw():
    store.put_bytes(KEY, RAW_BYTES)

def put_bz2(pbzip2=True):
    """Write <KEY>.bz2 with the compressor's exact tool (pbzip2) or python bz2."""
    comp = bz2.compress(RAW_BYTES) if not pbzip2 else _pbzip2_compress(RAW_BYTES)
    store.put_bytes(KEY + ".bz2", comp)

def _pbzip2_compress(data: bytes) -> bytes:
    p = subprocess.run(["pbzip2", "-9", "-c"], input=data,
                       stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
    return p.stdout

def clear():
    for k in (KEY, KEY + ".bz2"):
        try:
            store.delete(k)
        except Exception:
            pass

# ---------------------------------------------------------------------------
# (i) raw .npy present -> exact bytes via the fast path.
# ---------------------------------------------------------------------------
clear(); put_raw()
out = io.BytesIO()
pipeline.download_atom(KEY, out)
check("(i) raw .npy -> exact bytes", out.getvalue() == RAW_BYTES)
out.seek(0)
check("(i) raw .npy -> np.load round-trip",
      np.array_equal(np.load(out, allow_pickle=False), arr))

# ---------------------------------------------------------------------------
# (ii) only .npy.bz2 present (pbzip2) -> decompressed transparently.
# ---------------------------------------------------------------------------
clear(); put_bz2(pbzip2=True)
out = io.BytesIO()
pipeline.download_atom(KEY, out)               # pbzip2 -d path (binary present)
check("(ii) .npy.bz2 (pbzip2) -> exact decompressed bytes", out.getvalue() == RAW_BYTES)
out.seek(0)
check("(ii) .npy.bz2 (pbzip2) -> np.load round-trip",
      np.array_equal(np.load(out, allow_pickle=False), arr))

# ---------------------------------------------------------------------------
# (ii') only .npy.bz2 present, FORCE pure-python bz2 path (PBZIP2_BIN disabled).
# ---------------------------------------------------------------------------
clear(); put_bz2(pbzip2=False)
_saved = pipeline.PBZIP2_BIN
pipeline.PBZIP2_BIN = ""                        # force bz2.BZ2File streaming reader
try:
    out = io.BytesIO()
    pipeline.download_atom(KEY, out)
    check("(ii') .npy.bz2 (python bz2) -> exact decompressed bytes",
          out.getvalue() == RAW_BYTES)
finally:
    pipeline.PBZIP2_BIN = _saved

# ---------------------------------------------------------------------------
# (iii) neither present -> raises a NOT-FOUND (genuine miss), as before.
# ---------------------------------------------------------------------------
clear()
raised = None
try:
    pipeline.download_atom(KEY, io.BytesIO())
except Exception as e:
    raised = e
check("(iii) neither present -> raises", raised is not None)
check("(iii) the raise is a NOT-FOUND miss", raised is not None and pipeline._is_not_found(raised))

# ---------------------------------------------------------------------------
# (iv) streaming-decompress never buffers: write straight into a temp FILE (the
# real chunker/extractor pattern is download_atom(key, open(tmp,'wb'))).
# ---------------------------------------------------------------------------
clear(); put_bz2(pbzip2=True)
with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as tf:
    tmp_path = tf.name
try:
    with open(tmp_path, "wb") as f:
        pipeline.download_atom(KEY, f)
    loaded = np.load(tmp_path, allow_pickle=False)
    check("(iv) call-site pattern (open file + np.load) round-trips",
          np.array_equal(loaded, arr))
finally:
    os.unlink(tmp_path)

# ---------------------------------------------------------------------------
# (v) provisioner gate: COMPLETE when the .npy exists only as .npy.bz2.
# Import the provisioner module directly (no DB needed for gate_complete /
# _present_in_global — they operate on the in-memory `present` set).
# ---------------------------------------------------------------------------
prov_path = "/usr/src/savvy-provisioner/provisioner.py"
if not os.path.exists(prov_path):
    prov_path = os.environ.get("PROVISIONER_PATH", "")
import importlib.util  # noqa: E402
import types  # noqa: E402

# The provisioner imports pymysql at module top (DB driver) — irrelevant to the
# pure gate_complete / _present_in_global logic under test. Stub it if absent so
# the module imports without a live MariaDB driver.
if "pymysql" not in sys.modules:
    try:
        import pymysql  # noqa: F401
    except Exception:
        sys.modules["pymysql"] = types.ModuleType("pymysql")
        sys.modules["pymysql"].cursors = types.ModuleType("pymysql.cursors")

if prov_path and os.path.exists(prov_path):
    spec = importlib.util.spec_from_file_location("provisioner", prov_path)
    prov = importlib.util.module_from_spec(spec)
    try:
        spec.loader.exec_module(prov)
    except Exception as e:
        print("WARN provisioner import failed:", repr(e))
        prov = None
else:
    print("WARN provisioner.py not mounted; skipping gate tests")
    prov = None

if prov is not None:
    # _present_in_global is stubbed off (no global) so the gate decides on the
    # in-memory present set alone — exactly the COMPLETE-on-.bz2 acceptance.
    _real_present_in_global = prov._present_in_global
    prov._present_in_global = lambda k: False

    class _A:  # minimal atom stand-in; gate only reads atoms.keys()
        pass

    npy = "pipeline/ctx0/doc0/stage.hidden.json.npy"
    pt0 = "pipeline/ctx0/doc0/stage.hidden.json.win000.pt"
    expected = {npy, pt0, "pipeline/ctx0/doc0/meta.json"}

    # present set: the .npy is now .npy.bz2, the .pt is .pt.bz2, meta is raw.
    present_atoms = {
        npy + ".bz2": _A(),
        pt0 + ".bz2": _A(),
        "pipeline/ctx0/doc0/meta.json": _A(),
    }
    ok, missing = prov.gate_complete(expected, present_atoms)
    check("(v) gate COMPLETE when .npy is .npy.bz2 and .pt is .pt.bz2",
          ok and not missing)

    # Negative control: a genuinely-missing atom (no raw, no .bz2) -> INCOMPLETE.
    ok2, missing2 = prov.gate_complete(expected | {"pipeline/ctx0/doc0/absent.npy"},
                                       present_atoms)
    check("(v) gate INCOMPLETE when an atom is truly absent (no .bz2 either)",
          (not ok2) and missing2 == ["pipeline/ctx0/doc0/absent.npy"])

    # _present_in_global candidate generalization (restore the REAL function +
    # stub a global that holds only the .bz2 siblings).
    prov._present_in_global = _real_present_in_global
    have = {npy + ".bz2", pt0 + ".bz2"}
    class _G2:
        def head_object(self, Bucket=None, Key=None):
            if Key in have:
                return {}
            raise _StubNotFound()
    prov.global_s3 = lambda: _G2()
    check("(v) _present_in_global accepts .npy via .npy.bz2", prov._present_in_global(npy))
    check("(v) _present_in_global accepts .pt via .pt.bz2", prov._present_in_global(pt0))
    check("(v) _present_in_global rejects a truly-absent atom",
          not prov._present_in_global("pipeline/ctx0/doc0/absent.npy"))

print()
if failures:
    print("FAILURES:", failures)
    sys.exit(1)
print("ALL TESTS PASSED")
