#!/usr/bin/env python3
"""Решатель задачи 2 «У него есть прошлое» (CatFS/2: снимок + журнал транзакций).

Независимая реализация: только stdlib + Pillow + numpy. Никаких импортов
генератора и чтения файлов организатора.

Сценарий восстановления (всё публично, см. format.txt внутри ФС):
плитки 32x32 -> сектора 128 байт -> суперблок CATF версии 2 -> снимок
каталога (записи с generation) -> журнал транзакций -> replay ТОЛЬКО
зафиксированных транзакций, операции сопоставляются объектам по паре
(inode, generation) -> последняя зафиксированная версия nose_note.txt
перед её зафиксированным удалением -> флаг.

Использование: python solvers/solve_task2.py <participant.png>
Печатает извлечённый флаг caplag{...}; при любой ошибке — сообщение в stderr
и ненулевой код выхода.
"""
import re
import struct
import sys
import zlib

import numpy as np
from PIL import Image

TILE = 32
SECTOR = 128
HEADER = 16
PAYLOAD_CAP = SECTOR - HEADER
END_OF_CHAIN = 0xFFFFFFFF
DIR_ENTRY = 32
DIR_PER_SECTOR = SECTOR // DIR_ENTRY
VERSION = 2

FLAG_DELETED = 1
FLAG_ZLIB = 2

TXN_BEGIN = 1
OP_WRITE = 2
OP_DELETE = 3
TXN_COMMIT = 4
OP_NAMES = {TXN_BEGIN: "TXN_BEGIN", OP_WRITE: "OP_WRITE", OP_DELETE: "OP_DELETE", TXN_COMMIT: "TXN_COMMIT"}

MAX_RECORDS = 4096
MAX_OBJECT = 1 << 20  # лимит распакованных данных на объект/патч

TARGET_NAME = "nose_note.txt"


class CatfsError(Exception):
    pass


def fail(msg):
    print(f"solve_task2: ошибка: {msg}", file=sys.stderr)
    sys.exit(1)


def read_sectors(path):
    """Плитки 32x32 -> сектора 128 байт (чётность синего канала, MSB-first)."""
    try:
        img = Image.open(path).convert("RGB")
    except Exception as e:
        fail(f"не удалось открыть PNG: {e}")
    a = np.asarray(img, dtype=np.uint8)
    h, w, _ = a.shape
    tx, ty = w // TILE, h // TILE
    if tx == 0 or ty == 0:
        fail("изображение слишком мало для плиток 32x32")
    total = tx * ty
    # Собираем биты чётности синего канала по плиткам, обход по строкам.
    blue = (a[: ty * TILE, : tx * TILE, 2] & 1).reshape(ty, TILE, tx, TILE)
    blue = blue.transpose(0, 2, 1, 3).reshape(total, TILE * TILE)
    sectors = np.packbits(blue, axis=1)  # bitorder='big' = MSB-first
    return sectors, total


def parse_superblock(sec):
    if sec[0:4] != b"CATF":
        raise CatfsError("плохой magic суперблока")
    if sec[4] != VERSION:
        raise CatfsError(f"неподдерживаемая версия {sec[4]} (ожидалась {VERSION})")
    if sec[5] != 7:
        raise CatfsError(f"неожиданный sector_size_log2 {sec[5]}")
    (stored_crc,) = struct.unpack_from("<I", sec, 124)
    if zlib.crc32(bytes(sec[:124])) != stored_crc:
        raise CatfsError("CRC суперблока не сошёлся")
    (total, used, dir_start, dir_sectors, txn_start, txn_sectors, tile) = struct.unpack_from(
        "<7I", sec, 8
    )
    if tile != TILE:
        raise CatfsError("неожиданный tile_size в суперблоке")
    if used > total:
        raise CatfsError("sectors_used больше sectors_total")
    data_start = txn_start + txn_sectors
    if dir_start < 1 or dir_start + dir_sectors > txn_start or data_start > used:
        raise CatfsError("противоречивая раскладка областей в суперблоке")
    return {
        "total": total,
        "used": used,
        "dir_start": dir_start,
        "dir_sectors": dir_sectors,
        "txn_start": txn_start,
        "txn_sectors": txn_sectors,
        "data_start": data_start,
    }


