src/jobs.py

from __future__ import annotations

import os
import re
from typing import Any

from huggingface_hub import Volume, fetch_job_logs, inspect_job, run_job

from .config import bucket_uri_from_source, user_bucket_source, settings
from .version import resolve_app_version
from .eval_config import effective_eval_config
from .bucket import assert_user_bucket_ready, write_text
from .runs import make_run_id, utc_now_iso, validate_run_id
from .worker_payload import (
    universal_model_card_worker_script,
    validate_existing_space_worker_script,
    python_decode_and_run_command,
)

SPACE_SLUG_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{1,95}$")

AUTO_SPACE_HARDWARE_CHOICES = {"zero-a10g", "cpu-basic", "t4-small", "t4-medium", "a10g-small", "a10g-large", "l4x1", "l40sx1"}
# v198.20: Jobs running ASF/Pi are intentionally CPU-only. GPU hardware is
# reserved for generated target Spaces, not for the Launch Build Job itself.
CPU_JOB_FLAVOR_CHOICES = {"cpu-upgrade", "cpu-basic"}
DEFAULT_JOB_FLAVOR = "cpu-upgrade"
FALLBACK_JOB_FLAVOR = "cpu-basic"
DEFAULT_PREFERRED_SPACE_HARDWARE = "zero-a10g"
DEFAULT_FALLBACK_SPACE_HARDWARE = "a10g-large"

def normalize_auto_space_hardware(value: str | None, *, default: str) -> str:
    """Return a hardware flavor safe for automatic Space fallback.

    Expensive or restricted tiers such as A100/H200 are intentionally not used
    by the automatic fallback path. Users can still select them manually in the
    generated Space settings if their account or organization is allowed to.
    """
    candidate = (value or default).strip()
    return candidate if candidate in AUTO_SPACE_HARDWARE_CHOICES else default



def _runs_prefix() -> str:
    return settings.bucket_runs_prefix.strip().strip("/") or "runs"


def _runs_mount_path() -> str:
    return f"{settings.bucket_mount.rstrip('/')}/{_runs_prefix()}"


def _worker_script_bucket_path(*, bucket_source: str, run_id: str) -> str:
    return f"{bucket_uri_from_source(bucket_source)}/{_runs_prefix()}/{run_id}/_worker.py"


def _worker_script_mount_path(*, run_id: str) -> str:
    return f"{_runs_mount_path().rstrip('/')}/{run_id}/_worker.py"


def _persist_worker_script(*, token: str, bucket_source: str, run_id: str, script: str) -> str:
    """Persist the Job worker script in the mounted bucket before launch.

    Earlier versions passed the full worker through WORKER_SCRIPT_B64. As the
    agentic core grew, that environment payload became large enough for HF Jobs
    to fail at process start with `argument list too long`. Storing the worker
    as a run artifact keeps the Job command/env small and makes the launched
    code auditable next to the run artifacts.
    """
    write_text(_worker_script_bucket_path(bucket_source=bucket_source, run_id=run_id), script, token=token)
    return _worker_script_mount_path(run_id=run_id)



def _job_eval_config() -> dict[str, Any]:
    # Eval archive is published by the ASF backend, not by user Jobs.
    # Jobs only need to know whether to write a local anonymized eval_record.json
    # into the user's run bucket. Environment variables still work as an admin
    # override for the backend publisher, but they no longer add a second Job volume.
    if settings.eval_enabled and settings.eval_bucket_source:
        return {
            "enabled": True,
            "bucket_source": settings.eval_bucket_source,
            "bucket_path": settings.eval_bucket_path,
            "job_mount_path": settings.eval_bucket_mount,
            "salt": settings.eval_salt,
            "include_redacted_tails": bool(settings.eval_include_redacted_tails),
            "include_model_id": bool(os.getenv("ASF_EVAL_INCLUDE_MODEL_ID", "false").strip().lower() in {"1", "true", "yes", "on"}),
            "publish_mode": "backend",
        }
    return effective_eval_config()

def _base_env(*, run_id: str, username: str, bucket_source: str, worker_script_path: str) -> dict[str, str]:
    env = {
        "RUN_ID": run_id,
        "HF_USERNAME": username or "unknown",
        "BUCKET_SOURCE": bucket_source,
        "OUTPUT_ROOT": settings.bucket_mount,
        "WORKER_SCRIPT_PATH": worker_script_path,
        "LAUNCHED_AT": utc_now_iso(),
    }
    eval_cfg = _job_eval_config()
    if eval_cfg.get("enabled"):
        env.update({
            "ASF_EVAL_RECORD_ENABLED": "true",
            "ASF_EVAL_PUBLISH_MODE": "backend",
            "ASF_EVAL_SALT": str(eval_cfg.get("salt") or ""),
            "ASF_EVAL_INCLUDE_REDACTED_TAILS": "true" if eval_cfg.get("include_redacted_tails") else "false",
            "ASF_EVAL_INCLUDE_MODEL_ID": "true" if eval_cfg.get("include_model_id") else "false",
            "ASF_VERSION": resolve_app_version(getattr(settings, "app_version", None)),
        })
    return env

