"""
Standalone SAM3 concept-segmentation worker.

Runs under a SEPARATE Python environment (sam2_env, Python 3.10) from the
main Vase Matcher app (torch17_new, Python 3.8) -- transformers>=5.5.0
(required for Sam3Model/Sam3Processor) itself requires Python>=3.10, so
this can't be imported in-process by the Flask app. Instead it's invoked
as a one-shot subprocess per request (see pipeline/sam3_client.py) with a
JSON request file describing which prompts to run against which images,
and writes a JSON response + one PNG mask per detected instance.

Loaded in 4-bit (NF4) via bitsandbytes -- empirically ~700MB resident /
~1.9GB peak during inference on an 8GB card, vs. several GB unquantized,
and a single model load handles every job in the request (one process
per /api/count_flowers call, not per prompt/image).

Usage:
    python sam3_worker.py --request <request.json>

Request JSON:
    {
      "images": {"<image_key>": "<path to bgr PNG/JPG>", ...},
      "jobs": [{"image": "<image_key>", "prompt": "flower", "threshold": 0.5}, ...],
      "output_dir": "<dir to write masks + response.json into>"
    }

Response JSON (written to <output_dir>/response.json):
    {
      "results": [
        {"image": "...", "prompt": "...", "instances": [
            {"mask_file": "...", "score": 0.73, "box": [x1,y1,x2,y2]}, ...
        ]}, ...
      ],
      "error": null   # or a string on failure
    }
"""

import argparse
import json
import os
import sys
import traceback

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

HF_MODEL_NAME = "facebook/sam3"


def load_model():
    import torch
    from transformers import BitsAndBytesConfig, Sam3Model, Sam3Processor

    token = os.environ.get("HF_TOKEN")
    processor = Sam3Processor.from_pretrained(HF_MODEL_NAME, token=token)
    quant_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4"
    )
    model = Sam3Model.from_pretrained(
        HF_MODEL_NAME, token=token, quantization_config=quant_config,
        device_map="cuda:0", low_cpu_mem_usage=True,
    )
    model.eval()
    return model, processor


def run_job(model, processor, pil_img, prompt, threshold):
    import torch

    inputs = processor(images=pil_img, text=prompt, return_tensors="pt").to("cuda:0")
    with torch.no_grad():
        outputs = model(**inputs)
    results = processor.post_process_instance_segmentation(
        outputs, threshold=threshold, mask_threshold=0.5,
        target_sizes=[pil_img.size[::-1]],
    )[0]

    masks = results["masks"].cpu().numpy()
    scores = results["scores"].cpu().numpy().tolist()
    boxes = results["boxes"].cpu().numpy().tolist() if "boxes" in results else [None] * len(masks)
    return masks, scores, boxes


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--request", required=True)
    args = parser.parse_args()

    with open(args.request) as f:
        request = json.load(f)

    output_dir = request["output_dir"]
    os.makedirs(output_dir, exist_ok=True)
    response = {"results": [], "error": None}

    try:
        import cv2
        import numpy as np
        from PIL import Image

        model, processor = load_model()

        loaded_images = {}
        for key, path in request["images"].items():
            bgr = cv2.imread(path, cv2.IMREAD_COLOR)
            if bgr is None:
                raise ValueError(f"could not read image {path!r} for key {key!r}")
            loaded_images[key] = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))

        for job_idx, job in enumerate(request["jobs"]):
            image_key = job["image"]
            prompt = job["prompt"]
            threshold = job.get("threshold", 0.5)
            pil_img = loaded_images[image_key]

            masks, scores, boxes = run_job(model, processor, pil_img, prompt, threshold)

            instances = []
            for i, (m, score, box) in enumerate(zip(masks, scores, boxes)):
                m_bin = (m > 0.5).astype("uint8") * 255
                mask_file = f"{image_key}_{prompt.replace(' ', '_')}_{job_idx}_{i}.png"
                cv2.imwrite(os.path.join(output_dir, mask_file), m_bin)
                instances.append({"mask_file": mask_file, "score": round(float(score), 4), "box": box})

            response["results"].append({
                "image": image_key, "prompt": prompt, "instances": instances,
            })

    except Exception as e:
        response["error"] = f"{e}\n{traceback.format_exc()}"

    with open(os.path.join(output_dir, "response.json"), "w") as f:
        json.dump(response, f)

    if response["error"]:
        print(response["error"], file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()