#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path
import re
import subprocess


SOURCE_NAMES = (
    "compress.h",
    "decompressor.c",
    "decompressor_lz4.c",
    "decompressor_lzma.c",
    "decompressor_deflate.c",
    "decompressor_zstd.c",
    "zdata.c",
)


def fail(message: str) -> None:
    raise SystemExit(message)


def replace_once(source: str, old: str, new: str, label: str) -> str:
    count = source.count(old)
    if count != 1:
        fail(f"{label}: expected one source occurrence, found {count}")
    return source.replace(old, new, 1)


def extract_function(source: str, name: str) -> str:
    match = re.search(r"^" + re.escape(name) + r"\s*\(", source, re.MULTILINE)
    if not match:
        fail(f"missing function: {name}")
    name_line = source.rfind("\n", 0, match.start()) + 1
    start = source.rfind("\n", 0, name_line - 1) + 1
    brace = source.find("{", match.end())
    if brace < 0:
        fail(f"missing function body: {name}")
    depth = 0
    state = "code"
    index = brace
    while index < len(source):
        char = source[index]
        following = source[index + 1] if index + 1 < len(source) else ""
        if state == "code":
            if char == "/" and following == "*":
                state = "block"
                index += 2
                continue
            if char == "/" and following == "/":
                state = "line"
                index += 2
                continue
            if char == '"':
                state = "string"
            elif char == "'":
                state = "character"
            elif char == "{":
                depth += 1
            elif char == "}":
                depth -= 1
                if depth == 0:
                    return source[start : index + 1]
        elif state == "block" and char == "*" and following == "/":
            state = "code"
            index += 2
            continue
        elif state == "line" and char == "\n":
            state = "code"
        elif state in {"string", "character"}:
            if char == "\\":
                index += 2
                continue
            if (state == "string" and char == '"') or (
                state == "character" and char == "'"
            ):
                state = "code"
        index += 1
    fail(f"unterminated function: {name}")