def normalize_job_flavor(value: str | None, *, default: str = DEFAULT_JOB_FLAVOR) -> str:
    """Return a CPU-only HF Job flavor for the ASF/Pi Launch Build job.

    Target Space hardware remains independent and may still use ZeroGPU or fixed
    GPUs. This normalization prevents accidental GPU Job launches when an old
    env var such as ``SPACE_FACTORY_JOB_FLAVOR=a10g-small`` is still present.
    """
    candidate = (value or default or DEFAULT_JOB_FLAVOR).strip()
    return candidate if candidate in CPU_JOB_FLAVOR_CHOICES else DEFAULT_JOB_FLAVOR


def _run_job_with_flavor(*, token: str, env: dict[str, str], bucket_source: str, flavor: str, timeout: str) -> Any:
    return run_job(
        image=settings.job_image,
        command=python_decode_and_run_command(),
        flavor=flavor,
        timeout=timeout,
        env=env,
        # Keep the real token in Job secrets, not regular env metadata. Expose
        # both common aliases so Pi/HF tooling can authenticate inside the Job.
        secrets={"HF_TOKEN": token, "HUGGING_FACE_HUB_TOKEN": token},
        # Mount only the run-artifact prefix. The API persists `_worker.py` under
        # `hf://buckets/<source>/runs/<run_id>/_worker.py` before launching, so
        # the `runs/` prefix exists by the time the Job starts. Inside the Job,
        # that same prefix is mounted at `/output/runs`, keeping worker paths
        # stable as `/output/runs/<run_id>/...` without exposing the whole bucket.
        volumes=_job_volumes(bucket_source),
        token=token,
    )


def _job_volumes(bucket_source: str) -> list[Volume]:
    # User Jobs mount only the signed-in user's run bucket. The operator eval
    # archive is private and mounted on the ASF Space itself; the backend copies
    # anonymized records there after reading runs/<run_id>/eval_record.json.
    return [Volume(type="bucket", source=bucket_source, path=_runs_prefix(), mount_path=_runs_mount_path())]


def _launch_job(*, token: str, env: dict[str, str], bucket_source: str, flavor: str | None = None, timeout: str | None = None) -> Any:
    requested_flavor = normalize_job_flavor(flavor or settings.job_flavor)
    effective_timeout = timeout or settings.job_timeout
    try:
        return _run_job_with_flavor(
            token=token,
            env=env,
            bucket_source=bucket_source,
            flavor=requested_flavor,
            timeout=effective_timeout,
        )
    except Exception:
        if requested_flavor == FALLBACK_JOB_FLAVOR:
            raise
        return _run_job_with_flavor(
            token=token,
            env=env,
            bucket_source=bucket_source,
            flavor=FALLBACK_JOB_FLAVOR,
            timeout=effective_timeout,
        )


def _job_field(job: Any, name: str, default: Any = None) -> Any:
    if isinstance(job, dict):
        return job.get(name, default)
    return getattr(job, name, default)


def _job_status_stage(job: Any) -> Any:
    status = _job_field(job, "status")
    if isinstance(status, dict):
        return status.get("stage") or status.get("status") or status.get("state")
    return getattr(status, "stage", None) or getattr(status, "status", None) or getattr(status, "state", None)


def _job_url_from_id(*, job_id: str | None, bucket_source: str) -> str:
    if not job_id:
        return ""
    owner = str(bucket_source or "").split("/", 1)[0].strip()
    return f"https://huggingface.co/jobs/{owner}/{job_id}" if owner else ""


