from __future__ import annotations

import copy
import json
import os
import stat
import subprocess
import sys
from collections.abc import Callable, Iterator
from pathlib import Path

import pytest

from ops import run_local_corpus_knowledge_rebuild as rebuild
from ops import shard_corpus_extraction_work as sharder

REPO_ROOT = Path(__file__).resolve().parents[1]
CLI = REPO_ROOT / "ops" / "shard_corpus_extraction_work.py"
FIXTURE_ROOT = REPO_ROOT / "tests" / "fixtures" / "teleo_corpus_rebuild"
PRIVATE_TEXT = "PRIVATE SOURCE BODY MUST NOT ENTER SHARDS"


def _load_json(path: Path) -> dict:
    value = json.loads(path.read_text(encoding="utf-8"))
    assert isinstance(value, dict)
    return value


def _write_json(path: Path, value: object) -> None:
    path.write_bytes(sharder.rendered_json_bytes(value))
    path.chmod(0o600)


def _tree_bytes(root: Path) -> dict[str, bytes]:
    return {path.relative_to(root).as_posix(): path.read_bytes() for path in sorted(root.rglob("*")) if path.is_file()}


def _all_keys(value: object) -> Iterator[str]:
    if isinstance(value, dict):
        for key, child in value.items():
            yield key
            yield from _all_keys(child)
    elif isinstance(value, list):
        for child in value:
            yield from _all_keys(child)


def _build_inventory(tmp_path: Path, source_root: Path = FIXTURE_ROOT, *, name: str = "inventory") -> Path:
    output = tmp_path / name
    rebuild.build_source_inventory_bundle(source_root, output)
    return output


def _make_section_source(root: Path, count: int, *, private_text: str = PRIVATE_TEXT) -> Path:
    root.mkdir(parents=True, exist_ok=True)
    body = "".join(f"# Section {index:04d}\n\n{private_text} {index:04d}.\n\n" for index in range(count))
    source = root / "source.md"
    source.write_text(body, encoding="utf-8")
    return source


def _run_cli(*args: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [sys.executable, str(CLI), *args],
        cwd=REPO_ROOT,
        check=False,
        capture_output=True,
        text=True,
    )


def _shard_documents(bundle: Path) -> list[tuple[Path, dict]]:
    manifest = _load_json(bundle / sharder.MANIFEST_PATH)
    return [(bundle / record["path"], _load_json(bundle / record["path"])) for record in manifest["shards"]]


def _rehash_bundle(bundle: Path) -> None:
    manifest_path = bundle / sharder.MANIFEST_PATH
    receipt_path = bundle / sharder.RECEIPT_PATH
    manifest = _load_json(manifest_path)
    for record in manifest["shards"]:
        shard_path = bundle / record["path"]
        document = _load_json(shard_path)
        items = document["items"]
        work_ids = [item["work_item_id"] for item in items]
        document["work_item_count"] = len(items)
        document["work_item_ids_sha256"] = sharder.canonical_sha256(work_ids)
        _write_json(shard_path, document)
        raw = shard_path.read_bytes()
        record.update(
            bytes=len(raw),
            sha256=sharder.sha256_bytes(raw),
            work_item_count=len(items),
            work_item_ids_sha256=sharder.canonical_sha256(work_ids),
        )
    _write_json(manifest_path, manifest)

    receipt = _load_json(receipt_path)
    artifact_paths = [sharder.MANIFEST_PATH, *(record["path"] for record in manifest["shards"])]
    artifacts = sharder._artifact_records(bundle, artifact_paths)
    receipt["outputs"] = {
        "artifact_count": len(artifacts),
        "artifacts": artifacts,
        "manifest_path": sharder.MANIFEST_PATH,
    }
    receipt["hashes"] = {
        "bundle_tree_sha256": sharder.canonical_sha256(artifacts),
        "manifest_file_sha256": sharder.sha256_bytes(manifest_path.read_bytes()),
    }
    receipt["hashes"]["receipt_content_sha256"] = sharder.canonical_sha256(receipt)
    _write_json(receipt_path, receipt)