def parse_directory(sectors, sb):
    entries = []
    for s in range(sb["dir_start"], sb["dir_start"] + sb["dir_sectors"]):
        if s >= sb["used"]:
            raise CatfsError("каталог выходит за sectors_used")
        sec = bytes(sectors[s])
        for i in range(DIR_PER_SECTOR):
            rec = sec[i * DIR_ENTRY : (i + 1) * DIR_ENTRY]
            name = rec[0:16].split(b"\x00", 1)[0]
            if not name:
                continue
            inode, generation = struct.unpack_from("<IH", rec, 16)
            flags = rec[23]
            size, first = struct.unpack_from("<2I", rec, 24)
            entries.append(
                {
                    "name": name.decode("ascii", "replace"),
                    "inode": inode,
                    "generation": generation,
                    "size": size,
                    "first": first,
                    "flags": flags,
                }
            )
    return entries


def parse_journal(sectors, sb):
    """Разбор журнала транзакций: заголовок (type u8, reserved u8, len u16),
    полезная часть, crc32 на запись. Конец — type=0 или нулевой хвост."""
    blob = b"".join(
        bytes(sectors[s])
        for s in range(sb["txn_start"], sb["txn_start"] + sb["txn_sectors"])
    )
    if len(blob) != sb["txn_sectors"] * SECTOR:
        raise CatfsError("журнал выходит за sectors_used")
    out = []
    pos = 0
    while True:
        if pos + 8 > len(blob):
            if any(blob[pos:]):
                raise CatfsError(f"обрезанный заголовок записи журнала по смещению {pos}")
            return out
        rtype = blob[pos]
        if rtype == 0:
            return out
        if len(out) >= MAX_RECORDS:
            raise CatfsError(f"больше {MAX_RECORDS} записей журнала")
        (plen,) = struct.unpack_from("<H", blob, pos + 2)
        end = pos + 4 + plen + 4
        if end > len(blob):
            raise CatfsError(f"запись журнала по смещению {pos} выходит за область журнала")
        raw = blob[pos:end]
        (stored_crc,) = struct.unpack_from("<I", raw, 4 + plen)
        if zlib.crc32(raw[: 4 + plen]) != stored_crc:
            raise CatfsError(f"запись журнала по смещению {pos}: CRC не сошёлся")
        p = raw[4 : 4 + plen]
        rec = {"type": rtype, "offset": pos}
        if rtype in (TXN_BEGIN, TXN_COMMIT):
            if plen != 4:
                raise CatfsError(f"запись type={rtype} по смещению {pos}: длина {plen}")
            (rec["txn"],) = struct.unpack_from("<I", p, 0)
        elif rtype == OP_DELETE:
            if plen != 12:
                raise CatfsError(f"OP_DELETE по смещению {pos}: длина {plen}")
            rec["txn"], rec["inode"], rec["generation"] = struct.unpack_from("<IIH", p, 0)
        elif rtype == OP_WRITE:
            if plen < 22:
                raise CatfsError(f"OP_WRITE по смещению {pos}: длина {plen}")
            (rec["txn"], rec["inode"], rec["generation"]) = struct.unpack_from("<IIH", p, 0)
            (rec["wr_offset"], rec["replace_len"], dlen) = struct.unpack_from("<IIH", p, 12)
            if dlen != plen - 22:
                raise CatfsError(f"OP_WRITE по смещению {pos}: data_len не сходится")
            rec["data"] = p[22:]
        else:
            raise CatfsError(f"неизвестный type={rtype} записи журнала по смещению {pos}")
        out.append(rec)
        pos = end