def committed(root: Path, baseline: str, name: str) -> str:
    completed = subprocess.run(
        ["git", "-C", str(root), "show", f"{baseline}:repo-pre-15/src/{name}"],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if completed.returncode != 0:
        fail(f"cannot read B25 baseline {name}: {completed.stderr.strip()}")
    return completed.stdout


def require_once(source: str, needle: str, label: str) -> None:
    count = source.count(needle)
    if count != 1:
        fail(f"{label}: expected one occurrence, found {count}")


def verify_exact_small_transforms(current: dict[str, str], base: dict[str, str]) -> None:
    expected = replace_once(
        base["compress.h"],
        "\tconst char *name;\n",
        "\tconst char *name;\n"
        "\t/* Callbacks return zero or a positive FreeBSD errno. */\n",
        "callback errno contract",
    )
    if current["compress.h"] != expected:
        fail("compress.h differs from the declared callback contract annotation")

    expected = base["decompressor.c"]
    expected = replace_once(
        expected,
        "\t\tif (size < sizeof(*lz4))\n\t\t\treturn (EINVAL);",
        "\t\tif (size < sizeof(*lz4))\n\t\t\treturn (EINTEGRITY);",
        "short LZ4 config",
    )
    expected = replace_once(
        expected,
        "\t\t    (Z_EROFS_PCLUSTER_MAX_SIZE >> sbi->blkszbits))\n"
        "\t\t\treturn (EINVAL);",
        "\t\t    (Z_EROFS_PCLUSTER_MAX_SIZE >> sbi->blkszbits))\n"
        "\t\t\treturn (EOPNOTSUPP);",
        "unsupported LZ4 pcluster size",
    )
    if expected.count("\t\treturn (EOVERFLOW);\n") != 2:
        fail("baseline config arithmetic errno count changed")
    expected = expected.replace(
        "\t\treturn (EOVERFLOW);\n", "\t\treturn (EINTEGRITY);\n"
    )
    expected = replace_once(expected, "\tint ret;\n", "", "folding temporary")
    expected = replace_once(
        expected,
        "\tret = decompressor->decompress(&rq);\n"
        "\treturn (ret == 0 ? 0 : EIO);",
        "\treturn (decompressor->decompress(&rq));",
        "dispatch preservation",
    )
    if current["decompressor.c"] != expected:
        fail("decompressor.c differs from the declared B25 dispatch/config transforms")

    expected = base["decompressor_lz4.c"]
    if expected.count("return (-1);") != 12 or expected.count(": -1);") != 1:
        fail("baseline LZ4 private-status count changed")
    expected = expected.replace("return (-1);", "return (EINTEGRITY);")
    expected = expected.replace(": -1);", ": EINTEGRITY);")
    if current["decompressor_lz4.c"] != expected:
        fail("LZ4 differs from the exact private-status conversion")

    if current["zdata.c"] != base["zdata.c"]:
        fail("B25 changed the already-correct zdata I/O or cleanup path")


def verify_source_contracts(
    source: dict[str, str], fixture: dict[str, object]
) -> dict[str, object]:
    joined = "\n".join(source[name] for name in SOURCE_NAMES)
    if re.search(r"return\s*\(\s*-E[A-Z0-9_]+", joined):
        fail("negative Linux errno entered the FreeBSD codec path")
    for name in SOURCE_NAMES:
        if "return (-1)" in source[name] or re.search(r":\s*-1\s*\)", source[name]):
            fail(f"private -1 codec status remains in {name}")

    dispatch = extract_function(source["decompressor.c"], "z_erofs_decompress")
    if dispatch.count("return (decompressor->decompress(&rq));") != 2:
        fail("typed backend dispatch does not have two direct callback exits")
    if "ret == 0 ? 0 : EIO" in dispatch:
        fail("dispatcher still folds typed backend errors")
    require_once(dispatch, "return (EOPNOTSUPP);", "unsupported algorithm")
    if dispatch.count("return (EINTEGRITY);") != 2:
        fail("dispatch-local corruption checks changed")
    for case in fixture["dispatch"]:
        if case.get("backend_errno") != case.get("expected"):
            fail("dispatch fixture does not require exact errno preservation")

    config_expectations = {
        "z_erofs_load_lz4_config": ("decompressor.c", "EINTEGRITY", "EOPNOTSUPP"),
        "z_erofs_load_lzma_config": (
            "decompressor_lzma.c",
            "EINTEGRITY",
            "EOPNOTSUPP",
        ),
        "z_erofs_load_deflate_config": (
            "decompressor_deflate.c",
            "EINTEGRITY",
            "EOPNOTSUPP",
        ),
        "z_erofs_load_zstd_config": (
            "decompressor_zstd.c",
            "EINTEGRITY",
            "EOPNOTSUPP",
        ),
    }
    for function, (name, malformed, unsupported) in config_expectations.items():
        body = extract_function(source[name], function)
        if f"return ({malformed});" not in body:
            fail(f"{function} lacks malformed-config errno")
        if f"return ({unsupported});" not in body:
            fail(f"{function} lacks unsupported-config errno")
    if fixture["config"] != [
        {"class": "malformed", "errno": "EINTEGRITY"},
        {"class": "unsupported", "errno": "EOPNOTSUPP"},
    ]:
        fail("config fixture does not encode the B25 typed classes")

    lzma = extract_function(source["decompressor_lzma.c"], "z_erofs_lzma_decompress")
    require_once(lzma, "return (EOVERFLOW);", "LZMA ABI overflow")
    require_once(lzma, "return (ENOMEM);", "LZMA allocation failure")
    if not (
        lzma.index("xz_dec_microlzma_run")
        < lzma.index("xz_dec_microlzma_end")
        < lzma.index("return (error);")
    ):
        fail("LZMA state is not released before status publication")

    deflate = extract_function(
        source["decompressor_deflate.c"], "z_erofs_deflate_decompress"
    )
    require_once(deflate, "inflateEnd(&strm)", "Deflate cleanup")
    require_once(
        deflate,
        "if (error == 0 && endret != Z_OK)",
        "Deflate primary-error preservation",
    )
    if deflate.index("inflateEnd(&strm)") > deflate.index("return (error);"):
        fail("Deflate returns before releasing initialized state")

    zstd_source = source["decompressor_zstd.c"]
    enabled = extract_function(zstd_source, "z_erofs_zstd_decompress")
    require_once(enabled, "return (ENOMEM);", "Zstd allocation failure")
    require_once(
        enabled,
        "if (error == 0 && ZSTD_isError(ret))",
        "Zstd primary-error preservation",
    )
    require_once(enabled, "error = EOPNOTSUPP;", "Zstd parameter mapping")
    if enabled.count("error = EINTEGRITY;") != 4:
        fail("Zstd stream corruption/completion mappings changed")
    require_once(enabled, "error = EIO;", "Zstd release mapping")
    if "ZSTD_getErrorCode" in zstd_source:
        fail("B25 added a Zstd provider symbol outside the existing ABI")
    if zstd_source.count("return (EOPNOTSUPP);") < 3:
        fail("Zstd config/runtime disabled paths are not typed unsupported")
    expected_zstd = {
        "create_context": "ENOMEM",
        "set_window_limit": "EOPNOTSUPP",
        "decompress_stream": "EINTEGRITY",
        "free_context": "EIO",
    }
    actual_zstd = {
        case.get("operation"): case.get("expected") for case in fixture["zstd"]
    }
    if actual_zstd != expected_zstd or len(fixture["zstd"]) != len(expected_zstd):
        fail("Zstd phase fixture does not match the source contract")

    read_extent = extract_function(source["zdata.c"], "z_erofs_read_extent")
    decode_at = read_extent.index("error = z_erofs_decompress")
    meta_release_at = read_extent.index("erofs_put_metabuf(&buf)", decode_at)
    physical_release_at = read_extent.index("erofs_brelse(compressed)", decode_at)
    error_at = read_extent.index("if (error != 0)", decode_at)
    free_at = read_extent.index("free(decoded, M_EROFS)", error_at)
    if not (
        decode_at < meta_release_at < error_at < free_at
        and decode_at < physical_release_at < error_at < free_at
    ):
        fail("zdata input/output cleanup ordering changed")

    return {
        "callbacks": 2,
        "config_loaders": len(config_expectations),
        "dispatch_direct": True,
        "dispatch_fixture_cases": len(fixture["dispatch"]),
        "negative_errno": False,
        "zdata_unchanged": True,
    }


def compile_status_harness(
    source: dict[str, str], fixture: dict[str, object], artifacts: Path
) -> dict[str, int]:
    xz_helper = extract_function(source["decompressor_lzma.c"], "z_erofs_lzma_error")
    zlib_helper = extract_function(
        source["decompressor_deflate.c"], "z_erofs_deflate_error"
    )
    checks: list[str] = []
    for group, function in (
        ("xz", "z_erofs_lzma_error"),
        ("zlib", "z_erofs_deflate_error"),
    ):
        for case in fixture[group]:
            checks.append(
                f'\tcheck("{group}:{case["status"]}", '
                f'{function}({case["status"]}), {case["expected"]});'
            )
    program = f'''#include <stddef.h>
#include <stdio.h>

#define EIO 5
#define ENOMEM 12
#define EOVERFLOW 75
#define EOPNOTSUPP 95
#define EINTEGRITY 97

enum xz_ret {{
\tXZ_OK,
\tXZ_STREAM_END,
\tXZ_UNSUPPORTED_CHECK,
\tXZ_MEM_ERROR,
\tXZ_MEMLIMIT_ERROR,
\tXZ_FORMAT_ERROR,
\tXZ_OPTIONS_ERROR,
\tXZ_DATA_ERROR,
\tXZ_BUF_ERROR
}};

#define Z_OK 0
#define Z_STREAM_END 1
#define Z_NEED_DICT 2
#define Z_ERRNO (-1)
#define Z_STREAM_ERROR (-2)
#define Z_DATA_ERROR (-3)
#define Z_MEM_ERROR (-4)
#define Z_BUF_ERROR (-5)
#define Z_VERSION_ERROR (-6)

{xz_helper}

{zlib_helper}

static int failures;

static void
check(const char *name, int actual, int expected)
{{
\tif (actual != expected) {{
\t\tfprintf(stderr, "%s: %d != %d\\n", name, actual, expected);
\t\t++failures;
\t}}
}}

int
main(void)
{{
{chr(10).join(checks)}
\tif (failures != 0)
\t\treturn (1);
\tprintf("status-map cases={len(checks)}\\n");
\treturn (0);
}}
'''
    harness = artifacts / "B25-status-map.c"
    binary = artifacts / "B25-status-map"
    harness.write_text(program, encoding="ascii")
    compile_run = subprocess.run(
        ["cc", "-std=c11", "-Wall", "-Wextra", "-Werror", str(harness), "-o", str(binary)],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    (artifacts / "B25-status-map-build.stdout").write_text(
        compile_run.stdout, encoding="utf-8"
    )
    (artifacts / "B25-status-map-build.stderr").write_text(
        compile_run.stderr, encoding="utf-8"
    )
    if compile_run.returncode != 0:
        fail("B25 status-map harness did not compile")
    execute = subprocess.run(
        [str(binary)],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    (artifacts / "B25-status-map.stdout").write_text(execute.stdout, encoding="utf-8")
    (artifacts / "B25-status-map.stderr").write_text(execute.stderr, encoding="utf-8")
    if execute.returncode != 0:
        fail("B25 status-map harness failed")
    binary.unlink()
    return {
        "xz": len(fixture["xz"]),
        "zlib": len(fixture["zlib"]),
        "zstd": len(fixture["zstd"]),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--artifacts", required=True, type=Path)
    parser.add_argument("--baseline", required=True)
    parser.add_argument("--dut", required=True, type=Path)
    parser.add_argument("--fixture", required=True, type=Path)
    parser.add_argument("--root", required=True, type=Path)
    args = parser.parse_args()

    fixture = json.loads(args.fixture.read_text(encoding="ascii"))
    if fixture.get("schema") != 1 or fixture.get("candidate") != "P15-023":
        fail("invalid B25 fixture identity")
    for key in ("config", "dispatch", "xz", "zlib", "zstd"):
        if not isinstance(fixture.get(key), list) or not fixture[key]:
            fail(f"empty B25 fixture group: {key}")

    source_dir = args.dut / "src"
    current = {
        name: (source_dir / name).read_text(encoding="utf-8") for name in SOURCE_NAMES
    }
    base = {name: committed(args.root, args.baseline, name) for name in SOURCE_NAMES}
    verify_exact_small_transforms(current, base)
    contracts = verify_source_contracts(current, fixture)
    mappings = compile_status_harness(current, fixture, args.artifacts)

    report = {
        "candidate": "P15-023",
        "contracts": contracts,
        "mapping_cases": mappings,
        "status": "PASS",
    }
    (args.artifacts / "B25-codec-errors.json").write_text(
        json.dumps(report, ensure_ascii=True, indent=2, sort_keys=True) + "\n",
        encoding="ascii",
    )
    print(
        "B25 codec errno contract: "
        f"xz={mappings['xz']} zlib={mappings['zlib']} zstd={mappings['zstd']}"
    )


if __name__ == "__main__":
    main()