"""
Bridge from the main app (torch17_new, Python 3.8) to SAM3 (sam3_worker.py,
run under sam2_env's Python 3.10 -- see config.py's SAM3 section for why
this has to be a subprocess rather than an in-process import).

One subprocess invocation can batch multiple (image, prompt) jobs -- e.g.
"flower" on the upload and "vase" on both the upload and the matched
template -- so the ~6s model load only happens once per request, not once
per job.
"""

import json
import logging
import os
import shutil
import subprocess
import uuid

import cv2
import numpy as np

import config

logger = logging.getLogger(__name__)


class Sam3Error(RuntimeError):
    pass


def run_jobs(images: dict, jobs: list, workdir: str) -> dict:
    """images: {image_key: bgr_ndarray}. jobs: [{"image": key, "prompt": str,
    "threshold": float}, ...]. Returns {(image_key, prompt): [{"mask": bool
    ndarray, "score": float, "box": [x1,y1,x2,y2] or None}, ...]}.

    Raises Sam3Error on any failure (missing token, subprocess crash,
    timeout, malformed response) -- callers treat this the same as any
    other soft/optional-feature failure (caught, logged, reported back as
    an error string, never taking the whole request down)."""
    if not config.HF_TOKEN:
        raise Sam3Error(
            "No HF_TOKEN found (checked environment and .env) -- SAM3's "
            "weights are gated on Hugging Face and can't be downloaded "
            "without an access-granted token."
        )

    request_dir = os.path.join(workdir, f"sam3_{uuid.uuid4().hex[:8]}")
    os.makedirs(request_dir, exist_ok=True)
    try:
        image_paths = {}
        for key, bgr in images.items():
            path = os.path.join(request_dir, f"{key}.png")
            cv2.imwrite(path, bgr)
            image_paths[key] = path

        output_dir = os.path.join(request_dir, "out")
        request = {"images": image_paths, "jobs": jobs, "output_dir": output_dir}
        request_path = os.path.join(request_dir, "request.json")
        with open(request_path, "w") as f:
            json.dump(request, f)

        env = dict(os.environ)
        env["HF_TOKEN"] = config.HF_TOKEN

        proc = subprocess.run(
            [config.SAM3_PYTHON_BIN, config.SAM3_WORKER_SCRIPT, "--request", request_path],
            capture_output=True, text=True, timeout=config.SAM3_TIMEOUT_SECONDS, env=env,
        )

        response_path = os.path.join(output_dir, "response.json")
        if not os.path.isfile(response_path):
            raise Sam3Error(
                f"SAM3 worker produced no response (exit {proc.returncode}): "
                f"{proc.stderr[-2000:] if proc.stderr else '(no stderr)'}"
            )

        with open(response_path) as f:
            response = json.load(f)

        if response.get("error"):
            raise Sam3Error(f"SAM3 worker failed: {response['error'][:2000]}")

        out = {}
        for entry in response["results"]:
            key = (entry["image"], entry["prompt"])
            instances = []
            for inst in entry["instances"]:
                mask_path = os.path.join(output_dir, inst["mask_file"])
                mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
                if mask is None:
                    continue
                instances.append({
                    "mask": mask > 127,
                    "score": inst["score"],
                    "box": inst.get("box"),
                })
            out[key] = instances
        return out

    except subprocess.TimeoutExpired as e:
        raise Sam3Error(f"SAM3 worker timed out after {config.SAM3_TIMEOUT_SECONDS}s") from e
    finally:
        shutil.rmtree(request_dir, ignore_errors=True)