def test_build_is_deterministic_private_complete_and_source_text_free(tmp_path: Path) -> None:
    source_root = tmp_path / "source"
    _make_section_source(source_root, 12)
    inventory_bundle = _build_inventory(tmp_path, source_root)
    first = tmp_path / "shards-a"
    second = tmp_path / "shards-b"

    first_result = sharder.build_shard_bundle(inventory_bundle, first, 4)["verification"]
    second_result = sharder.build_shard_bundle(inventory_bundle, second, 4)["verification"]

    assert first_result["status"] == second_result["status"] == "pass"
    assert first_result["work_item_count"] == second_result["work_item_count"] == 12
    assert first_result["shard_count"] == second_result["shard_count"] == 4
    assert first_result["bundle_tree_sha256"] == second_result["bundle_tree_sha256"]
    assert _tree_bytes(first) == _tree_bytes(second)
    persisted = b"".join(_tree_bytes(first).values())
    assert PRIVATE_TEXT.encode() not in persisted
    assert str(source_root).encode() not in persisted
    assert str(inventory_bundle).encode() not in persisted
    assert str(first).encode() not in persisted

    forbidden_keys = {"claim", "claims", "proposal", "proposals", "admitted", "database_rows"}
    for document in [_load_json(first / sharder.MANIFEST_PATH), _load_json(first / sharder.RECEIPT_PATH)]:
        assert forbidden_keys.isdisjoint(_all_keys(document))
        for key, expected in sharder.NO_WRITE_FLAGS.items():
            assert document[key] is expected

    for path in first.rglob("*"):
        expected_mode = 0o600 if path.is_file() else 0o700
        assert stat.S_IMODE(path.stat().st_mode) == expected_mode
    assert stat.S_IMODE(first.stat().st_mode) == 0o700
    assert not (first / sharder.INCOMPLETE_MARKER).exists()


def test_hash_assignment_covers_every_exact_inventory_item_once(tmp_path: Path) -> None:
    source_root = tmp_path / "source"
    _make_section_source(source_root, 40)
    inventory_bundle = _build_inventory(tmp_path, source_root)
    output = tmp_path / "shards"
    sharder.build_shard_bundle(inventory_bundle, output, 7)
    inventory_items = _load_json(inventory_bundle / rebuild.EXTRACTION_WORK_INDEX_PATH)["items"]

    assigned: list[dict] = []
    for index, (_path, document) in enumerate(_shard_documents(output)):
        ids = [item["work_item_id"] for item in document["items"]]
        assert ids == sorted(ids, key=lambda value: value.encode("utf-8"))
        assert all(sharder.assignment_for(work_id, 7) == index for work_id in ids)
        assigned.extend(document["items"])

    assert len(assigned) == len(inventory_items)
    assert {item["work_item_id"] for item in assigned} == {item["work_item_id"] for item in inventory_items}
    assert {item["work_item_id"]: item for item in assigned} == {item["work_item_id"]: item for item in inventory_items}


@pytest.mark.parametrize("shard_count", [0, 1025])
def test_shard_count_must_be_within_hard_bounds(tmp_path: Path, shard_count: int) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    with pytest.raises(sharder.ShardError) as error:
        sharder.build_shard_bundle(inventory_bundle, tmp_path / "shards", shard_count)
    assert error.value.code == "invalid_shard_count"
    assert not (tmp_path / "shards").exists()


