src/edge_impulse.py

"""Upload a generated dataset directly to an Edge Impulse project.

Uses the Edge Impulse ingestion REST API with a project API key
(Project → Dashboard → Keys). WAV filenames use the ``label.<id>.wav``
convention, so Edge Impulse auto-assigns labels from the filename prefix.

Ingestion endpoints:
    POST https://ingestion.edgeimpulse.com/api/training/files
    POST https://ingestion.edgeimpulse.com/api/testing/files
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, List, Optional

import requests

_INGESTION_URL = "https://ingestion.edgeimpulse.com/api/{category}/files"
_TIMEOUT = 60
_BATCH_SIZE = 25

ProgressFn = Callable[[str], None]


@dataclass
class UploadResult:
    uploaded: int = 0
    failed: int = 0
    errors: List[str] = field(default_factory=list)


def verify_api_key(api_key: str) -> bool:
    """Return True if the ingestion API accepts the key.

    A key with no files still returns a non-auth error, so we treat any
    non-401/403 response as "key looks usable".
    """
    if not api_key or not api_key.strip():
        return False
    try:
        response = requests.post(
            _INGESTION_URL.format(category="training"),
            headers={"x-api-key": api_key.strip()},
            timeout=_TIMEOUT,
        )
    except requests.RequestException:
        return False
    return response.status_code not in (401, 403)


def _upload_batch(
    category: str,
    api_key: str,
    wav_paths: List[Path],
    allow_duplicates: bool,
) -> tuple[int, Optional[str]]:
    headers = {"x-api-key": api_key}
    if not allow_duplicates:
        headers["x-disallow-duplicates"] = "1"

    files = []
    handles = []
    try:
        for path in wav_paths:
            handle = path.open("rb")
            handles.append(handle)
            files.append(("data", (path.name, handle, "audio/wav")))
        response = requests.post(
            _INGESTION_URL.format(category=category),
            headers=headers,
            files=files,
            timeout=_TIMEOUT,
        )
    finally:
        for handle in handles:
            handle.close()

    if response.status_code == 200:
        return len(wav_paths), None
    return 0, f"HTTP {response.status_code}: {response.text[:200]}"


def upload_dataset(
    dataset_dir: str,
    api_key: str,
    allow_duplicates: bool = False,
    progress: Optional[ProgressFn] = None,
) -> UploadResult:
    """Upload the ``edge_impulse_upload/{training,testing}`` WAVs to a project."""
    def log(message: str) -> None:
        if progress:
            progress(message)
        else:
            print(message)

    api_key = (api_key or "").strip()
    if not api_key:
        raise ValueError("An Edge Impulse API key is required to upload.")

    base = Path(dataset_dir) / "edge_impulse_upload"
    result = UploadResult()

    for src_split, category in (("training", "training"), ("testing", "testing")):
        split_dir = base / src_split
        if not split_dir.exists():
            continue
        wavs = sorted(split_dir.glob("*.wav"))
        if not wavs:
            continue

        log(f"Uploading {len(wavs)} {category} file(s) to Edge Impulse...")
        for start in range(0, len(wavs), _BATCH_SIZE):
            batch = wavs[start:start + _BATCH_SIZE]
            uploaded, error = _upload_batch(category, api_key, batch, allow_duplicates)
            if error is None:
                result.uploaded += uploaded
                log(f"  {category}: {result.uploaded} uploaded")
            else:
                result.failed += len(batch)
                result.errors.append(f"{category} batch @ {start}: {error}")
                log(f"  WARNING: {category} batch failed: {error}")

    log(f"Edge Impulse upload complete: {result.uploaded} uploaded, {result.failed} failed.")
    return result