import pytest

from tht.corpus.chunk import ChunkPolicy
from tht.corpus.pipeline import CorpusPipeline, PipelineError
from tht.corpus.store import CorpusStore
from tht.ports.evidence import AcquiredDocument, SourceObject
from tht.ports.vector import VectorCapabilities


class Source:
    def __init__(self, documents):
        self.documents = documents
        self.acquire_calls = []

    def discover(self):
        return [item[0] for item in self.documents]

    def acquire(self, item):
        self.acquire_calls.append(item.source_id)
        payload = next(payload for source, payload in self.documents if source.source_id == item.source_id)
        if isinstance(payload, Exception):
            raise payload
        return AcquiredDocument(source=item, content=payload.encode())


class Embedder:
    def __init__(self, dim=3, fail=False):
        self.dim = dim
        self.fail = fail
        self.calls = []

    def embed_documents(self, texts):
        self.calls.extend(texts)
        if self.fail:
            raise RuntimeError("embed failed")
        return [[float(i) for i in range(self.dim)] for _ in texts]


class Vectors:
    capabilities = VectorCapabilities(search=True, existing_hashes=True, upsert=True)

    def __init__(self, fail=False):
        self.fail = fail
        self.records = []

    def upsert(self, collection, records):
        self.records.extend(records[:1] if self.fail else records)
        if self.fail:
            raise RuntimeError("partial write")
        return len(records)

    def existing_hashes(self, collection, kinds):
        return {
            value.record.id: value.content_hash for value in self.records
        }

    def delete_generation(self, collection, generation):
        self.records = [
            value for value in self.records
            if value.record.metadata["vector_generation"] != generation
        ]
        return 0


class InterruptingVectors(Vectors):
    def __init__(self):
        super().__init__()
        self.batches = []
        self.interrupt = True

    def upsert(self, collection, records):
        self.batches.append([value.record.id for value in records])
        if self.interrupt:
            self.interrupt = False
            self.records.append(records[0])
            raise KeyboardInterrupt("process interruption after partial write")
        self.records.extend(records)
        return len(records)


def item(name, fingerprint):
    return SourceObject(
        source_id=f"fs:{name}", uri=f"file:///safe/{name}.md", fingerprint=f"sha256:{fingerprint}"
    )


def pipeline(tmp_path, source, *, embedder=None, vectors=None, model="model-a", policy=None):
    return CorpusPipeline(
        store=CorpusStore(tmp_path / "corpus"), sources=[source],
        embedder=embedder or Embedder(), vector_store=vectors or Vectors(),
        embedding_model=model, embedding_dimensions=3,
        chunk_policy=policy or ChunkPolicy(version="chunk-v1", max_chars=100),
        pipeline_version="evidence-v1",
    )


def test_unchanged_documents_skip_acquire_normalize_chunk_and_embed(tmp_path):
    one = item("one", "a")
    first_source = Source([(one, "hello")])
    first = pipeline(tmp_path, first_source)
    first.run()
    second_source = Source([(one, "ignored")])
    second_embedder = Embedder()
    result = pipeline(tmp_path, second_source, embedder=second_embedder).run()
    assert result.unchanged == ("fs:one",)
    assert second_source.acquire_calls == []
    assert second_embedder.calls == []


def test_removed_documents_are_marked_and_absent_from_new_manifest(tmp_path):
    one, two = item("one", "a"), item("two", "b")
    pipeline(tmp_path, Source([(one, "one"), (two, "two")])).run()
    result = pipeline(tmp_path, Source([(one, "one")])).run()
    assert result.removed == ("fs:two",)
    assert {doc.source_id for doc in result.manifest.documents} == {"fs:one"}


def test_model_or_chunk_policy_change_forces_full_rebuild(tmp_path):
    one = item("one", "a")
    pipeline(tmp_path, Source([(one, "hello")])).run()
    source = Source([(one, "hello")])
    changed = pipeline(tmp_path, source, model="model-b").run()
    assert changed.changed == ("fs:one",)
    assert source.acquire_calls == ["fs:one"]


