# SPDX-License-Identifier: Apache-2.0
"""Where the tests find data that does not ship in code/: the AutowareFoundation/bevfusion weights, the research
scripts and goldens of the workspace, and the bundle's own small goldens. Missing data skips the test that needs it
(never a failure): the container image and a Hub clone have neither the workspace nor, offline, the weights.

- weights: ``$BEVFUSION_WEIGHTS_DIR`` (holds ``bevfusion_lidar.onnx`` and the yaml files), else the workspace download
  ``assets/bevfusion/hf_bevfusion`` (sha256-verified, S:bevfusion:53-65), else the local HF cache snapshot at the
  pinned revision (``local_files_only``: tests never download);
- research: ``$BEVFUSION_RESEARCH_DIR``, else ``research/bevfusion`` of the workspace (scripts, the sample-rosbag
  frame and its goldens, public_data inputs and goldens);
- large goldens: ``$BEVFUSION_GOLDEN_DIR``, else ``research/bevfusion/goldens`` (written by
  ``code/scripts/make_goldens.py``; never shipped);
- samples (``sample_path`` / ``find_sample``): the shipped ``samples/`` (``test_pcd.npz``, Apache-2.0), then the
  git-ignored ``staging_samples_pandaset/`` of the bundle (``pandaset_019_f40.npz`` and its stored reference:
  PandaSet-derived, not shipped until the user approves PandaSet samples, research/DATASETS.md N5;
  ``code/scripts/make_samples.py`` and ``make_goldens.py`` write it there). The one exception to "missing data
  skips": inside the workspace a missing sample FAILS (``missing``), so that the frozen sample gates never silently
  stop running on it.
"""
from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Optional

import numpy as np
import pytest

PKG_DIR = Path(__file__).resolve().parents[1]
SAMPLES = PKG_DIR / "samples"
GOLDENS = Path(__file__).resolve().parent / "goldens"
WORKSPACE = Path("/home/ubuntu/experiments/tt-models")
# The PandaSet sample waits in the bundle's git-ignored staging directory until it may ship (DATASETS.md N5); it stays
# the frame of the frozen ``sample.*`` gates of test_e2e_device.py. staging_samples_pandaset/README.md: swap-in steps.
PANDASET_STAGING = PKG_DIR.parents[1] / "staging_samples_pandaset"
SAMPLE_PANDASET_NAME = "pandaset_019_f40.npz"
SAMPLE_TESTPCD = SAMPLES / "test_pcd.npz"      # the shipped sample: autoware_universe test.pcd (Apache-2.0)
ONNX = "bevfusion_lidar.onnx"
# Development frames (goldens by code/scripts/make_goldens.py; every precision choice of the port is made on them):
DEV_FRAMES = ("pandaset_019_f40", "test_pcd", "rosbag_0144", "public_ps090_f40", "public_ns0103_k20",
              "public_ns0916_k10", "public_ns0103_k00")
# HELD-OUT frames (the research public-data goldens of research/bevfusion/public_data/golden, made by the research
# reference, SPEC 6.2): never used to choose anything; the device e2e test's generalisation check.
HELDOUT = ("ps019_f31", "ps019_f47", "ps090_f33", "ps090_f46", "ns0103_k07", "ns0103_k15", "ns0103_k33",
           "ns0916_k05", "ns0916_k25", "ns0916_k37")


def _first_dir(*candidates: Optional[Path]) -> Optional[Path]:
    for c in candidates:
        if c is not None and c.is_dir():
            return c
    return None


def in_workspace() -> bool:
    """True inside the porting workspace (the research tree exists): there, missing sample data is an error."""
    return (WORKSPACE / "research" / "bevfusion" / "SPEC.md").is_file()


def missing(message: str):
    """Missing sample data: FAIL inside the workspace (a gate must never run on less data than it was frozen on), skip
    elsewhere (a Hub clone or the container image has only the shipped sample)."""
    if in_workspace():
        pytest.fail(f"{message} [inside the workspace missing sample data is an error, not a skip]")
    pytest.skip(message)


def find_sample(name: str) -> Optional[Path]:
    """A sample file by name (input or stored reference): the shipped ``samples/``, then the PandaSet staging copy."""
    for d in (SAMPLES, PANDASET_STAGING):
        if (d / name).is_file():
            return d / name
    return None