def read_chain(sectors, sb, start, what):
    """Обход цепочки секторов с полными проверками. Возвращает полезные байты."""
    if start < sb["data_start"] or start >= sb["used"]:
        raise CatfsError(f"{what}: первый сектор {start} вне области данных")
    out = bytearray()
    seen = set()
    s = start
    while True:
        if s in seen:
            raise CatfsError(f"{what}: цикл в цепочке на секторе {s}")
        seen.add(s)
        if len(seen) > sb["total"]:
            raise CatfsError(f"{what}: цепочка длиннее sectors_total")
        if s < sb["data_start"] or s >= sb["used"]:
            raise CatfsError(f"{what}: сектор {s} вне области данных")
        sec = bytes(sectors[s])
        sector_id, next_sector, used_len = struct.unpack_from("<IIH", sec, 0)
        if sector_id != s:
            raise CatfsError(f"{what}: сектор {s}: sector_id={sector_id}")
        if used_len > PAYLOAD_CAP:
            raise CatfsError(f"{what}: сектор {s}: used_length={used_len}")
        payload = sec[HEADER : HEADER + PAYLOAD_CAP]
        (pcrc,) = struct.unpack_from("<I", sec, 12)
        if zlib.crc32(payload) != pcrc:
            raise CatfsError(f"{what}: сектор {s}: payload_crc не сошёлся")
        out += payload[:used_len]
        if next_sector == END_OF_CHAIN:
            return bytes(out)
        s = next_sector


def zlib_limited(comp, size, what):
    """Распаковка zlib с жёстким лимитом (потоково)."""
    d = zlib.decompressobj()
    parts = []
    got = 0
    data = comp
    while True:
        chunk = d.decompress(data, 65536)
        got += len(chunk)
        if got > size:
            raise CatfsError(f"{what}: распакованные данные превышают лимит ({size})")
        parts.append(chunk)
        data = d.unconsumed_tail
        if not data:
            if d.eof:
                break
            raise CatfsError(f"{what}: zlib-поток оборван раньше времени")
    return b"".join(parts)


