# backend/data_manager.py

import os
import base64
import io
from typing import List, Optional, Any, Dict

import numpy as np
from datasets import load_dataset

from .config import AUDIO_DATASET_ID
from .models import Clip

try:
    import soundfile as sf
except ImportError:
    sf = None


class DataManager:
    """Handles loading and processing data from Hugging Face."""

    def __init__(self, dataset_id: str = AUDIO_DATASET_ID):
        self.dataset_id = dataset_id
        self._clips: Optional[List[Clip]] = None
        self._loading = False

    def _audio_to_data_url(self, audio_val) -> Optional[str]:
        """
        Accepts:
        - torchcodec AudioDecoder
        - dict-like with 'path' / 'array' / 'sampling_rate'
        Returns data:audio/wav;base64,... or None.
        """
        # 1) Try to get a real file path and read it
        try:
            path = None
            if isinstance(audio_val, dict) and "path" in audio_val:
                path = audio_val["path"]
            else:
                # mapping-like: try __getitem__ then attribute
                try:
                    path = audio_val["path"]  # works on some decoders
                except Exception:
                    path = getattr(audio_val, "path", None)

            if isinstance(path, str) and os.path.exists(path):
                with open(path, "rb") as f:
                    audio_bytes = f.read()
                b64 = base64.b64encode(audio_bytes).decode("ascii")
                return f"data:audio/wav;base64,{b64}"
        except Exception as e:
            print(f"[WARN] Failed to build data URL from path: {e}")

        # 2) Fallback: use array + sampling_rate and render WAV in-memory
        try:
            array = None
            sr = None

            if isinstance(audio_val, dict):
                array = audio_val.get("array")
                sr = audio_val.get("sampling_rate")
            if array is None or sr is None:
                # try mapping-style then attributes
                try:
                    array = audio_val["array"]
                    sr = audio_val["sampling_rate"]
                except Exception:
                    array = getattr(audio_val, "array", None)
                    sr = getattr(audio_val, "sampling_rate", None)

            if array is not None and sr is not None and sf is not None:
                buf = io.BytesIO()
                sf.write(buf, np.array(array), int(sr), format="WAV")
                b64 = base64.b64encode(buf.getvalue()).decode("ascii")
                return f"data:audio/wav;base64,{b64}"
        except Exception as e:
            print(f"[WARN] Failed to build data URL from array/sr: {e}")

        print("[WARN] Could not build audio data URL for this example")
        return None

    def load_clips(self) -> List[Clip]:
        if self._clips is not None:
            return self._clips

        if self._loading:
            print("Dataset loading already in progress...")
            return []

        self._loading = True

        print(f"Loading dataset {self.dataset_id}...")
        dataset = load_dataset(self.dataset_id, split="train")

        clips: List[Clip] = []
        for row in dataset:
            audio_val = row.get("audio")

            audio_url = self._audio_to_data_url(audio_val)
            if audio_url is None:
                print(f"[WARN] Skipping clip {row.get('exercise_id')} – could not build audio URL")
                continue

            clip = Clip(
                id=f"{row['model']}_{row['speaker']}_{row['exercise_id']}",
                model=row["model"],
                speaker=row["speaker"],
                exercise=row["exercise"],
                exercise_id=row["exercise_id"],
                transcript=row["rt"],
                audio_url=audio_url,  # string usable in <audio src="...">
            )
            clips.append(clip)

        self._clips = clips
        self._loading = False
        print(f"Loaded {len(clips)} clips")
        return clips