def test_partial_vector_failure_never_changes_active_or_exposes_generation(tmp_path):
    one = item("one", "a")
    good = pipeline(tmp_path, Source([(one, "old")]))
    old = good.run().generation
    changed = item("one", "b")
    vectors = Vectors(fail=True)
    broken = pipeline(tmp_path, Source([(changed, "new")]), vectors=vectors)
    with pytest.raises(PipelineError):
        broken.run()
    assert broken.store.active_generation() == old
    assert vectors.records[0].record.metadata["vector_generation"] != old


def test_dimension_mismatch_fails_before_vector_write_and_publish(tmp_path):
    one = item("one", "a")
    vectors = Vectors()
    candidate = pipeline(tmp_path, Source([(one, "hello")]), embedder=Embedder(dim=2), vectors=vectors)
    with pytest.raises(PipelineError, match="dimension"):
        candidate.run()
    assert vectors.records == []
    assert candidate.store.active_generation() is None


def test_dry_run_and_failed_acquire_never_change_active(tmp_path):
    one = item("one", "a")
    active = pipeline(tmp_path, Source([(one, "old")])).run().generation
    changed = item("one", "b")
    dry = pipeline(tmp_path, Source([(changed, "new")])).run(dry_run=True)
    assert dry.published is False
    assert dry.generation is None
    assert dry.manifest.documents[0].content == "old"
    with pytest.raises(PipelineError):
        pipeline(tmp_path, Source([(changed, RuntimeError("boom"))])).run()
    assert CorpusStore(tmp_path / "corpus").active_generation() == active


def test_job_pipeline_uses_ordered_plan_and_returns_run_id(tmp_path):
    one = item("one", "a")
    candidate = pipeline(tmp_path, Source([(one, "hello")]))
    result = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
    )
    assert result.status == "succeeded"
    assert result.run_id and len(result.run_id) == 32
    checkpoint = tmp_path / ".tht-jobs" / "evidence" / "runs" / result.run_id / "checkpoint.json"
    payload = __import__("json").loads(checkpoint.read_text())
    assert [stage["name"] for stage in payload["stages"]] == [
        "discover", "acquire_normalize_chunk", "embed", "vector_upsert",
        "stage_validate", "publish", "retention_cleanup",
    ]


def test_job_pipeline_dry_run_only_discovers_and_reports_changes(tmp_path):
    one = item("one", "a")
    source = Source([(one, "hello")])
    embedder = Embedder()
    vectors = Vectors()
    result = pipeline(tmp_path, source, embedder=embedder, vectors=vectors).run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
        dry_run=True,
    )
    assert result.changed == ("fs:one",)
    assert source.acquire_calls == []
    assert embedder.calls == []
    assert vectors.records == []
    assert result.generation is None and result.published is False


@pytest.mark.parametrize("crash_stage", [
    "discover", "acquire_normalize_chunk", "embed", "vector_upsert",
    "stage_validate", "publish", "retention_cleanup",
])
def test_job_pipeline_crash_after_each_stage_resumes_without_duplicate_effects(tmp_path, crash_stage):
    one = item("one", "a")
    source = Source([(one, "hello")])
    embedder = Embedder()
    vectors = Vectors()
    candidate = pipeline(tmp_path, source, embedder=embedder, vectors=vectors)

    class Crash(BaseException):
        pass

    def fault(_context, stage):
        if stage == crash_stage:
            raise Crash()

    with pytest.raises(Crash):
        candidate.run_as_job(
            workspace_id="demo", workspace_root=tmp_path,
            config_fingerprint="sha256:" + "1" * 64,
            input_fingerprint="sha256:" + "2" * 64,
            after_stage_return=fault,
        )
    runs = tmp_path / ".tht-jobs" / "evidence" / "runs"
    crashed = next(runs.iterdir()).name
    result = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
        resume_run_id=crashed,
    )
    assert result.status == "succeeded"
    assert source.acquire_calls == ["fs:one"]
    assert len(embedder.calls) == 1
    assert len(vectors.records) == 1