def _job_result(job: Any, *, run_id: str, kind: str, bucket_source: str, extra: dict[str, Any] | None = None) -> dict[str, Any]:
    job_id = str(_job_field(job, "id") or _job_field(job, "job_id") or _job_field(job, "jobId") or "").strip()
    job_url = str(_job_field(job, "url") or _job_field(job, "job_url") or "").strip()
    if not job_url:
        job_url = _job_url_from_id(job_id=job_id, bucket_source=bucket_source)
    payload: dict[str, Any] = {
        "run_id": run_id,
        "kind": kind,
        "job_id": job_id,
        "job_url": job_url,
        "status": _job_status_stage(job),
        "bucket_source": bucket_source,
        "bucket_uri": bucket_uri_from_source(bucket_source),
        "created_by": str(bucket_source or "").split("/", 1)[0].strip(),
        "created_at": str(_job_field(job, "created_at") or utc_now_iso()),
        "started_at": str(_job_field(job, "started_at") or _job_field(job, "created_at") or utc_now_iso()),
    }
    if extra:
        payload.update(extra)
    return payload


def normalize_target_space(*, username: str, target_slug: str | None, run_id: str) -> str:
    """Return `username/slug`, constrained to the signed-in user's namespace."""
    slug = (target_slug or "").strip()
    if not slug:
        slug = f"space-factory-{run_id}".lower()[:80]
    if "/" in slug:
        namespace, repo = slug.split("/", 1)
        if namespace != username:
            raise ValueError("The target Space must be created in your own namespace.")
        slug = repo
    if not SPACE_SLUG_RE.match(slug):
        raise ValueError("Invalid target Space name. Use letters, numbers, dots, underscores, or dashes.")
    return f"{username}/{slug}"


def _clean_repo_id(value: str | None, *, repo_kind: str) -> str:
    cleaned = (value or "").strip()
    cleaned = cleaned.replace("https://huggingface.co/spaces/", "")
    cleaned = cleaned.replace("https://huggingface.co/", "")
    cleaned = cleaned.strip("/")
    if "/" not in cleaned:
        raise ValueError(f"{repo_kind} must look like owner/name or a Hugging Face URL.")
    return cleaned


def launch_universal_model_card_job(
    *,
    token: str,
    username: str,
    target_slug: str | None = None,
    model_id: str | None = None,
    pi_model: str | None = None,
    preferred_space_hardware: str | None = None,
    fallback_space_hardware: str | None = None,
    allow_fixed_gpu_fallback: bool = True,
    try_zero_gpu_first: bool = True,
    implementation_mode: str | None = None,
    expected_output_type: str | None = None,
    run_id: str | None = None,
    bucket_name: str | None = None,
) -> dict[str, Any]:
    """Launch the public product builder: model card → private Space attempt."""
    if not token:
        raise ValueError("Missing OAuth token. Please sign in with Hugging Face first.")
    safe_run_id = validate_run_id(run_id) if run_id else make_run_id("universal")
    target_space_auto_generated = not bool((target_slug or "").strip())
    target_space_id = normalize_target_space(username=username, target_slug=target_slug, run_id=safe_run_id)
    clean_model_id = _clean_repo_id(model_id, repo_kind="Model ID")
    bucket_source = user_bucket_source(username=username, bucket_name=bucket_name)
    assert_user_bucket_ready(username=username, bucket_name=bucket_name, token=token)

    worker_script_path = _persist_worker_script(
        token=token,
        bucket_source=bucket_source,
        run_id=safe_run_id,
        script=universal_model_card_worker_script(),
    )
    env = _base_env(
        run_id=safe_run_id,
        username=username,
        bucket_source=bucket_source,
        worker_script_path=worker_script_path,
    )
    env["TARGET_SPACE_ID"] = target_space_id
    env["TARGET_SPACE_AUTO_GENERATED"] = "true" if target_space_auto_generated else "false"
    env["MODEL_ID"] = clean_model_id
    env["PI_MODEL"] = (pi_model or "Qwen/Qwen3-Coder-Next").strip()
    env["PREFERRED_SPACE_HARDWARE"] = normalize_auto_space_hardware(preferred_space_hardware, default=DEFAULT_PREFERRED_SPACE_HARDWARE)
    env["FALLBACK_SPACE_HARDWARE"] = normalize_auto_space_hardware(fallback_space_hardware, default=DEFAULT_FALLBACK_SPACE_HARDWARE)
    env["ALLOW_FIXED_GPU_FALLBACK"] = "true" if allow_fixed_gpu_fallback else "false"
    env["TRY_ZERO_GPU_FIRST"] = "true" if try_zero_gpu_first else "false"
    env["IMPLEMENTATION_MODE"] = (implementation_mode or "full-inference-gated").strip()
    env["EXPECTED_OUTPUT_TYPE"] = (expected_output_type or "any").strip()
    job = _launch_job(token=token, env=env, bucket_source=bucket_source, timeout="60m")
    return _job_result(
        job,
        run_id=safe_run_id,
        kind="universal_model_card_builder",
        bucket_source=bucket_source,
        extra={
            "target_space": target_space_id,
            "target_space_url": f"https://huggingface.co/spaces/{target_space_id}",
            "target_space_auto_generated": target_space_auto_generated,
            "model_id": clean_model_id,
            "pi_model": env["PI_MODEL"],
            "preferred_space_hardware": env["PREFERRED_SPACE_HARDWARE"],
            "fallback_space_hardware": env["FALLBACK_SPACE_HARDWARE"],
            "allow_fixed_gpu_fallback": allow_fixed_gpu_fallback,
            "try_zero_gpu_first": try_zero_gpu_first,
            "implementation_mode": env["IMPLEMENTATION_MODE"],
            "expected_output_type": env["EXPECTED_OUTPUT_TYPE"],
            "anonymous_eval_enabled": bool(_job_eval_config().get("enabled")),
            "anonymous_eval_publish_mode": "backend",
        },
    )



