"""Read the native Q1 artifacts needed by the unified alignment adapter."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np

FEATURE_DIR = Path(__file__).resolve().parents[1] / "output" / "q1"


def load_sample(sample_id: str, feature_dir: Path = FEATURE_DIR) -> dict[str, Any]:
    """Load a Q1 sample by its video/clip ID from a feature directory."""
    feature_dir = Path(feature_dir)
    manifest = feature_dir / "manifest_q1.jsonl"
    for line in manifest.read_text(encoding="utf-8").splitlines():
        row = json.loads(line)
        if row["sample_id"] != sample_id:
            continue
        path = feature_dir / Path(row["feature_path"]).name
        with np.load(path, allow_pickle=False) as archive:
            result = {name: archive[name] for name in archive.files}
        result["_path"] = path
        result["_manifest"] = row
        result["_meta"] = json.loads(str(result["meta_json"]))
        return result
    raise KeyError(f"sample_id not found: {sample_id}")


def native_arrays(sample: dict[str, Any], modality: str):
    """Return values, observed dimensions, time intervals and quality evidence."""
    if modality == "text":
        values = sample["native_text_features"].astype(np.float32)
        return (
            values,
            np.broadcast_to(sample["native_text_observed"][:, None], values.shape),
            sample["native_text_intervals"].astype(np.float32),
            sample["native_text_quality_effective"].astype(np.float32),
            sample["native_text_quality_available"].astype(bool),
        )
    if modality in ("audio", "vision"):
        prefix = f"native_{modality}_"
        return (
            sample[prefix + "features"].astype(np.float32),
            sample[prefix + "mask"].astype(bool),
            sample[prefix + "intervals"].astype(np.float32),
            sample[prefix + "quality"].astype(np.float32),
            sample[prefix + "quality_available"].astype(bool),
        )
    raise ValueError(f"unknown modality: {modality}")