import json
import os
from typing import Any, Dict, List, Tuple


def load_teacher_units(teacher_units_dir: str) -> Tuple[Dict[str, Dict[str, Any]], List[str]]:
    if not os.path.isdir(teacher_units_dir):
        raise FileNotFoundError(
            f"Teacher units not found: {teacher_units_dir}. Please generate units offline."
        )

    units: Dict[str, Dict[str, Any]] = {}
    for name in os.listdir(teacher_units_dir):
        unit_dir = os.path.join(teacher_units_dir, name)
        if not os.path.isdir(unit_dir):
            continue

        text_path = os.path.join(unit_dir, "text.json")
        align_words_path = os.path.join(unit_dir, "align_words.json")
        video_path = os.path.join(unit_dir, "video.mp4")
        audio_path = os.path.join(unit_dir, "audio.wav")
        align_phones_path = os.path.join(unit_dir, "align_phones.json")

        if not all(os.path.exists(p) for p in (text_path, align_words_path, video_path, audio_path)):
            continue

        with open(text_path, "r", encoding="utf-8") as f:
            text_data = json.load(f)
        with open(align_words_path, "r", encoding="utf-8") as f:
            align_words = json.load(f)

        align_phones = []
        if os.path.exists(align_phones_path):
            with open(align_phones_path, "r", encoding="utf-8") as f:
                align_phones = json.load(f)

        units[name] = {
            "unit_id": name,
            "text": text_data,
            "words": align_words,
            "video_path": video_path,
            "audio_path": audio_path,
            "phones": align_phones,
            "unit_index": text_data.get("unit_index", 0),
        }

    if not units:
        raise FileNotFoundError("No valid teacher units found.")

    ordered = sorted(units.values(), key=lambda x: (x["unit_index"], x["unit_id"]))
    unit_ids = [u["unit_id"] for u in ordered]
    return units, unit_ids