VALIDATION_PAYLOAD_JSON_MAX_BYTES = 16_000
VALIDATION_TIMEOUT_MIN_SECONDS = 30
VALIDATION_TIMEOUT_MAX_SECONDS = 3600


def clamp_validation_timeout_seconds(value: int | float | str | None, *, default: int = 1800) -> int:
    """Return a safe validation timeout bounded for browser/API inputs."""
    try:
        seconds = int(float(value if value is not None and str(value).strip() != "" else default))
    except Exception as exc:  # noqa: BLE001
        raise ValueError("live_timeout_seconds must be a number of seconds") from exc
    return max(VALIDATION_TIMEOUT_MIN_SECONDS, min(seconds, VALIDATION_TIMEOUT_MAX_SECONDS))


def ensure_validation_payload_json_fits_env(value: str | None, *, field_name: str) -> str:
    """Validate JSON payload size before passing it through HF Job env vars."""
    text = (value or "").strip()
    size = len(text.encode("utf-8"))
    if size > VALIDATION_PAYLOAD_JSON_MAX_BYTES:
        raise ValueError(
            f"{field_name} is too large for validation launch ({size} bytes; "
            f"max {VALIDATION_PAYLOAD_JSON_MAX_BYTES} bytes). Store large data in the Space or bucket instead."
        )
    return text


