"""External read-only proof of current KFC4 entity attribute storage.

Uses the original owner chain, engine registration indices and reflected layouts.
Does not load the rebuilt DLL, execute game functions or write process memory.
"""
import json
import struct
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[6]
sys.path.insert(0, str(ROOT / "src/loader/runtime/profile-tools/dev/tools/enshrouded"))
from discovery.pe import PE
from discovery.live import Process


def capture(pid, role="client"):
    profile = json.loads((ROOT / "src/loader/runtime/profiles/enshrouded/client/25966634-resource-read.json").read_text())
    reflection = json.loads((ROOT / "src/loader/runtime/profile-tools/devdata/enshrouded-client-25966634-kfc4/live-discovery/reflection.json").read_text())
    catalog = json.loads((ROOT / "target/enshrouded-attribute-resolution.json").read_text())
    types = {ty["qualifiedName"]: ty for ty in reflection["types"]}
    layouts = {layout["qualified_name"]: layout for layout in catalog["component_layouts"]}
    roots = {}
    for entry in catalog["entries"]:
        roots.setdefault(entry["root_hash"], []).append(entry)
    for entries in roots.values():
        entries.sort(key=lambda entry: entry["index"])
    pe = PE(Path(r"C:\Program Files (x86)\Steam\steamapps\common\Enshrouded\enshrouded.exe"))
    assert pe.sha256 == profile["image"]["sha256"]
    process = Process(pid, pe)
    by_metadata = {process.base + ty["metadataRva"]: ty for ty in reflection["types"]}
    by_index = {ty["index"]: ty for ty in reflection["types"]}
    counters = Counter()
    def read(address, size):
        value = process.read(address, size)
        counters["read_calls"] += 1
        counters["read_bytes"] += size
        return value
    def q(address):
        return struct.unpack("<Q", read(address, 8))[0]
    def u(address):
        return struct.unpack("<I", read(address, 4))[0]
    def chain():
        p = profile["ownerLayout"]
        singleton = process.base + p["rootObjectRva"]
        assert q(singleton) == process.base + p["rootVtableRva"]
        client = q(singleton + p["clientPointerOffset"])
        assert q(client) == process.base + p["clientVtableRva"]
        controller = q(client + p["controllerPointerOffset"])
        assert q(controller) == process.base + p["controllerVtableRva"]
        assert q(controller + p["controllerClientBackPointerOffset"]) == client
        main = q(controller + p["controllerMainOwnerOffset"])
        ownership = dict(singleton=singleton, client=client, controller=controller)
        if role == "host":
            scene = q(client + 0x630)
            data, count, capacity = struct.unpack("<QQQ", read(scene + 0x4e0, 24))
            assert 0 < count <= capacity <= 4096
            active = []
            for index in range(count):
                entry = read(data + index * 0x90, 0x90)
                world, uuid0, uuid1 = struct.unpack_from("<QQQ", entry)
                if world and entry[0x18] == 5:
                    active.append((index, world, uuid0, uuid1, entry[0x18]))
            assert len(active) == 1, "host capture requires one proven active scene"
            index, main, uuid0, uuid1, state = active[0]
            ownership.update(scene=scene, activeData=data, activeCount=count, activeCapacity=capacity,
                activeIndex=index, worldUuid=[uuid0, uuid1], activeState=state)
        ecs = profile["ecsReadLayout"]
        execution = q(main + ecs["mainExecutionOwnerOffset"])
        owner = execution + ecs["executionRegistryOwnerOffset"]
        manager = execution + ecs["executionEntityManagerOffset"]
        assert q(manager + ecs["entityManagerOwnerOffset"]) == owner
        ownership.update(main=main, execution=execution, registryOwner=owner, manager=manager)
        return ownership
    try:
        ownership = chain()
        owner, manager = ownership["registryOwner"], ownership["manager"]
        ecs = profile["ecsReadLayout"]
        registry_headers = read(owner, 296)
        registry_count = struct.unpack_from("<Q", registry_headers, 16)[0]
        assert 1 <= registry_count <= 1024
        records = struct.unpack_from("<Q", registry_headers, 8)[0]
        storage_sizes = struct.unpack_from("<Q", registry_headers, 232)[0]
        storage_types = struct.unpack_from("<Q", registry_headers, 256)[0]
        assert all(struct.unpack_from("<Q", registry_headers, at)[0] == registry_count for at in (240, 264, 288))
        record_data = read(records, registry_count * 256)
        sizes_data = read(storage_sizes, registry_count * 2)
        types_data = read(storage_types, registry_count * 8)
        targets = {}
        registration_names = set()
        for index in range(registry_count):
            row = record_data[index * 256:(index + 1) * 256]
            name_pointer, name_length = struct.unpack_from("<QQ", row, 16)
            name = read(name_pointer, name_length).rstrip(b"\0").decode()
            assert name not in registration_names
            registration_names.add(name)
            assert name in types
            assert struct.unpack_from("<I", row, 32)[0] == types[name]["qualifiedHash"]
            storage_pointer = struct.unpack_from("<Q", row, 40)[0]
            configuration_pointer = struct.unpack_from("<Q", row, 48)[0]
            assert not configuration_pointer or configuration_pointer in by_metadata
            if not storage_pointer:
                assert not struct.unpack_from("<H", sizes_data, index * 2)[0]
                assert not struct.unpack_from("<Q", types_data, index * 8)[0]
                continue
            assert storage_pointer in by_metadata
            ty = by_metadata[storage_pointer]
            assert ty["qualifiedName"] not in targets
            expected_type = process.base + ty["metadataRva"]
            assert by_metadata[configuration_pointer if configuration_pointer else storage_pointer]["qualifiedName"] == name
            assert struct.unpack_from("<Q", row, 40)[0] == expected_type
            assert struct.unpack_from("<Q", types_data, index * 8)[0] == expected_type
            assert struct.unpack_from("<I", row, 56)[0] & 0xffff == ty["size"]
            assert struct.unpack_from("<H", sizes_data, index * 2)[0] == ty["size"]
            descriptor = read(expected_type, 0x90)
            assert struct.unpack_from("<I", descriptor, 0x40)[0] == ty["size"]
            assert struct.unpack_from("<I", descriptor, 0x50)[0] == ty["qualifiedHash"]
            assert read(expected_type + 0x40, 24) == pe.read(ty["metadataRva"] + 0x40, 24)
            targets[ty["qualifiedName"]] = dict(index=index, size=ty["size"], registrationName=name, qualifiedHash=ty["qualifiedHash"], metadataRva=ty["metadataRva"], recordPrefix=row[:88].hex())
        assert "keen::ecs::Health" in targets and "keen::ecs::Stamina" in targets
        count = q(manager + ecs["entityManagerCount"])
        table = q(manager + ecs["entityManagerTable"])
        assert 0 < count <= ecs["maximumEntities"]
        table_data = read(table, count * 8)
        archetypes = {}
        archetype_columns = {}
        snapshots = []
        sample_types = set()
        observed = Counter()
        racing = Counter()
        read_failed = Counter()
        stride_mismatches = {}
        failures = Counter()
        def entity_view(raw):
            return tuple(struct.unpack_from(format_, raw, ecs[key])[0] for key, format_ in (
                ("entityId", "<I"), ("entityGeneration", "<I"), ("entityLayout", "<Q"),
                ("entityStorage", "<Q"), ("entityDefinitionPointer", "<Q"), ("entityRow", "<I")))
        for table_slot, (pointer,) in enumerate(struct.iter_unpack("<Q", table_data)):
            if not pointer:
                continue
            try:
                entity_raw = read(pointer, 56)
                entity_id, generation, layout, storage, definition, row_index = entity_view(entity_raw)
                if not entity_id or not layout or not storage or row_index > (1 << 24):
                    continue
                if layout not in archetypes:
                    archetypes[layout] = read(layout, ecs["componentStrides"] + registry_count * 2)
                    columns = []
                    for name, registration in targets.items():
                        index = registration["index"]
                        word = struct.unpack_from("<Q", archetypes[layout], ecs["componentBits"] + (index // 64) * 8)[0]
                        if word & (1 << (index % 64)):
                            offset = struct.unpack_from("<H", archetypes[layout], ecs["componentOffsets"] + index * 2)[0]
                            stride = struct.unpack_from("<H", archetypes[layout], ecs["componentStrides"] + index * 2)[0]
                            columns.append((name, registration, word, offset, stride))
                    archetype_columns[layout] = columns
                layout_raw = archetypes[layout]
                for name, registration, word, column_offset, stride in archetype_columns[layout]:
                    index, size = registration["index"], registration["size"]
                    observed[name] += 1
                    if name in sample_types and name not in ("keen::ecs::Health", "keen::ecs::Stamina"):
                        continue
                    if stride != size:
                        failures["stride_mismatch"] += 1
                        stride_mismatches.setdefault(name, set()).add((stride, size, column_offset))
                        continue
                    address = storage + column_offset + row_index * stride
                    try:
                        raw = read(address, size)
                        raw_again = read(address, size)
                    except OSError:
                        read_failed[name] += 1
                        continue
                    if entity_view(read(pointer, 56)) != entity_view(entity_raw) or q(table + table_slot * 8) != pointer:
                        racing[name] += 1
                        continue
                    live_word = q(layout + ecs["componentBits"] + (index // 64) * 8)
                    live_column = read(layout + ecs["componentOffsets"] + index * 2, 2)
                    live_stride = read(layout + ecs["componentStrides"] + index * 2, 2)
                    if live_word != word or int.from_bytes(live_column, "little") != column_offset or int.from_bytes(live_stride, "little") != stride:
                        racing[name] += 1
                        continue
                    sample = dict(entityId=entity_id, generation=generation, tableSlot=table_slot,
                        entityPointer=hex(pointer), layoutPointer=hex(layout), storagePointer=hex(storage),
                        componentPointer=hex(address), row=row_index, component=name, registrationIndex=index,
                        registrationName=registration["registrationName"], qualifiedHash=registration["qualifiedHash"],
                        size=size, columnOffset=column_offset, stride=stride, bytes=raw.hex(),
                        repeatBytes=raw_again.hex(), valuesStable=raw == raw_again,
                        generationTableLayoutAndColumnStable=True)
                    if name not in layouts:
                        fields = {}
                        ty = types[name]
                        visited = set()
                        while ty["index"] not in visited:
                            visited.add(ty["index"])
                            for field in ty.get("structFields", {}).values():
                                ft = by_index[field["type"]]
                                format_ = {"UINT8":"<B", "UINT16":"<H", "UINT32":"<I", "UINT64":"<Q", "SINT8":"<b", "SINT16":"<h", "SINT32":"<i", "SINT64":"<q", "FLOAT32":"<f", "FLOAT64":"<d", "BOOL":"<?"}.get(ft["primitiveType"])
                                if format_ and field["dataOffset"] + ft["size"] <= len(raw):
                                    fields[field["name"]] = struct.unpack_from(format_, raw, field["dataOffset"])[0]
                            if ty["primitiveType"] not in ("STRUCT", "TYPEDEF") or ty.get("innerType") is None:
                                break
                            ty = by_index[ty["innerType"]]
                        sample["inlineScalarFields"] = fields
                        snapshots.append(sample)
                        sample_types.add(name)
                        continue
                    reflected = layouts[name]
                    def header_value(key, data=raw):
                        field = reflected["header"][key]
                        return int.from_bytes(data[field["offset"]:field["offset"] + field["size"]], "little")
                    root_id = header_value("rootId")
                    storage_offset = header_value("storageOffset")
                    storage_count = header_value("storageCount")
                    flags = header_value("flags")
                    if not flags & reflected["header"]["flags"]["initialized_mask"]:
                        failures["not_initialized"] += 1
                        continue
                    if root_id not in roots or storage_offset != reflected["storage_offset"] or storage_count * 4 != reflected["storage_size"]:
                        failures["root_or_layout_mismatch"] += 1
                        continue
                    entries = roots[root_id]
                    assert len(entries) == storage_count
                    format_ = {"keen::sint32":"<i", "keen::uint32":"<I", "keen::float32":"<f"}[reflected["storage_scalar_type"]]
                    values = {entry["name"]:struct.unpack_from(format_, raw, storage_offset + entry["index"] * 4)[0] for entry in entries}
                    if raw[:storage_offset] != raw_again[:storage_offset]:
                        failures["attribute_header_changed"] += 1
                        racing[name] += 1
                        continue
                    sample.update(rootId=root_id, rootName=entries[0]["name"], storageOffset=storage_offset,
                        storageCount=storage_count, flags=flags, values=values)
                    snapshots.append(sample)
                    sample_types.add(name)
            except (OSError, ValueError, struct.error) as error:
                failures[type(error).__name__] += 1
        assert chain() == ownership
        assert q(manager + ecs["entityManagerTable"]) == table
        count_after = q(manager + ecs["entityManagerCount"])
        assert count_after == count
        assert read(storage_types, registry_count * 8) == types_data
        assert read(storage_sizes, registry_count * 2) == sizes_data
        for registration in targets.values():
            assert read(records + registration["index"] * 256, 88).hex() == registration["recordPrefix"]
        return dict(capturedAt=datetime.now(timezone.utc).isoformat(), pid=pid,
            executable=pe.identity(), source="external-read-only-Process-memory-capture",
            newDllOrLuaExecutionVerified=False, engineCalls=False, processWrites=False, hooksInstalled=False,
            worldRole=role, ownerChain=ownership, registryCount=registry_count,
            registrations=targets, entityTable=hex(table), entityCount=count, entityCountAfter=count_after,
            originalOwnershipRegistryAndTableStable=True, accepted=len(snapshots),
            registeredStorageTypeCount=len(targets), observedTypeCount=len(observed),
            readSuccessTypeCount=len(sample_types), observedEntityOccurrences=dict(observed),
            unobservedTypes=sorted(set(targets) - set(observed)),
            observedWithoutValidatedSample=sorted(set(observed) - sample_types),
            racingAttemptsByType=dict(racing), unreadableAttemptsByType=dict(read_failed),
            strideMismatchesByType={name:sorted(values) for name, values in stride_mismatches.items()},
            componentCounts=dict(Counter(snapshot["component"] for snapshot in snapshots)),
            failures=dict(failures), reads=dict(counters), snapshots=snapshots)
    finally:
        process.close()


if __name__ == "__main__":
    role = sys.argv[2] if len(sys.argv) > 2 else "client"
    assert role in ("client", "host")
    result = capture(int(sys.argv[1]), role)
    output = Path(__file__).with_name("current-attributes-host-live-20261016.json" if role == "host" else "current-attributes-live-20261016.json")
    output.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(json.dumps({key:result[key] for key in ("pid", "accepted", "componentCounts", "entityCount", "failures", "reads")}, ensure_ascii=False))
