"""Worker-only offline SegFormer comparison; no live control authority."""

import argparse
import hashlib
import json
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    processor = SegformerImageProcessor.from_pretrained("/assets", local_files_only=True)
    model = (
        SegformerForSemanticSegmentation.from_pretrained("/assets", local_files_only=True)
        .cuda()
        .eval()
    )
    image = Image.open(args.image).convert("RGB").crop((100, 0, 700, 600))
    tensor = processor(images=image, return_tensors="pt")["pixel_values"].cuda()
    times = []
    with torch.inference_mode():
        for _ in range(6):
            torch.cuda.synchronize()
            start = time.monotonic()
            logits = model(tensor).logits
            logits = torch.nn.functional.interpolate(
                logits, size=(512, 512), mode="bilinear", align_corners=False
            )
            labels = logits.argmax(1)[0].cpu().numpy().astype(np.uint8)
            times.append((time.monotonic() - start) * 1000)
    args.output.mkdir(parents=True, exist_ok=True)
    Image.fromarray(labels).save(args.output / "labels.png")
    ids, counts = np.unique(labels, return_counts=True)
    report = {
        "source_sha256": hashlib.sha256(args.image.read_bytes()).hexdigest(),
        "classes": {model.config.id2label[int(i)]: int(counts[n]) for n, i in enumerate(ids)},
        "inference_ms": times,
    }
    (args.output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report))


if __name__ == "__main__":
    main()