from __future__ import annotations

import os
import random
import tempfile
from pathlib import Path

import hvce


def make_tree(root: Path) -> None:
    root.mkdir(parents=True, exist_ok=True)
    (root / "a.txt").write_text("hello heaven vector\n" * 1000, encoding="utf-8")
    (root / "empty.bin").write_bytes(b"")
    (root / "sub").mkdir()
    (root / "sub" / "ramp.bin").write_bytes(bytes([i & 255 for i in range(8192)]))
    rng = random.Random(1)
    base = bytearray(rng.getrandbits(8) for _ in range(128 * 1024))
    (root / "media.mp4").write_bytes(bytes(base))
    base[0] ^= 1
    (root / "media_v2.mp4").write_bytes(bytes(base))


def test_roundtrip_plain():
    with tempfile.TemporaryDirectory() as td:
        src = Path(td) / "src"; out = Path(td) / "out"; arc = Path(td) / "x.hvce"
        make_tree(src)
        m, p = hvce.build_manifest(src, "fast", 64 * 1024, True, 32 * 1024)
        arc.write_bytes(hvce.append_recovery(hvce.make_plain_container(m, p), 0))
        hvce.extract_archive(arc, out, overwrite=True)
        hvce.compare_trees(src, out)


def test_roundtrip_encrypted():
    with tempfile.TemporaryDirectory() as td:
        src = Path(td) / "src"; out = Path(td) / "out"; arc = Path(td) / "x.hvce"
        make_tree(src)
        m, p = hvce.build_manifest(src, "balanced", 64 * 1024, True, 32 * 1024)
        arc.write_bytes(hvce.make_encrypted_container(m, p, "secret", 5000))
        hvce.extract_archive(arc, out, password="secret", overwrite=True)
        hvce.compare_trees(src, out)


def test_recovery_one_shard():
    with tempfile.TemporaryDirectory() as td:
        src = Path(td) / "src"; out = Path(td) / "out"; arc = Path(td) / "x.hvce"; bad = Path(td) / "bad.hvce"; fixed = Path(td) / "fixed.hvce"
        make_tree(src)
        m, p = hvce.build_manifest(src, "fast", 64 * 1024, True, 32 * 1024)
        arc.write_bytes(hvce.append_recovery(hvce.make_plain_container(m, p), 20))
        blob = bytearray(arc.read_bytes())
        pre, rec, tail = hvce.parse_recovery_tail(bytes(blob))
        assert rec is not None
        blob[min(len(pre) - 1, int(rec["shard_size"]) + 10)] ^= 0x42
        bad.write_bytes(bytes(blob))
        result = hvce.repair_archive(bad, fixed)
        assert result["repaired"] is True
        hvce.extract_archive(fixed, out, overwrite=True)
        hvce.compare_trees(src, out)


def test_recipe_polyword():
    data = bytearray()
    v, d1, d2 = 5, 7, 3
    for _ in range(1000):
        data.extend(v.to_bytes(4, "little"))
        v = (v + d1) & 0xffffffff
        d1 = (d1 + d2) & 0xffffffff
    cand = hvce.choose_representation(bytes(data), "balanced")
    dec = hvce.decode_candidate(cand.method, cand.payload, len(data), cand.params)
    assert dec == bytes(data)


def test_password_rejects_wrong_password():
    with tempfile.TemporaryDirectory() as td:
        src = Path(td) / "src"; arc = Path(td) / "x.hvce"
        make_tree(src)
        m, p = hvce.build_manifest(src, "fast", 64 * 1024, True, 32 * 1024)
        arc.write_bytes(hvce.make_encrypted_container(m, p, "secret", 5000))
        try:
            hvce.read_archive(arc, "wrong")
        except hvce.HVCEError:
            pass
        else:
            raise AssertionError("wrong password accepted")