"""Unit suite for common/pipeline_pool.py — PipelinePoolWorker.Execute, the
standard pipeline-stage contract that bridges PoolWorker (HTTP/heartbeat/state)
with the Garage-S3 + Redis result-atom pattern.

Pure unit:
  * boto3 / botocore are stubbed in conftest so the import chain
    (pipeline_pool -> common.pipeline) resolves without S3.
  * PoolWorker.__init__ does real I/O (Redis ping, HTTP bind, heartbeat thread,
    signal handlers). We NEVER call it — workers are built via a tiny subclass
    whose __init__ sets only the attributes Execute touches (_redis,
    output_stage). This isolates Execute's control flow.
  * The collaborating functions Execute calls (read_task, parse_task,
    write_result, stage_exists, read_stage, write_stage, storage_key) are
    imported into the pipeline_pool namespace, so we monkeypatch them THERE —
    that is the correct seam.
"""
import pytest

import pipeline_pool
from _commonfakes import FakeRedis


# ── A test worker that bypasses the heavy PoolWorker.__init__ ─────────────────

class _StageWorker(pipeline_pool.PipelinePoolWorker):
    """Concrete PipelinePoolWorker built without PoolWorker.__init__.

    Records process() calls so tests can assert on the computation path.
    """

    output_stage = "tokenized"

    def __init__(self, redis_client, output_stage="tokenized",
                 process_result=None, process_raises=None):
        # Deliberately DO NOT call super().__init__() — it opens sockets/threads.
        self._redis = redis_client
        self.output_stage = output_stage
        self._process_result = process_result if process_result is not None \
            else {"out": "payload"}
        self._process_raises = process_raises
        self.process_calls = []
        self.cleanup_calls = 0

    def process(self, doc_id, context_uuid, metadata, payload):
        self.process_calls.append((doc_id, context_uuid, metadata, payload))
        if self._process_raises is not None:
            raise self._process_raises
        return self._process_result

    def on_cleanup(self):
        self.cleanup_calls += 1


@pytest.fixture
def patch_seams(monkeypatch):
    """Install controllable fakes for every collaborator Execute calls.

    Returns a dict of recorders the test can inspect / configure.
    """
    rec = {
        "task_atom": None,          # what read_task returns
        "parsed": None,             # what parse_task returns
        "stage_exists": False,      # checkpoint decision
        "read_stage_payload": {"raw": "data"},
        "write_stage_key": "pipeline/ctx/doc/stage.tokenized.json",
        "storage_key_value": "pipeline/ctx/doc/stage.tokenized.json",
        "results": [],              # captured write_result calls
        "read_task_calls": [],
        "stage_exists_calls": [],
        "read_stage_calls": [],
        "write_stage_calls": [],
        "storage_key_calls": [],
    }

    def fake_read_task(r, uuid):
        rec["read_task_calls"].append((r, uuid))
        return rec["task_atom"]

    def fake_parse_task(atom):
        # Default parse if the test didn't override.
        if rec["parsed"] is not None:
            return rec["parsed"]
        return (atom["doc_id"], atom["stage"], atom["storage_key"],
                atom["context_uuid"], atom.get("metadata", {}))

    def fake_stage_exists(ctx, doc_id, out_stage):
        rec["stage_exists_calls"].append((ctx, doc_id, out_stage))
        return rec["stage_exists"]

    def fake_read_stage(in_key):
        rec["read_stage_calls"].append(in_key)
        return rec["read_stage_payload"]

    def fake_write_stage(ctx, doc_id, out_stage, output):
        rec["write_stage_calls"].append((ctx, doc_id, out_stage, output))
        return rec["write_stage_key"]

    def fake_storage_key(ctx, doc_id, out_stage):
        rec["storage_key_calls"].append((ctx, doc_id, out_stage))
        return rec["storage_key_value"]

    def fake_write_result(r, uuid, out_stage, out_key, **kw):
        rec["results"].append({
            "uuid": uuid, "out_stage": out_stage, "out_key": out_key, "kw": kw,
        })

    monkeypatch.setattr(pipeline_pool, "read_task", fake_read_task)
    monkeypatch.setattr(pipeline_pool, "parse_task", fake_parse_task)
    monkeypatch.setattr(pipeline_pool, "stage_exists", fake_stage_exists)
    monkeypatch.setattr(pipeline_pool, "read_stage", fake_read_stage)
    monkeypatch.setattr(pipeline_pool, "write_stage", fake_write_stage)
    monkeypatch.setattr(pipeline_pool, "storage_key", fake_storage_key)
    monkeypatch.setattr(pipeline_pool, "write_result", fake_write_result)
    return rec