def sample_path(name: str) -> Path:
    """:func:`find_sample`, or :func:`missing` (fail in the workspace, skip elsewhere)."""
    p = find_sample(name)
    if p is None:
        missing(f"sample {name} found neither in {SAMPLES} nor in {PANDASET_STAGING} (code/scripts/make_samples.py; "
                "the stored references: code/scripts/make_goldens.py)")
    return p


def pandaset_sample() -> Path:
    """The PandaSet 019 f40 sample (213,007 points, 2 sweeps, 95,109 voxels, bucket ``m``): ``samples/`` once it
    ships, until then the staging copy."""
    return sample_path(SAMPLE_PANDASET_NAME)


def find_weights_dir() -> Optional[Path]:
    env = os.environ.get("BEVFUSION_WEIGHTS_DIR")
    found = _first_dir(Path(env) if env else None, WORKSPACE / "assets" / "bevfusion" / "hf_bevfusion")
    if found is not None and (found / ONNX).is_file():
        return found
    try:
        from huggingface_hub import snapshot_download

        from tt_bevfusion.api import BEVFusion

        snap = Path(snapshot_download(BEVFusion.DEFAULT_REPO, revision=BEVFusion.DEFAULT_REVISION,
                                      allow_patterns=list(BEVFusion.ALLOW_PATTERNS), local_files_only=True))
        if (snap / ONNX).is_file():
            return snap
    except Exception:  # noqa: BLE001 -- not cached / no huggingface_hub: the caller skips
        pass
    return None


def weights_dir() -> Path:
    d = find_weights_dir()
    if d is None:
        pytest.skip("AutowareFoundation/bevfusion weights not found (set BEVFUSION_WEIGHTS_DIR)")
    return d


def research_dir() -> Path:
    env = os.environ.get("BEVFUSION_RESEARCH_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "bevfusion")
    if d is None or not (d / "scripts" / "bevfusion_ref.py").is_file():
        pytest.skip("research/bevfusion (the research reference scripts) not found")
    return d


def research_module(name: str = "bevfusion_ref"):
    """Import a research script by name (``bevfusion_ref``); its directory is workspace code, not a download. The
    research reference imports onnx / onnxruntime / torch / yaml at module level: callers skip without them."""
    pytest.importorskip("onnxruntime")
    pytest.importorskip("torch")
    pytest.importorskip("yaml")
    scripts = str(research_dir() / "scripts")
    if scripts not in sys.path:
        sys.path.insert(0, scripts)
    import importlib

    return importlib.import_module(name)


def research_cloud(name: str):
    """A research input as a ``PointCloud``: ``rosbag_0144`` (the Autoware sample-rosbag top scan #144, ground-levelled,
    single sweep, x, y, z; license unstated: research only) or a public-data id (``ps019_f40``, ``ns0103_k20``, ...:
    x, y, z, time_lag)."""
    from tt_bevfusion import io as tio

    r = research_dir()
    if name == "rosbag_0144":
        path = r / "sample_inputs" / "rosbag_top144_concat_leveled.npz"
        fields = ("x", "y", "z")
    else:
        path = r / "public_data" / "inputs" / f"{name}.npz"
        fields = ("x", "y", "z", "time_lag")
    if not path.is_file():
        pytest.skip(f"research input {path} not found")
    with np.load(path, allow_pickle=False) as z:
        pts = z["points"][:, : len(fields)]
    return tio.PointCloud(np.ascontiguousarray(pts, dtype=np.float32), fields)


def research_golden(name: str):
    """The research reference's golden npz of a frame: ``rosbag_0144`` -> ``golden_leveled/golden_lidar.npz``, a
    public id -> ``public_data/golden/<id>_lidar.npz``."""
    r = research_dir()
    path = (r / "golden_leveled" / "golden_lidar.npz" if name == "rosbag_0144"
            else r / "public_data" / "golden" / f"{name}_lidar.npz")
    if not path.is_file():
        pytest.skip(f"research golden {path} not found")
    return np.load(path, allow_pickle=False)


def golden_dir() -> Path:
    env = os.environ.get("BEVFUSION_GOLDEN_DIR")
    d = _first_dir(Path(env) if env else None, WORKSPACE / "research" / "bevfusion" / "goldens")
    if d is None:
        pytest.skip("large goldens (research/bevfusion/goldens) not found")
    return d