def test_job_pipeline_raw_upsert_failure_compensates_and_resumes_with_new_generation(tmp_path):
    one = item("one", "a")
    vectors = Vectors(fail=True)
    candidate = pipeline(tmp_path, Source([(one, "hello")]), vectors=vectors)
    first = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
    )
    assert first.status == "failed"
    assert vectors.records == []
    old_generation = first.generation
    vectors.fail = False
    resumed = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
        resume_run_id=first.run_id,
    )
    assert resumed.status == "succeeded", resumed
    assert resumed.generation != old_generation
    assert candidate.store.active_generation() == resumed.generation


def test_job_pipeline_raw_stage_failure_compensates_vectors_and_resumes(tmp_path, monkeypatch):
    one = item("one", "a")
    vectors = Vectors()
    candidate = pipeline(tmp_path, Source([(one, "hello")]), vectors=vectors)
    real_stage = candidate.store.stage
    calls = 0

    def fail_once(*args, **kwargs):
        nonlocal calls
        calls += 1
        if calls == 1:
            raise OSError("raw stage failure")
        return real_stage(*args, **kwargs)

    monkeypatch.setattr(candidate.store, "stage", fail_once)
    first = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
    )
    assert first.status == "failed" and vectors.records == []
    resumed = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
        resume_run_id=first.run_id,
    )
    assert resumed.status == "succeeded", resumed


@pytest.mark.parametrize("stage,filename", [
    ("discover", "plan.json"),
    ("acquire_normalize_chunk", "manifest.json"),
    ("embed", "embeddings.json"),
])
@pytest.mark.parametrize("mutation", ["missing", "tampered"])
def test_job_pipeline_rejects_corrupt_required_artifacts_before_resume(
    tmp_path, stage, filename, mutation,
):
    one = item("one", "a")
    candidate = pipeline(tmp_path, Source([(one, "hello")]))

    class Crash(BaseException):
        pass

    with pytest.raises(Crash):
        candidate.run_as_job(
            workspace_id="demo", workspace_root=tmp_path,
            config_fingerprint="sha256:" + "1" * 64,
            input_fingerprint="sha256:" + "2" * 64,
            after_stage_return=lambda _context, name: (
                (_ for _ in ()).throw(Crash()) if name == stage else None
            ),
        )
    runs = tmp_path / ".tht-jobs" / "evidence" / "runs"
    crashed = next(runs.iterdir())
    target = crashed / "artifacts" / filename
    target.unlink() if mutation == "missing" else target.write_text("tampered")
    from tht.jobs.runner import CorruptCheckpointError
    with pytest.raises(CorruptCheckpointError, match="artifact"):
        candidate.run_as_job(
            workspace_id="demo", workspace_root=tmp_path,
            config_fingerprint="sha256:" + "1" * 64,
            input_fingerprint="sha256:" + "2" * 64,
            resume_run_id=crashed.name,
        )


def test_vector_intent_is_reconciled_after_process_interruption_without_duplicate_upsert(tmp_path):
    one = item("one", "a")
    vectors = InterruptingVectors()
    candidate = pipeline(
        tmp_path, Source([(one, "a" * 250)]), vectors=vectors,
        policy=ChunkPolicy(version="chunk-v1", max_chars=100),
    )
    with pytest.raises(KeyboardInterrupt):
        candidate.run_as_job(
            workspace_id="demo", workspace_root=tmp_path,
            config_fingerprint="sha256:" + "1" * 64,
            input_fingerprint="sha256:" + "2" * 64,
        )
    runs = tmp_path / ".tht-jobs" / "evidence" / "runs"
    interrupted = next(runs.iterdir())
    checkpoint = __import__("json").loads((interrupted / "checkpoint.json").read_text())
    vector_stage = checkpoint["stages"][3]
    assert vector_stage["status"] == "running"
    assert vector_stage["effect_state"] == "intent"
    first_written = vectors.batches[0][0]

    result = candidate.run_as_job(
        workspace_id="demo", workspace_root=tmp_path,
        config_fingerprint="sha256:" + "1" * 64,
        input_fingerprint="sha256:" + "2" * 64,
        resume_run_id=interrupted.name,
    )
    assert result.status == "succeeded" and result.published is True
    assert first_written not in vectors.batches[1]
    assert len(vectors.records) == 3