def _atom(doc_id="doc-1", stage="raw",
          in_key="pipeline/ctx/doc-1/stage.raw.json",
          ctx="ctx-1", metadata=None):
    return {
        "doc_id": doc_id,
        "stage": stage,
        "storage_key": in_key,
        "context_uuid": ctx,
        "metadata": metadata if metadata is not None else {"title": "T"},
    }


# ── Happy path: full stage execution ─────────────────────────────────────────

def test_execute_full_path_runs_process_and_writes_result(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False
    patch_seams["read_stage_payload"] = {"tokens": [1, 2, 3]}
    patch_seams["write_stage_key"] = "pipeline/ctx-1/doc-1/stage.tokenized.json"

    w = _StageWorker(FakeRedis(), process_result={"tokenized": "ok"})
    ok = w.Execute("task-uuid-1")

    assert ok is True, "successful stage returns True"
    # process() was called once with the parsed fields + the loaded payload.
    assert len(w.process_calls) == 1
    doc_id, ctx, meta, payload = w.process_calls[0]
    assert doc_id == "doc-1"
    assert ctx == "ctx-1"
    assert meta == {"title": "T"}
    assert payload == {"tokens": [1, 2, 3]}, "process receives the read_stage payload"


def test_execute_full_path_writes_correct_result_atom(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False
    patch_seams["write_stage_key"] = "OUT/KEY"

    w = _StageWorker(FakeRedis(), output_stage="tokenized")
    w.Execute("u1")

    assert len(patch_seams["results"]) == 1, "exactly one result atom written"
    res = patch_seams["results"][0]
    assert res["uuid"] == "u1"
    assert res["out_stage"] == "tokenized"
    assert res["out_key"] == "OUT/KEY", "result out_key is the write_stage return"


def test_execute_full_path_uses_input_key_for_read_stage(patch_seams):
    patch_seams["task_atom"] = _atom(in_key="THE/INPUT/KEY")
    patch_seams["stage_exists"] = False

    w = _StageWorker(FakeRedis())
    w.Execute("u1")

    assert patch_seams["read_stage_calls"] == ["THE/INPUT/KEY"], \
        "read_stage must be called with the task's storage_key, not the output key"


def test_execute_full_path_write_stage_gets_output_of_process(patch_seams):
    patch_seams["task_atom"] = _atom(ctx="ctx-9", doc_id="d9")
    patch_seams["stage_exists"] = False

    sentinel_output = {"computed": [9, 9, 9]}
    w = _StageWorker(FakeRedis(), output_stage="hidden",
                     process_result=sentinel_output)
    w.Execute("u1")

    assert len(patch_seams["write_stage_calls"]) == 1
    ctx, doc_id, out_stage, output = patch_seams["write_stage_calls"][0]
    assert ctx == "ctx-9"
    assert doc_id == "d9"
    assert out_stage == "hidden", "write_stage uses the worker's output_stage"
    assert output is sentinel_output, "write_stage receives exactly process()'s return"


def test_execute_full_path_does_not_call_cleanup_itself(patch_seams):
    # on_cleanup is the PoolWorker /execute handler's job, NOT Execute's.
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False
    w = _StageWorker(FakeRedis())
    w.Execute("u1")
    assert w.cleanup_calls == 0, "Execute must not invoke on_cleanup (handler does)"


# ── Checkpoint hit: output already in Garage ─────────────────────────────────

def test_execute_checkpoint_hit_skips_process(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = True   # output already present
    patch_seams["storage_key_value"] = "pipeline/ctx-1/doc-1/stage.tokenized.json"

    w = _StageWorker(FakeRedis())
    ok = w.Execute("u1")

    assert ok is True, "checkpoint hit still returns True"
    assert w.process_calls == [], "process() must be skipped on checkpoint hit"
    assert patch_seams["read_stage_calls"] == [], "no Garage read on checkpoint"
    assert patch_seams["write_stage_calls"] == [], "no Garage write on checkpoint"


def test_execute_checkpoint_hit_writes_result_with_existing_key(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = True
    patch_seams["storage_key_value"] = "EXISTING/OUT/KEY"

    w = _StageWorker(FakeRedis(), output_stage="tokenized")
    w.Execute("u-cp")

    assert len(patch_seams["results"]) == 1
    res = patch_seams["results"][0]
    assert res["uuid"] == "u-cp"
    assert res["out_stage"] == "tokenized"
    assert res["out_key"] == "EXISTING/OUT/KEY", \
        "checkpoint result key comes from storage_key(), not write_stage"


def test_execute_checkpoint_uses_storage_key_not_write_stage(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = True
    w = _StageWorker(FakeRedis())
    w.Execute("u1")
    assert len(patch_seams["storage_key_calls"]) == 1, \
        "checkpoint path resolves the key via storage_key()"


def test_execute_stage_exists_checked_with_output_stage(patch_seams):
    patch_seams["task_atom"] = _atom(ctx="CC", doc_id="DD")
    patch_seams["stage_exists"] = True
    w = _StageWorker(FakeRedis(), output_stage="embedded")
    w.Execute("u1")
    assert patch_seams["stage_exists_calls"] == [("CC", "DD", "embedded")], \
        "checkpoint is keyed on (context, doc, OUTPUT stage)"


# ── Task-not-found path ──────────────────────────────────────────────────────

def test_execute_missing_task_returns_false(patch_seams):
    patch_seams["task_atom"] = None    # read_task -> None
    w = _StageWorker(FakeRedis())
    ok = w.Execute("ghost")
    assert ok is False, "missing task atom -> Execute returns False"


def test_execute_missing_task_does_no_work(patch_seams):
    patch_seams["task_atom"] = None
    w = _StageWorker(FakeRedis())
    w.Execute("ghost")
    assert w.process_calls == [], "no process on missing task"
    assert patch_seams["results"] == [], "no result atom written on missing task"
    assert patch_seams["stage_exists_calls"] == [], "no checkpoint check on missing task"
    assert patch_seams["read_stage_calls"] == [], "no Garage read on missing task"


def test_execute_reads_task_with_given_uuid(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = True   # short-circuit; we only care about read_task
    w = _StageWorker(FakeRedis())
    w.Execute("the-uuid-xyz")
    assert len(patch_seams["read_task_calls"]) == 1
    _, uuid = patch_seams["read_task_calls"][0]
    assert uuid == "the-uuid-xyz", "Execute reads the atom for the dispatched uuid"


# ── Error propagation ────────────────────────────────────────────────────────

def test_execute_propagates_process_exception(patch_seams):
    # Execute does NOT swallow process() errors — the PoolWorker /execute
    # handler is responsible for rolling state back to FREE and re-raising.
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False
    w = _StageWorker(FakeRedis(), process_raises=RuntimeError("boom"))
    with pytest.raises(RuntimeError, match="boom"):
        w.Execute("u1")


def test_execute_no_result_written_when_process_raises(patch_seams):
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False
    w = _StageWorker(FakeRedis(), process_raises=ValueError("bad"))
    with pytest.raises(ValueError):
        w.Execute("u1")
    assert patch_seams["results"] == [], "a failed process must not emit a success result"
    assert patch_seams["write_stage_calls"] == [], "no Garage write when process fails"


# ── Ordering invariants (the documented Execute contract) ────────────────────

def test_execute_order_read_then_checkpoint_then_read_write_then_result(patch_seams):
    order = []
    patch_seams["task_atom"] = _atom()
    patch_seams["stage_exists"] = False

    # Wrap each seam to record invocation order.
    import pipeline_pool as pp
    orig = {n: getattr(pp, n) for n in
            ("read_task", "stage_exists", "read_stage", "write_stage", "write_result")}

    def wrap(name):
        f = orig[name]

        def inner(*a, **k):
            order.append(name)
            return f(*a, **k)
        return inner

    for n in orig:
        setattr(pp, n, wrap(n))
    try:
        w = _StageWorker(FakeRedis())
        w.Execute("u1")
    finally:
        for n, f in orig.items():
            setattr(pp, n, f)

    assert order == ["read_task", "stage_exists", "read_stage",
                     "write_stage", "write_result"], \
        "Execute must read atom -> checkpoint -> load -> compute/store -> result"


# ── process() / on_cleanup() base behaviour ──────────────────────────────────

def test_base_process_raises_not_implemented():
    # The base class process() is abstract-by-convention (raises).
    w = object.__new__(pipeline_pool.PipelinePoolWorker)
    with pytest.raises(NotImplementedError):
        w.process("d", "c", {}, "payload")


def test_base_on_cleanup_is_noop_returns_none():
    w = object.__new__(pipeline_pool.PipelinePoolWorker)
    assert w.on_cleanup() is None, "default on_cleanup is a no-op"


def test_pipeline_pool_worker_is_poolworker_subclass():
    from pool_worker import PoolWorker
    assert issubclass(pipeline_pool.PipelinePoolWorker, PoolWorker), \
        "PipelinePoolWorker must extend PoolWorker for HTTP/heartbeat/state"


# ── metadata threading through to process() ──────────────────────────────────

def test_execute_threads_parsed_metadata_to_process(patch_seams):
    meta = {"author": "A", "year": 2021, "tags": ["x"]}
    patch_seams["task_atom"] = _atom(metadata=meta)
    patch_seams["stage_exists"] = False
    w = _StageWorker(FakeRedis())
    w.Execute("u1")
    assert w.process_calls[0][2] == meta, "metadata dict reaches process() unchanged"