def test_shard_count_must_not_exceed_work_item_count(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    work_count = _load_json(inventory_bundle / rebuild.EXTRACTION_WORK_INDEX_PATH)["work_item_count"]
    with pytest.raises(sharder.ShardError) as error:
        sharder.build_shard_bundle(inventory_bundle, tmp_path / "shards", work_count + 1)
    assert error.value.code == "invalid_shard_count"


def test_existing_output_is_never_overwritten(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"
    output.mkdir(mode=0o700)
    sentinel = output / "sentinel"
    sentinel.write_text("preserve", encoding="utf-8")

    with pytest.raises(sharder.ShardError) as error:
        sharder.build_shard_bundle(inventory_bundle, output, 2)

    assert error.value.code == "output_exists"
    assert sentinel.read_text(encoding="utf-8") == "preserve"


def test_marker_write_failure_closes_and_removes_reserved_output(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"
    original_write = sharder._write_bytes_at

    def fail_after_write(directory_fd: int, name: str, payload: bytes) -> None:
        original_write(directory_fd, name, payload)
        raise OSError("synthetic marker publication failure")

    monkeypatch.setattr(sharder, "_write_bytes_at", fail_after_write)
    with pytest.raises(OSError, match="synthetic marker publication failure"):
        sharder.build_shard_bundle(inventory_bundle, output, 2)

    assert not output.exists()


def test_receipt_is_published_while_incomplete_marker_exists(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"

    def inspect(reserved: sharder.ReservedOutput) -> None:
        marker = os.stat(sharder.INCOMPLETE_MARKER, dir_fd=reserved.root_fd, follow_symlinks=False)
        assert stat.S_ISREG(marker.st_mode)
        with pytest.raises(FileNotFoundError):
            os.stat(sharder.RECEIPT_PATH, dir_fd=reserved.root_fd, follow_symlinks=False)

    monkeypatch.setattr(sharder, "_before_receipt_publication", inspect)
    sharder.build_shard_bundle(inventory_bundle, output, 2)
    assert not (output / sharder.INCOMPLETE_MARKER).exists()
    assert (output / sharder.RECEIPT_PATH).is_file()


def test_destination_swap_does_not_overwrite_or_cleanup_replacement(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"
    moved = tmp_path / "reserved-moved"

    def swap(_reserved: sharder.ReservedOutput) -> None:
        output.rename(moved)
        output.mkdir(mode=0o700)
        sentinel = output / "replacement-sentinel"
        sentinel.write_text("preserve replacement", encoding="utf-8")
        sentinel.chmod(0o600)

    monkeypatch.setattr(sharder, "_after_output_reservation", swap)
    with pytest.raises(sharder.ShardError) as error:
        sharder.build_shard_bundle(inventory_bundle, output, 2)

    assert error.value.code == "output_binding_changed"
    assert (output / "replacement-sentinel").read_text(encoding="utf-8") == "preserve replacement"
    assert moved.is_dir()
    assert list(moved.iterdir()) == []


def test_inventory_change_before_publication_rejects_and_cleans_owned_output(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"

    def mutate(root: Path) -> None:
        receipt = root / rebuild.RECEIPT_PATH
        receipt.write_bytes(receipt.read_bytes() + b" ")

    monkeypatch.setattr(sharder, "_before_final_publication", mutate)
    with pytest.raises(sharder.ShardError) as error:
        sharder.build_shard_bundle(inventory_bundle, output, 2)

    assert error.value.code == "inventory_changed"
    assert not output.exists()


def _mutate_unknown(bundle: Path) -> None:
    for path, document in _shard_documents(bundle):
        if document["items"]:
            document["items"][0]["work_item_id"] = "f" * 64
            document["items"].sort(key=lambda item: item["work_item_id"].encode("utf-8"))
            _write_json(path, document)
            return
    raise AssertionError("fixture unexpectedly produced no shard items")


def _mutate_missing(bundle: Path) -> None:
    for path, document in _shard_documents(bundle):
        if document["items"]:
            document["items"].pop()
            _write_json(path, document)
            return
    raise AssertionError("fixture unexpectedly produced no shard items")


def _mutate_duplicate(bundle: Path) -> None:
    for path, document in _shard_documents(bundle):
        if document["items"]:
            document["items"].append(copy.deepcopy(document["items"][0]))
            document["items"].sort(key=lambda item: item["work_item_id"].encode("utf-8"))
            _write_json(path, document)
            return
    raise AssertionError("fixture unexpectedly produced no shard items")


def _mutate_overclaim(bundle: Path) -> None:
    receipt_path = bundle / sharder.RECEIPT_PATH
    receipt = _load_json(receipt_path)
    receipt["database_write_performed"] = True
    _write_json(receipt_path, receipt)


@pytest.mark.parametrize(
    ("mutator", "expected_codes"),
    [
        (_mutate_unknown, {"unknown_work_item"}),
        (_mutate_missing, {"missing_work_item"}),
        (_mutate_duplicate, {"duplicate_work_item", "invalid_shard_order"}),
        (_mutate_overclaim, {"false_sharding_claim"}),
    ],
)
def test_verifier_rejects_fully_rehashed_semantic_tampering(
    tmp_path: Path,
    mutator: Callable[[Path], None],
    expected_codes: set[str],
) -> None:
    source_root = tmp_path / "source"
    _make_section_source(source_root, 20)
    inventory_bundle = _build_inventory(tmp_path, source_root)
    output = tmp_path / "shards"
    sharder.build_shard_bundle(inventory_bundle, output, 4)

    mutator(output)
    _rehash_bundle(output)

    with pytest.raises(sharder.ShardError) as error:
        sharder.verify_shard_bundle(inventory_bundle, output)
    assert error.value.code in expected_codes


def test_verifier_rejects_wrong_inventory_even_with_same_cardinality(tmp_path: Path) -> None:
    first_source = tmp_path / "source-a"
    second_source = tmp_path / "source-b"
    _make_section_source(first_source, 8, private_text="first corpus")
    _make_section_source(second_source, 8, private_text="second corpus")
    first_inventory = _build_inventory(tmp_path, first_source, name="inventory-a")
    second_inventory = _build_inventory(tmp_path, second_source, name="inventory-b")
    output = tmp_path / "shards"
    sharder.build_shard_bundle(first_inventory, output, 3)

    with pytest.raises(sharder.ShardError) as error:
        sharder.verify_shard_bundle(second_inventory, output)
    assert error.value.code == "inventory_mismatch"


def test_verifier_rejects_extra_file_and_incomplete_marker(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    extra_output = tmp_path / "extra-shards"
    sharder.build_shard_bundle(inventory_bundle, extra_output, 2)
    extra = extra_output / "extra.json"
    extra.write_text("{}\n", encoding="utf-8")
    extra.chmod(0o600)
    with pytest.raises(sharder.ShardError) as extra_error:
        sharder.verify_shard_bundle(inventory_bundle, extra_output)
    assert extra_error.value.code == "unexpected_bundle_entry"

    incomplete_output = tmp_path / "incomplete-shards"
    sharder.build_shard_bundle(inventory_bundle, incomplete_output, 2)
    marker = incomplete_output / sharder.INCOMPLETE_MARKER
    marker.write_text("incomplete\n", encoding="utf-8")
    marker.chmod(0o600)
    with pytest.raises(sharder.ShardError) as incomplete_error:
        sharder.verify_shard_bundle(inventory_bundle, incomplete_output)
    assert incomplete_error.value.code == "incomplete_shard_bundle"


def test_verifier_rejects_symlinked_shard(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"
    sharder.build_shard_bundle(inventory_bundle, output, 2)
    manifest = _load_json(output / sharder.MANIFEST_PATH)
    target = output / manifest["shards"][0]["path"]
    retained = target.with_suffix(".retained")
    target.rename(retained)
    try:
        target.symlink_to(retained.name)
    except OSError as exc:
        pytest.skip(f"symlinks unavailable: {exc}")

    with pytest.raises(sharder.ShardError) as error:
        sharder.verify_shard_bundle(inventory_bundle, output)
    assert error.value.code == "unsafe_bundle_entry"


def test_cli_build_and_verify_return_path_free_receipts(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    output = tmp_path / "shards"
    build = _run_cli(
        "--inventory-bundle",
        str(inventory_bundle),
        "--output-dir",
        str(output),
        "--shard-count",
        "2",
    )
    verify = _run_cli(
        "--inventory-bundle",
        str(inventory_bundle),
        "--verify-bundle",
        str(output),
    )

    assert build.returncode == verify.returncode == 0, (build.stderr, build.stdout, verify.stderr, verify.stdout)
    build_status = json.loads(build.stdout)
    verify_status = json.loads(verify.stdout)
    assert build_status["status"] == verify_status["status"] == "pass"
    assert build_status["bundle_tree_sha256"] == verify_status["bundle_tree_sha256"]
    assert build_status["work_item_count"] == verify_status["work_item_count"] == 3
    assert str(inventory_bundle) not in build.stdout + verify.stdout
    assert str(output) not in build.stdout + verify.stdout
    for key, expected in sharder.NO_WRITE_FLAGS.items():
        assert build_status[key] is verify_status[key] is expected


def test_cli_rejects_invalid_mode_arguments_without_output(tmp_path: Path) -> None:
    inventory_bundle = _build_inventory(tmp_path)
    result = _run_cli(
        "--inventory-bundle",
        str(inventory_bundle),
        "--verify-bundle",
        str(tmp_path / "missing"),
        "--shard-count",
        "2",
    )
    status = json.loads(result.stdout)
    assert result.returncode == 2
    assert status["status"] == "rejected"
    assert status["reason"] == "invalid_arguments"
