from __future__ import annotations

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 .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"}
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": False,
            "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": "v181",
        })
    return env


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:
    return run_job(
        image=settings.job_image,
        command=python_decode_and_run_command(),
        flavor=flavor or settings.job_flavor,
        timeout=timeout or settings.job_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_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",
        },
    )


def launch_validate_existing_space_job(
    *,
    token: str,
    username: str,
    target_space_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,
    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")
    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["API_NAME"] = (api_name or "/generate").strip()
    env["TEST_ARGS_JSON"] = (test_args_json or '["a cinematic robot cat astronaut, detailed, studio lighting"]').strip()
    env["TEST_KWARGS_JSON"] = (test_kwargs_json or "{}").strip()
    env["EXPECTED_OUTPUT_TYPE"] = (expected_output_type or "image").strip()
    env["LIVE_TIMEOUT_SECONDS"] = str(int(live_timeout_seconds or 1800))
    job = _launch_job(token=token, env=env, bucket_source=bucket_source, timeout="60m")
    return _job_result(
        job,
        run_id=safe_run_id,
        kind="validate_existing_space",
        bucket_source=bucket_source,
        extra={
            "target_space": target,
            "target_space_url": f"https://huggingface.co/spaces/{target}",
            "api_name": env["API_NAME"],
            "expected_output_type": env["EXPECTED_OUTPUT_TYPE"],
            "test_args_json": env["TEST_ARGS_JSON"],
            "test_kwargs_json": env["TEST_KWARGS_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)}