import argparse
import ctypes
from ctypes import wintypes
import gc
import hashlib
import json
import os
import subprocess
import sys
import time
from pathlib import Path

import avro
import fastavro


ROOT = Path(__file__).resolve().parent
ARTIFACT = ROOT / "artifacts" / "avro_bzip2_dos.avro"
RESULTS_DIR = ROOT / "results"
GENERATOR = ROOT / "generate_poc.py"


class PROCESS_MEMORY_COUNTERS(ctypes.Structure):
    _fields_ = [
        ("cb", wintypes.DWORD),
        ("PageFaultCount", wintypes.DWORD),
        ("PeakWorkingSetSize", ctypes.c_size_t),
        ("WorkingSetSize", ctypes.c_size_t),
        ("QuotaPeakPagedPoolUsage", ctypes.c_size_t),
        ("QuotaPagedPoolUsage", ctypes.c_size_t),
        ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t),
        ("QuotaNonPagedPoolUsage", ctypes.c_size_t),
        ("PagefileUsage", ctypes.c_size_t),
        ("PeakPagefileUsage", ctypes.c_size_t),
    ]


def working_set() -> dict[str, int]:
    get_process_memory_info = ctypes.windll.psapi.GetProcessMemoryInfo
    get_process_memory_info.argtypes = [wintypes.HANDLE, ctypes.POINTER(PROCESS_MEMORY_COUNTERS), wintypes.DWORD]
    get_process_memory_info.restype = wintypes.BOOL
    counters = PROCESS_MEMORY_COUNTERS()
    counters.cb = ctypes.sizeof(PROCESS_MEMORY_COUNTERS)
    ok = get_process_memory_info(
        ctypes.windll.kernel32.GetCurrentProcess(),
        ctypes.byref(counters),
        counters.cb,
    )
    if not ok:
        raise ctypes.WinError()
    return {
        "working_set": int(counters.WorkingSetSize),
        "peak_working_set": int(counters.PeakWorkingSetSize),
        "pagefile": int(counters.PagefileUsage),
    }


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def run_child(library: str) -> dict:
    cmd = [sys.executable, str(Path(__file__).resolve()), "--child", library, "--artifact", str(ARTIFACT)]
    proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
    return json.loads(proc.stdout)


def child_main(library: str, artifact: Path) -> None:
    gc.collect()
    before = working_set()
    t0 = time.perf_counter()

    if library == "avro":
        import avro.io
        from avro.datafile import DataFileReader

        with artifact.open("rb") as handle:
            reader = DataFileReader(handle, avro.io.DatumReader())
            record = next(iter(reader))
            reader.close()
    elif library == "fastavro":
        with artifact.open("rb") as handle:
            record = next(fastavro.reader(handle))
    else:
        raise SystemExit(f"unknown library: {library}")

    elapsed = time.perf_counter() - t0
    after = working_set()
    result = {
        "library": library,
        "elapsed_seconds": round(elapsed, 4),
        "working_set_before": before["working_set"],
        "working_set_after": after["working_set"],
        "peak_working_set": after["peak_working_set"],
        "pagefile_after": after["pagefile"],
        "tensor_name": record["tensor_name"],
        "tensor_bytes_len": len(record["tensor_bytes"]),
    }
    print(json.dumps(result))


def run_modelscan(artifact: Path) -> dict:
    report_path = RESULTS_DIR / "modelscan_report.json"
    cmd = [
        str(ROOT.parent / ".venv" / "Scripts" / "modelscan.exe"),
        "scan",
        "-p",
        str(artifact),
        "--show-skipped",
        "-r",
        "json",
        "-o",
        str(report_path),
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    report = None
    if report_path.exists():
        report = json.loads(report_path.read_text(encoding="utf-8"))
    raw = {
        "command": cmd,
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "report_path": str(report_path),
        "report": report,
    }
    return raw


def ensure_artifact() -> None:
    if ARTIFACT.exists():
        return
    subprocess.run(
        [sys.executable, str(GENERATOR), "--output", str(ARTIFACT)],
        cwd=str(ROOT),
        check=True,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify the Avro benign DoS PoC.")
    parser.add_argument("--child", choices=["avro", "fastavro"])
    parser.add_argument("--artifact", default=str(ARTIFACT))
    args = parser.parse_args()

    artifact = Path(args.artifact)
    if args.child:
        child_main(args.child, artifact)
        return

    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    ensure_artifact()

    avro_result = run_child("avro")
    fastavro_result = run_child("fastavro")
    modelscan_result = run_modelscan(artifact)

    summary = {
        "artifact": str(artifact),
        "artifact_bytes": artifact.stat().st_size,
        "artifact_sha256": sha256_file(artifact),
        "payload_bytes": avro_result["tensor_bytes_len"],
        "compression_ratio": round(avro_result["tensor_bytes_len"] / artifact.stat().st_size, 2),
        "versions": {
            "python": sys.version,
            "avro": avro.__version__,
            "fastavro": fastavro.__version__,
        },
        "results": {
            "avro": avro_result,
            "fastavro": fastavro_result,
        },
        "modelscan": modelscan_result,
    }

    (RESULTS_DIR / "results.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    (RESULTS_DIR / "modelscan_stdout.txt").write_text(modelscan_result["stdout"], encoding="utf-8")
    (RESULTS_DIR / "modelscan_stderr.txt").write_text(modelscan_result["stderr"], encoding="utf-8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()