def launch_validate_existing_space_job(
    *,
    token: str,
    username: str,
    target_space_id: str,
    parent_build_run_id: str,
    api_name: str | None = None,
    test_args_json: str | None = None,
    test_kwargs_json: str | None = None,
    expected_output_type: str | None = None,
    live_timeout_seconds: int = 1800,
    validation_mode: str | None = None,
    effective_status_on_success: str | None = None,
    parent_replay_source_json: str | None = None,
    validation_launch_payload_json: str | None = None,
    run_id: str | None = None,
    bucket_name: str | None = None,
) -> dict[str, Any]:
    """Launch the public product validator for an existing generated Space."""
    if not token:
        raise ValueError("Missing OAuth token. Please sign in with Hugging Face first.")
    safe_run_id = validate_run_id(run_id) if run_id else make_run_id("validate")
    parent_run = validate_run_id(parent_build_run_id)
    target = _clean_repo_id(target_space_id, repo_kind="Target Space")
    namespace, _ = target.split("/", 1)
    if namespace != username:
        raise ValueError("For this version, target Space validation is limited to your own namespace.")
    bucket_source = user_bucket_source(username=username, bucket_name=bucket_name)
    assert_user_bucket_ready(username=username, bucket_name=bucket_name, token=token)
    worker_script_path = _persist_worker_script(
        token=token,
        bucket_source=bucket_source,
        run_id=safe_run_id,
        script=validate_existing_space_worker_script(),
    )
    env = _base_env(
        run_id=safe_run_id,
        username=username,
        bucket_source=bucket_source,
        worker_script_path=worker_script_path,
    )
    env["TARGET_SPACE_ID"] = target
    env["PARENT_BUILD_RUN_ID"] = parent_run
    env["VALIDATION_LINK_MODE"] = "build_run"
    env["SPACE_TEST_POLICY_MODE"] = (validation_mode or "complete").strip()
    env["EFFECTIVE_STATUS_ON_SUCCESS"] = (effective_status_on_success or "validated_after_manual_space_test").strip()
    env["API_NAME"] = (api_name if api_name is not None else "/generate").strip()
    env["ENDPOINT_DISCOVERY_REQUIRED"] = "true" if not env["API_NAME"] else "false"
    env["TEST_ARGS_JSON"] = ensure_validation_payload_json_fits_env(test_args_json or '["a cinematic robot cat astronaut, detailed, studio lighting"]', field_name="test_args_json")
    env["TEST_KWARGS_JSON"] = ensure_validation_payload_json_fits_env(test_kwargs_json or "{}", field_name="test_kwargs_json")
    if parent_replay_source_json:
        env["PARENT_REPLAY_SOURCE_JSON"] = ensure_validation_payload_json_fits_env(parent_replay_source_json, field_name="parent_replay_source_json")
    if validation_launch_payload_json:
        env["VALIDATION_LAUNCH_PAYLOAD_JSON"] = ensure_validation_payload_json_fits_env(validation_launch_payload_json, field_name="validation_launch_payload_json")
        try:
            _launch_payload = json.loads(validation_launch_payload_json)
            if isinstance(_launch_payload, dict) and _launch_payload.get("payload_source"):
                env["PAYLOAD_SOURCE"] = str(_launch_payload.get("payload_source") or "")
        except Exception:
            pass
    if parent_replay_source_json:
        env["PAYLOAD_SOURCE"] = "parent_automatic_smoke"
    env["EXPECTED_OUTPUT_TYPE"] = (expected_output_type or "image").strip()
    env["LIVE_TIMEOUT_SECONDS"] = str(clamp_validation_timeout_seconds(live_timeout_seconds))
    job = _launch_job(token=token, env=env, bucket_source=bucket_source, timeout="60m")
    return _job_result(
        job,
        run_id=safe_run_id,
        kind="linked_space_validation",
        bucket_source=bucket_source,
        extra={
            "target_space": target,
            "target_space_url": f"https://huggingface.co/spaces/{target}",
            "parent_build_run_id": parent_run,
            "source_kind": "build_run_prefill",
            "linked_target_space": target,
            "space_test_policy_mode": env.get("SPACE_TEST_POLICY_MODE"),
            "effective_status_on_success": env.get("EFFECTIVE_STATUS_ON_SUCCESS"),
            "api_name": env["API_NAME"] or "<discover>",
            "endpoint_discovery_required": env.get("ENDPOINT_DISCOVERY_REQUIRED") == "true",
            "expected_output_type": env["EXPECTED_OUTPUT_TYPE"],
            "test_args_json": env["TEST_ARGS_JSON"],
            "test_kwargs_json": env["TEST_KWARGS_JSON"],
            "payload_source": env.get("PAYLOAD_SOURCE") or "",
            "has_validation_launch_payload": bool(env.get("VALIDATION_LAUNCH_PAYLOAD_JSON")),
            "anonymous_eval_enabled": bool(_job_eval_config().get("enabled")),
            "anonymous_eval_publish_mode": "backend",
        },
    )


def inspect_job_safe(job_id: str, token: str | None = None) -> dict[str, Any]:
    if not job_id:
        return {"error": "Missing job_id"}
    try:
        info = inspect_job(job_id=job_id, token=token)
        status = getattr(info, "status", None)
        return {
            "id": info.id,
            "url": getattr(info, "url", None),
            "stage": getattr(status, "stage", None),
            "message": getattr(status, "message", None),
            "flavor": getattr(info, "flavor", None),
            "created_at": str(getattr(info, "created_at", "")),
            "started_at": str(getattr(info, "started_at", "")),
            "finished_at": str(getattr(info, "finished_at", "")),
        }
    except Exception as exc:  # noqa: BLE001
        return {"error": str(exc)}


def fetch_recent_logs_safe(job_id: str, token: str | None = None, max_lines: int = 120) -> str:
    if not job_id:
        return ""
    try:
        logs = list(fetch_job_logs(job_id=job_id, token=token))
        return "\n".join(str(line).rstrip("\n") for line in logs[-max_lines:])
    except Exception as exc:  # noqa: BLE001
        return f"Could not fetch job logs: {exc}"


def cancel_job_safe(job_id: str, *, namespace: str | None = None, token: str | None = None) -> dict[str, Any]:
    """Cancel a running HF Job. Returns a JSON-safe status payload.

    Import cancel_job lazily so older cached Space environments do not crash the
    whole app at import time. The endpoint reports a clear error if the installed
    huggingface_hub build does not expose Jobs cancellation.
    """
    if not job_id:
        return {"ok": False, "error": "Missing job_id"}
    try:
        from huggingface_hub import cancel_job  # type: ignore
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "job_id": job_id, "error": f"huggingface_hub.cancel_job is unavailable: {exc}"}
    try:
        cancel_job(job_id=job_id, namespace=namespace, token=token)
        return {"ok": True, "job_id": job_id, "status": "cancelled"}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "job_id": job_id, "error": str(exc)}