def replay(objects, history, series, records, log):
    """Replay журнала транзакций по публичным правилам format.txt.

    objects: inode -> {"gen": int, "data": bytearray}
    history: inode -> [ {"series", "gen", "data"|None, "deleted", "cause"} ]
    series:  inode -> номер текущей серии
    Возвращает список отклонённых записей (строки).
    """
    rejected = []

    def apply(op, txn_id):
        typ = op["type"]
        inode = op.get("inode")
        if typ == OP_WRITE:
            try:
                data = zlib_limited(op["data"], MAX_OBJECT, f"txn {txn_id} OP_WRITE")
            except CatfsError as e:
                rejected.append(str(e))
                return
            gen = op["generation"]
            if gen == 0:
                if inode in objects:
                    rejected.append(f"txn {txn_id} OP_WRITE inode={inode} gen=0: серия уже существует")
                    return
                if op["wr_offset"] != 0 or op["replace_len"] != 0:
                    rejected.append(
                        f"txn {txn_id} OP_WRITE inode={inode} gen=0: создание требует offset=0, replace_len=0"
                    )
                    return
                series[inode] = series.get(inode, 0) + 1
                objects[inode] = {"gen": 1, "data": bytearray(data)}
                history.setdefault(inode, []).append(
                    {"series": series[inode], "gen": 1, "data": bytes(data),
                     "deleted": False, "cause": f"txn {txn_id} OP_WRITE (create)"}
                )
                return
            obj = objects.get(inode)
            if obj is None or obj["gen"] != gen:
                rejected.append(
                    f"txn {txn_id} OP_WRITE inode={inode} gen={gen}: нет живого объекта с такой generation"
                )
                return
            base = obj["data"]
            off, rlen = op["wr_offset"], op["replace_len"]
            if off > len(base) or off + rlen > len(base):
                rejected.append(
                    f"txn {txn_id} OP_WRITE inode={inode}: диапазон [{off}:{off + rlen}] вне {len(base)} байт"
                )
                return
            nd = bytearray(base[:off]) + data + base[off + rlen :]
            if len(nd) > MAX_OBJECT:
                rejected.append(f"txn {txn_id} OP_WRITE inode={inode}: объект превысил лимит")
                return
            obj["data"] = nd
            obj["gen"] += 1
            history.setdefault(inode, []).append(
                {"series": series.get(inode, 1), "gen": obj["gen"], "data": bytes(nd),
                 "deleted": False, "cause": f"txn {txn_id} OP_WRITE"}
            )
        elif typ == OP_DELETE:
            obj = objects.get(inode)
            gen = op["generation"]
            if obj is None or obj["gen"] != gen:
                rejected.append(
                    f"txn {txn_id} OP_DELETE inode={inode} gen={gen}: нет живого объекта с такой generation"
                )
                return
            history.setdefault(inode, []).append(
                {"series": series.get(inode, 1), "gen": obj["gen"], "data": None,
                 "deleted": True, "cause": f"txn {txn_id} OP_DELETE"}
            )
            del objects[inode]

    open_txn = None
    ops = []
    for rec in records:
        if rec["type"] == TXN_BEGIN:
            # BEGIN при открытой транзакции: прежняя не зафиксирована (сбой) —
            # отбрасывается целиком.
            if open_txn is not None:
                log(f"    txn {open_txn}: нет COMMIT (сбой) — отброшена")
            open_txn, ops = rec["txn"], []
        elif rec["type"] == TXN_COMMIT:
            if open_txn is None or rec["txn"] != open_txn:
                raise CatfsError(
                    f"TXN_COMMIT {rec['txn']} не соответствует открытой транзакции {open_txn}"
                )
            for op in ops:
                apply(op, rec["txn"])
            log(f"    txn {rec['txn']}: COMMIT, операций {len(ops)}")
            open_txn, ops = None, []
        else:
            if open_txn is None:
                raise CatfsError("операция журнала вне транзакции")
            ops.append(rec)
    if open_txn is not None:
        log(f"    txn {open_txn}: нет COMMIT (сбой) — отброшена")
    return rejected


def main():
    if len(sys.argv) != 2:
        print("использование: python solvers/solve_task2.py <participant.png>", file=sys.stderr)
        sys.exit(2)
    sectors, total = read_sectors(sys.argv[1])
    print(f"[+] носитель: {total} секторов по {SECTOR} байт")
    try:
        sb = parse_superblock(bytes(sectors[0]))
        print(
            f"[+] суперблок ok (v{VERSION}): used={sb['used']} dir={sb['dir_start']}+{sb['dir_sectors']} "
            f"txn={sb['txn_start']}+{sb['txn_sectors']} data с {sb['data_start']}"
        )
        entries = parse_directory(sectors, sb)
        print(f"[+] снимок каталога: {len(entries)} записей")
        for e in entries:
            mark = "УДАЛЁН" if e["flags"] & FLAG_DELETED else "жив"
            print(
                f"    inode={e['inode']} gen={e['generation']} {e['name']!r} "
                f"size={e['size']} first={e['first']} [{mark}]"
            )

        # Живой format.txt — там публичная спецификация журнала транзакций.
        fmt = next((e for e in entries if e["name"] == "format.txt"), None)
        if fmt is None or (fmt["flags"] & FLAG_DELETED):
            raise CatfsError("живой format.txt не найден")
        fmt_text = read_chain(sectors, sb, fmt["first"], "format.txt").decode("utf-8")
        print("[+] format.txt прочитан (спецификация журнала транзакций и правил replay):")
        for line in fmt_text.splitlines()[:4]:
            print("    | " + line)

        # Целевой файл в снимке (базовая версия, generation из каталога).
        target = next(
            (e for e in entries if e["name"] == TARGET_NAME and not (e["flags"] & FLAG_DELETED)),
            None,
        )
        if target is None:
            raise CatfsError(f"{TARGET_NAME} не найден в снимке каталога")
        base = read_chain(sectors, sb, target["first"], TARGET_NAME)
        if target["flags"] & FLAG_ZLIB:
            base = zlib_limited(base, target["size"], TARGET_NAME)
        if len(base) != target["size"]:
            raise CatfsError(
                f"{TARGET_NAME}: базовая версия {len(base)} байт, каталог обещал {target['size']}"
            )
        print(
            f"[+] базовая версия из снимка: inode={target['inode']} "
            f"gen={target['generation']} ({len(base)} байт логических данных)"
        )

        records = parse_journal(sectors, sb)
        print(f"[+] журнал транзакций: {len(records)} записей (CRC каждой сошёлся)")
        for rec in records:
            line = f"    @{rec['offset']:<4} {OP_NAMES[rec['type']]:<10} txn={rec['txn']}"
            if rec["type"] in (OP_WRITE, OP_DELETE):
                line += f" inode={rec['inode']} gen={rec['generation']}"
            if rec["type"] == OP_WRITE:
                line += f" offset={rec['wr_offset']} replace_len={rec['replace_len']}"
            print(line)

        print("[+] replay зафиксированных транзакций (правила — format.txt):")
        logs = []
        objects = {target["inode"]: {"gen": target["generation"], "data": bytearray(base)}}
        history = {
            target["inode"]: [
                {"series": 1, "gen": target["generation"], "data": bytes(base),
                 "deleted": False, "cause": "snapshot"}
            ]
        }
        # Остальные живые объекты снимка тоже участвуют в состоянии (у них могут
        # быть свои операции); загружаем их для полноты replay.
        for e in entries:
            if e["inode"] == target["inode"] or (e["flags"] & FLAG_DELETED):
                continue
            raw = read_chain(sectors, sb, e["first"], e["name"])
            if e["flags"] & FLAG_ZLIB:
                raw = zlib_limited(raw, e["size"], e["name"])
            objects[e["inode"]] = {"gen": e["generation"], "data": bytearray(raw)}
            history[e["inode"]] = [
                {"series": 1, "gen": e["generation"], "data": bytes(raw),
                 "deleted": False, "cause": "snapshot"}
            ]
        rejected = replay(objects, history, {e["inode"]: 1 for e in entries}, records, logs.append)
        for line in logs:
            print(line)
        for line in rejected:
            print(f"    отклонено: {line}")

        hist = history.get(target["inode"], [])
        # Цель (публичная): последняя зафиксированная версия серии снимка
        # перед её зафиксированным удалением.
        series1 = [p for p in hist if p["series"] == 1]
        if not series1:
            raise CatfsError("в истории нет серии снимка для целевого inode")
        last = series1[-1]
        if last["deleted"]:
            if len(series1) < 2:
                raise CatfsError("перед удалением нет ни одной зафиксированной версии")
            point = series1[-2]
            print(
                f"[+] объект серии 1 зафиксировано удалён (gen={last['gen']}); "
                f"последняя зафиксированная версия перед удалением — gen={point['gen']}"
            )
        else:
            point = last
            print(f"[!] объект серии 1 не удалён; последняя зафиксированная версия — gen={point['gen']}")
        print("[+] история целевого inode:")
        for p in hist:
            if p["deleted"]:
                print(f"    серия {p['series']} gen={p['gen']} УДАЛЁН ({p['cause']})")
            else:
                print(f"    серия {p['series']} gen={p['gen']} {len(p['data'])} байт ({p['cause']})")
        text = point["data"].decode("utf-8")
    except CatfsError as e:
        fail(str(e))

    m = re.search(r"caplag\{[^}\s]+\}", text)
    if not m:
        fail("флаг caplag{...} в восстановленной версии не найден")
    print("[+] восстановленная записка (фрагмент):")
    tail = text[max(0, m.start() - 120) : m.end() + 40]
    for line in tail.splitlines():
        print("    > " + line)
    print()
    print("FLAG: " + m.group(0))


if __name__ == "__main__":
    main()